[EPLB] Fix balancedness metric computation and add verbose reporting

docs/serving/expert_parallel_deployment.md

@@ -147,12 +147,15 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T
 
 | Parameter | Description | Default |
 | --------- | ----------- | ------- |
-| `window_size` | Number of engine steps to track for rebalancing decisions | 1000 |
-| `step_interval` | Frequency of rebalancing (every N engine steps) | 3000 |
-| `log_balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` |
+| `window_size` | Number of engine steps to track for rebalancing decisions | `1000` |
+| `step_interval` | Frequency of rebalancing (every N engine steps) | `3000` |
 | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` |
+| `log_balancedness` | Log balancedness metric: `max(tokens per rank) / mean(tokens per rank)`, averaged over layers. 1.0 = perfect balance | `false` |
+| `log_balancedness_interval` | How often to log balancedness (every N steps, requires `log_balancedness`) | `1` |
+| `log_balancedness_verbose` | Log a detailed multi-line report (per-rank layout, per-layer token table) | `false` |
 | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` |
 | `policy` | The policy type for expert parallel load balancing | `"default"` |
+| `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, or `null` (auto) | `null` |
 
 For example:
 

vllm/config/parallel.py

@@ -75,6 +75,13 @@ class EPLBConfig:
     """
     Interval for logging the balancedness.
     """
+    log_balancedness_verbose: bool = False
+    """
+    When True (together with ``log_balancedness``), log a multi-line EPLB report
+    per interval step: per-rank routed expert layout, replication summary,
+    and a per-layer / per-rank token table.
+    Off by default to avoid large logs and extra CPU work on rank 0.
+    """
     use_async: bool = False
     """
     Whether to use non-blocking EPLB.

vllm/distributed/eplb/eplb_state.py

@@ -26,6 +26,7 @@
 physical experts.
 """
 
+import sys
 import threading
 from collections.abc import Sequence
 from dataclasses import dataclass
@@ -48,6 +49,7 @@
 
 from .async_worker import start_async_worker
 from .eplb_communicator import EplbCommunicator, create_eplb_communicator
+from .eplb_utils import heat_cell
 from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy
 from .rebalance_execute import (
     RecvMetadata,
@@ -516,6 +518,106 @@ def add_model(
         self.model_states[model_config.compute_hash()] = model_state
         self.num_valid_physical_experts = model.num_physical_experts
 
+    def _dump(
+        self,
+        num_tokens_per_rank: torch.Tensor,
+        verbose: bool = False,
+    ) -> None:
+        """Print EPLB balancedness stats to stderr."""
+        layer_means = num_tokens_per_rank.mean(dim=1, dtype=torch.float64)
+        layer_maxes = num_tokens_per_rank.max(dim=1).values.to(torch.float64)
+        valid_layers = layer_means > 0
+
+        # Compute balancedness ratio:
+        # for each layer:
+        #   (max load across ranks) / (mean load across ranks)
+        # then average over active layers.
+        if valid_layers.any():
+            layer_ratios = layer_maxes[valid_layers] / layer_means[valid_layers]
+            balancedness = layer_ratios.mean().item()
+        else:
+            balancedness = 1.0
+        steps_left = (
+            self.expert_rearrangement_step_interval - self.expert_rearrangement_step
+        )
+        step = self.expert_rearrangement_step
+
+        if not verbose:
+            print(
+                f"EPLB: step={step}, balancedness={balancedness:.2f}x, "
+                f"\nNext rearrangement in {steps_left} steps",
+                file=sys.stderr,
+                flush=True,
+            )
+            return
+
+        # ---- verbose dump ----
+        lines: list[str] = [
+            "=== EPLB dump ===",
+            f"step={step}, balancedness={balancedness:.2f}x",
+            f"Next rearrangement in {steps_left} steps",
+        ]
+
+        num_ranks = num_tokens_per_rank.shape[1]
+        if num_ranks == 0:
+            lines += ["No EP ranks.", "=== end EPLB dump ==="]
+            print("\n".join(lines), file=sys.stderr, flush=True)
+            return
+
+        # Per-layer / per-rank token table
+        table = num_tokens_per_rank.long().cpu()  # [n_layers, num_ranks]
+        n_layers = table.shape[0]
+        total = int(table.sum().item())
+        if total > 0:
+            row_sums = table.sum(dim=1)
+            col_sums = table.sum(dim=0)
+
+            # Auto-size columns to fit the widest value
+            biggest_value = max(int(table.max().item()), int(col_sums.max().item()))
+            val_w = max(6, len(str(biggest_value)))
+            sum_w = max(6, len(str(total)))
+            label_w = len(f"Layer{n_layers - 1}")
+
+            # Header
+            lines.append("Tokens per MoE layer and EP rank:")
+            rank_hdr = " ".join(f"{'rank' + str(r):>{val_w}}" for r in range(num_ranks))
+            lines.append(f"{'':{label_w}} {rank_hdr}  {'sum':>{sum_w}}  max/mean")
+
+            # One row per MoE layer
+            for i in range(n_layers):
+                row = table[i]
+                row_min = int(row.min().item())
+                row_max = int(row.max().item())
+                row_total = int(row_sums[i].item())
+                row_mean = row_total / num_ranks
+                ratio = row_max / row_mean if row_mean > 0 else float("inf")
+                vals = [int(row[r].item()) for r in range(num_ranks)]
+                cells = " ".join(
+                    heat_cell(f"{v:>{val_w}}", v, row_min, row_max) for v in vals
+                )
+                lines.append(
+                    f"{'Layer' + str(i):{label_w}} {cells}"
+                    f"  {row_total:>{sum_w}}  {ratio:.2f}x"
+                )
+
+            # Totals row
+            mean_per_rank = total / num_ranks
+            totals_max = int(col_sums.max().item())
+            totals_ratio = (
+                totals_max / mean_per_rank if mean_per_rank > 0 else float("inf")
+            )
+            totals_cells = " ".join(
+                f"{int(col_sums[r].item()):>{val_w}}" for r in range(num_ranks)
+            )
+            sigma = "\u03a3"
+            lines.append(
+                f"{sigma:{label_w}} {totals_cells}"
+                f"  {total:>{sum_w}}  {totals_ratio:.2f}x"
+            )
+
+        lines.append("=== end EPLB dump ===")
+        print("\n".join(lines), file=sys.stderr, flush=True)
+
     def step(
         self,
         is_dummy: bool = False,
@@ -534,12 +636,6 @@ def step(
                 `profile_run` to reserve enough memory
                 for the communication buffer.
             log_stats (bool): If `True`, log the expert load metrics.
-
-        # Stats
-            The metrics are all summed up across layers.
-            - `avg_tokens`: The average load across ranks.
-            - `max_tokens`: The maximum load across ranks.
-            - `balancedness`: The ratio of average load to maximum load.
         """
         ep_group = get_ep_group().device_group
         if is_profile:
@@ -573,31 +669,10 @@ def step(
                     .float()
                 )
 
-                # Compute balancedness ratio:
-                # for each layer:
-                #   (mean load across ranks) / (max load across ranks)
-                avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
-                max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(dim=0)
-
-                # Just to make type checker happy
-                tokens_tensors: list[float] = torch.stack(
-                    [avg_tokens_tensor, max_tokens_tensor]
-                ).tolist()
-                avg_tokens, max_tokens = tokens_tensors
-                balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
-
                 if ep_group.rank() == 0:
-                    logger.info(
-                        "EPLB step: %d for model %s: avg_tokens=%.2f, "
-                        "max_tokens=%d, balancedness=%.4f, "
-                        "steps until the next rearrangement: %d",
-                        self.expert_rearrangement_step,
-                        eplb_model_state.model_name,
-                        avg_tokens,
-                        max_tokens,
-                        balancedness,
-                        self.expert_rearrangement_step_interval
-                        - self.expert_rearrangement_step,
+                    self._dump(
+                        num_tokens_per_rank,
+                        verbose=self.parallel_config.eplb_config.log_balancedness_verbose,
                     )
 
         # Update the expert load sliding window

vllm/distributed/eplb/eplb_utils.py

@@ -2,7 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Utility functions for EPLB (Expert Parallel Load Balancing)."""
 
+import functools
 import os
+import sys
 
 from vllm.config import ParallelConfig
 from vllm.logger import init_logger
@@ -63,3 +65,62 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None:
             "deepep_low_latency backend",
             scope="global",
         )
+
+
+# ---------------------------------------------------------------------------
+# Formatting helpers for EPLB dump output
+# ---------------------------------------------------------------------------
+
+
+@functools.lru_cache(maxsize=1)
+def _use_heat_color() -> bool:
+    return (
+        not os.environ.get("NO_COLOR", "")
+        and os.environ.get("TERM", "") != "dumb"
+        and sys.stderr.isatty()
+    )
+
+
+def heat_cell(text: str, val: float, vmin: float, vmax: float) -> str:
+    """Wrap *text* in green-to-red ANSI color based on *val* in [vmin, vmax]."""
+    if not _use_heat_color() or vmin >= vmax:
+        return text
+    t = max(0.0, min(1.0, (val - vmin) / (vmax - vmin)))
+    r, g = int(220 * t), int(220 * (1 - t))
+    return f"\033[38;2;{r};{g};30m{text}\033[0m"
+
+
+def human_tokens(n: float) -> str:
+    """1234 -> '1234', 12345 -> '12k', 1234567 -> '1235k'."""
+    v = int(round(n))
+    return str(v) if v < 10_000 else f"{round(v / 1000)}k"
+
+
+def compact_int_list(items: list) -> str:
+    """Format mixed str/int list with run compression: [shared, 0..63, 123]."""
+    if not items:
+        return "[]"
+    parts: list[str] = []
+    rs: int | None = None
+    re: int | None = None
+
+    def _flush() -> None:
+        if rs is not None:
+            parts.append(str(rs) if rs == re else f"{rs}..{re}")
+
+    for item in items:
+        if isinstance(item, str):
+            _flush()
+            rs = re = None
+            parts.append(item)
+        else:
+            x = int(item)
+            if rs is None or re is None:
+                rs = re = x
+            elif x == re + 1:
+                re = x
+            elif x != re:
+                _flush()
+                rs = re = x
+    _flush()
+    return "[" + ", ".join(parts) + "]"