From 72832b4418e0920acd03e062a9659fb8b527c6bb Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Tue, 7 Apr 2026 14:36:39 +0300 Subject: [PATCH 01/10] [EPLB] Fix balancedness computation: use per-layer mean/max across ranks The comment says "for each layer: (mean load across ranks) / (max load across ranks)" but the code was using dim=0 (averaging/maxing across layers) instead of dim=1 (across ranks within each layer). Fix to match the documented intent: compute mean/max across EP ranks for each MoE layer independently, then average the per-layer ratios over active layers. Signed-off-by: Artem Perevedentsev --- vllm/distributed/eplb/eplb_state.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index c56f8b0364aa..eccf76082ac9 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -576,15 +576,18 @@ def step( # 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 + # then average over active layers. + per_layer_mean = num_tokens_per_rank.mean(dim=1) + per_layer_max = num_tokens_per_rank.max(dim=1).values + active = per_layer_max > 0 + if active.any(): + balancedness = ( + per_layer_mean[active] / per_layer_max[active] + ).mean().item() + else: + balancedness = 0.0 + avg_tokens = per_layer_mean.sum().item() + max_tokens = per_layer_max.sum().item() if ep_group.rank() == 0: logger.info( From 3bd5159506a22207041574e71c71fe51e2916ced Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Tue, 7 Apr 2026 15:03:21 +0300 Subject: [PATCH 02/10] [EPLB] Use float64 for balancedness summation to avoid precision loss Compute per-layer mean and max in float64 to prevent precision loss when summing token counts across many MoE layers. Signed-off-by: Artem Perevedentsev --- vllm/distributed/eplb/eplb_state.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index eccf76082ac9..0d214c6736e3 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -577,13 +577,13 @@ def step( # for each layer: # (mean load across ranks) / (max load across ranks) # then average over active layers. - per_layer_mean = num_tokens_per_rank.mean(dim=1) - per_layer_max = num_tokens_per_rank.max(dim=1).values + per_layer_mean = num_tokens_per_rank.mean(dim=1, dtype=torch.float64) + per_layer_max = num_tokens_per_rank.max(dim=1).values.to(torch.float64) active = per_layer_max > 0 if active.any(): balancedness = ( - per_layer_mean[active] / per_layer_max[active] - ).mean().item() + (per_layer_mean[active] / per_layer_max[active]).mean().item() + ) else: balancedness = 0.0 avg_tokens = per_layer_mean.sum().item() From 29c8d6623cb9b7845aad6727ba75165a7a4e5253 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Thu, 9 Apr 2026 13:10:41 +0300 Subject: [PATCH 03/10] [EPLB] Fix balancedness metric computation and add verbose reporting Signed-off-by: Artem Perevedentsev --- docs/serving/expert_parallel_deployment.md | 9 +- vllm/config/parallel.py | 7 ++ vllm/distributed/eplb/eplb_state.py | 138 ++++++++++++++++----- vllm/distributed/eplb/eplb_utils.py | 61 +++++++++ 4 files changed, 179 insertions(+), 36 deletions(-) diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index d75ae7feb49e..55652bcce383 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/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: diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 107dcfa273eb..fe8e370de5b0 100644 --- a/vllm/config/parallel.py +++ b/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. diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 0d214c6736e3..cea94484d3b5 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/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,34 +669,10 @@ def step( .float() ) - # Compute balancedness ratio: - # for each layer: - # (mean load across ranks) / (max load across ranks) - # then average over active layers. - per_layer_mean = num_tokens_per_rank.mean(dim=1, dtype=torch.float64) - per_layer_max = num_tokens_per_rank.max(dim=1).values.to(torch.float64) - active = per_layer_max > 0 - if active.any(): - balancedness = ( - (per_layer_mean[active] / per_layer_max[active]).mean().item() - ) - else: - balancedness = 0.0 - avg_tokens = per_layer_mean.sum().item() - max_tokens = per_layer_max.sum().item() - 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 diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index f499b3e518b8..4bce531c7e3e 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/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) + "]" From 44a1887f144387c3d1e11ac9f222c6e25496526a Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Fri, 10 Apr 2026 18:49:44 +0300 Subject: [PATCH 04/10] Address review comments on balancedness reporting Signed-off-by: Artem Perevedentsev --- vllm/distributed/eplb/eplb_state.py | 47 +++++++++++++++++------------ vllm/distributed/eplb/eplb_utils.py | 36 ---------------------- 2 files changed, 28 insertions(+), 55 deletions(-) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index cea94484d3b5..42d3eb0e6223 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -518,43 +518,45 @@ def add_model( self.model_states[model_config.compute_hash()] = model_state self.num_valid_physical_experts = model.num_physical_experts - def _dump( + def _log_balancedness( self, num_tokens_per_rank: torch.Tensor, + model_name: str, 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 - ) + # (mean load across ranks) / (max load across ranks) + avg_tokens = layer_means.sum().item() + max_tokens = layer_maxes.sum().item() + balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 + step = self.expert_rearrangement_step + steps_left = self.expert_rearrangement_step_interval - 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, + logger.info( + "EPLB step: %d for model %s: avg_tokens=%.2f, " + "max_tokens=%d, balancedness=%.4f, " + "steps until the next rearrangement: %d", + step, + model_name, + avg_tokens, + max_tokens, + balancedness, + steps_left, ) return # ---- verbose dump ---- lines: list[str] = [ "=== EPLB dump ===", - f"step={step}, balancedness={balancedness:.2f}x", + f"model={model_name}, step={step}, avg_tokens={avg_tokens:.2f}, " + f"max_tokens={int(max_tokens)}, balancedness={balancedness:.4f}", f"Next rearrangement in {steps_left} steps", ] @@ -636,6 +638,12 @@ 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 summed 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: @@ -670,8 +678,9 @@ def step( ) if ep_group.rank() == 0: - self._dump( + self._log_balancedness( num_tokens_per_rank, + model_name=eplb_model_state.model_name, verbose=self.parallel_config.eplb_config.log_balancedness_verbose, ) diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 4bce531c7e3e..edbca1e67a5c 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -88,39 +88,3 @@ def heat_cell(text: str, val: float, vmin: float, vmax: float) -> str: 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) + "]" From 9b1ee1883d8c5038554a411bb67864ef447aaab2 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Thu, 30 Apr 2026 14:37:20 +0300 Subject: [PATCH 05/10] [EPLB] Add imbalance, global_step, async-overdue, and JSONL dump Signed-off-by: Artem Perevedentsev --- docs/serving/expert_parallel_deployment.md | 2 + vllm/config/parallel.py | 13 ++ vllm/distributed/eplb/eplb_state.py | 151 +++++++++++++++++---- 3 files changed, 140 insertions(+), 26 deletions(-) diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index fef4df770fa3..4c5479d03cf2 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -150,6 +150,8 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `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` | +| `log_balancedness_verbose` | When set with `log_balancedness`, emits a per-layer / per-rank token table to stderr instead of the one-line summary | `false` | +| `expert_load_dump_dir` | Directory to which JSONL records of per-step expert load are appended (one file per MoE model). Only EP rank 0 writes. Independent of `log_balancedness` and `log_balancedness_verbose` — setting it alone is enough | `null` | | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` | | `policy` | The policy type for expert parallel load balancing | `"default"` | diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index d02c3beeb815..e2cbfd9f6808 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -84,6 +84,19 @@ class EPLBConfig: and a per-layer / per-rank token table. Off by default to avoid large logs and extra CPU work on rank 0. """ + expert_load_dump_dir: str | None = None + """ + Directory to which JSONL records of per-step expert load are appended + on every ``log_balancedness_interval`` step (one file per MoE model, + named ``{model_name}_expert_load.jsonl``). Only EP rank 0 writes. + Independent of ``log_balancedness`` and ``log_balancedness_verbose``; + setting it alone is enough to start writing — the on-screen log and + the JSONL dump can be toggled separately or together. + + Example:: + + --eplb-config '{"expert_load_dump_dir":"./eplb_stats"}' + """ use_async: bool = False """ Whether to use non-blocking EPLB. diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index ca3b4b8617e5..eb5f72b240e0 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -26,10 +26,12 @@ physical experts. """ +import json import sys import threading from collections.abc import Sequence from dataclasses import dataclass +from pathlib import Path import torch from torch.distributed import ProcessGroup, all_reduce @@ -243,6 +245,11 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device): Otherwise, the rearrangement will hang at collective communication calls. """ + self.global_step: int = 0 + """ + Monotonic step counter that is **not** reset on rearrangement. + Used to distinguish rearrangement cycles in EPLB logs. + """ self.expert_rearrangement_step_interval: int = 0 """ Interval for expert rearrangement steps. @@ -471,6 +478,46 @@ def add_model( self.model_states[model_config.compute_hash()] = model_state self.num_valid_physical_experts = model.num_physical_experts + def _dump_expert_load( + self, + window_load_list: list[torch.Tensor], + latest_load_list: list[torch.Tensor], + ep_size: int, + ) -> None: + """Append expert-load snapshots to a per-model JSONL file.""" + dump_dir = self.parallel_config.eplb_config.expert_load_dump_dir + if dump_dir is None: + return + + dump_path = Path(dump_dir) + dump_path.mkdir(parents=True, exist_ok=True) + + for window_load, latest_load, eplb_model_state in zip( + window_load_list, latest_load_list, self.model_states.values() + ): + model = eplb_model_state.model + safe_name = eplb_model_state.model_name.replace("/", "_") + file_path = dump_path / f"{safe_name}_expert_load.jsonl" + + record = { + "model_name": eplb_model_state.model_name, + "world_size": ep_size, + "num_moe_layers": model.num_moe_layers, + "num_physical_experts": model.num_physical_experts, + "num_logical_experts": model.num_logical_experts, + "num_redundant_experts": model.num_redundant_experts, + "window_size": self.expert_load_window_size, + "step": self.global_step, + "window_expert_load": window_load.cpu().tolist(), + "latest_expert_load": latest_load.cpu().tolist(), + "physical_to_logical_map": ( + eplb_model_state.physical_to_logical_map.cpu().tolist() + ), + } + + with open(file_path, "a") as f: + f.write(json.dumps(record, separators=(",", ":")) + "\n") + def _log_balancedness( self, num_tokens_per_rank: torch.Tensor, @@ -488,29 +535,51 @@ def _log_balancedness( max_tokens = layer_maxes.sum().item() balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 - step = self.expert_rearrangement_step - steps_left = self.expert_rearrangement_step_interval - step + # Per-layer (max - mean) / mean averaged across active layers + # (matches TRT-LLM's imbalance metric). + active = layer_means > 0 + if active.any(): + active_means = layer_means[active] + active_maxes = layer_maxes[active] + imbalance = ((active_maxes - active_means) / active_means).mean().item() + else: + imbalance = 0.0 + + rearrangement_step = self.expert_rearrangement_step + steps_left = self.expert_rearrangement_step_interval - rearrangement_step + # Negative `steps_left` means async EPLB is overdue (counter not reset). + if steps_left >= 0: + schedule_line = f"steps_until_next_rearrangement={steps_left}" + else: + schedule_line = ( + f"async rearrangement overdue by {-steps_left} steps " + "(async worker is slow or dead)" + ) if not verbose: logger.info( - "EPLB step: %d for model %s: avg_tokens=%.2f, " - "max_tokens=%d, balancedness=%.4f, " - "steps until the next rearrangement: %d", - step, + "EPLB stats: model=%s, global_step=%d\n" + " rearrangement_step=%d, %s\n" + " avg_tokens=%.2f, max_tokens=%d\n" + " balancedness=%.4f, imbalance=%.4f", model_name, + self.global_step, + rearrangement_step, + schedule_line, avg_tokens, max_tokens, balancedness, - steps_left, + imbalance, ) return # ---- verbose dump ---- lines: list[str] = [ "=== EPLB dump ===", - f"model={model_name}, step={step}, avg_tokens={avg_tokens:.2f}, " - f"max_tokens={int(max_tokens)}, balancedness={balancedness:.4f}", - f"Next rearrangement in {steps_left} steps", + f"model={model_name}, global_step={self.global_step}", + f"rearrangement_step={rearrangement_step}, {schedule_line}", + f"avg_tokens={avg_tokens:.2f}, max_tokens={int(max_tokens)}", + f"balancedness={balancedness:.4f}, imbalance={imbalance:.4f}", ] num_ranks = num_tokens_per_rank.shape[1] @@ -608,8 +677,14 @@ def step( for eplb_model_state in self.model_states.values(): eplb_model_state.expert_load_pass.zero_() + # Terminal log and JSONL dump are independent — each can be enabled + # on its own; both share the same `log_balancedness_interval` cadence + # and the same one synced expert-load pass. + dump_dir_set = self.parallel_config.eplb_config.expert_load_dump_dir is not None + need_periodic_stats = log_stats or dump_dir_set + if ( - log_stats + need_periodic_stats and self.expert_rearrangement_step % self.parallel_config.eplb_config.log_balancedness_interval == 0 @@ -618,28 +693,51 @@ def step( # expert_load_pass: (num_moe_layers, num_physical_experts) expert_load_pass_list = self._sync_load_pass() ep_group = get_ep_group().device_group - for expert_load_pass, eplb_model_state in zip( - expert_load_pass_list, self.model_states.values() - ): - # num_tokens_per_rank: (num_moe_layers, num_ranks) - num_tokens_per_rank = ( - expert_load_pass.reshape( - expert_load_pass.shape[0], ep_group.size(), -1 + + if log_stats: + for expert_load_pass, eplb_model_state in zip( + expert_load_pass_list, self.model_states.values() + ): + # num_tokens_per_rank: (num_moe_layers, num_ranks) + num_tokens_per_rank = ( + expert_load_pass.reshape( + expert_load_pass.shape[0], ep_group.size(), -1 + ) + .sum(dim=-1) + .float() ) - .sum(dim=-1) - .float() - ) + if ep_group.rank() == 0: + self._log_balancedness( + num_tokens_per_rank, + model_name=eplb_model_state.model_name, + verbose=( + self.parallel_config.eplb_config.log_balancedness_verbose + ), + ) + + # All-reduce the per-model window aggregates so the dumped window + # is the global view (consistent with the synced latest_load); + # only rank 0 actually writes the files. + if dump_dir_set: + window_load_list = self._allreduce_list( + [ + s.expert_load_window.sum(dim=0) + for s in self.model_states.values() + ] + ) if ep_group.rank() == 0: - self._log_balancedness( - num_tokens_per_rank, - model_name=eplb_model_state.model_name, - verbose=self.parallel_config.eplb_config.log_balancedness_verbose, + self._dump_expert_load( + window_load_list, + expert_load_pass_list, + ep_group.size(), ) # Update the expert load sliding window if not is_dummy: - should_record = self._should_record_current_step(log_stats=log_stats) + should_record = self._should_record_current_step( + log_stats=need_periodic_stats + ) for eplb_model_state in self.model_states.values(): if should_record: eplb_model_state.expert_load_window[ @@ -657,6 +755,7 @@ def step( # rearrangement step and perform rearrangement to ensure all ranks are # performing collective communication. self.expert_rearrangement_step += 1 + self.global_step += 1 if self.is_async: # Run _move_to_workspace if all ranks have finished transferring the From 7e8395211e475bd8ac1696b83dbd60ef9bc9a905 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Thu, 30 Apr 2026 14:57:07 +0300 Subject: [PATCH 06/10] [Doc] Document log_balancedness_interval Signed-off-by: Artem Perevedentsev --- docs/serving/expert_parallel_deployment.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index 4c5479d03cf2..4e94a3f33e7f 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -150,6 +150,7 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `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` | +| `log_balancedness_interval` | Cadence (in engine steps) for the balancedness log and the JSONL dump | `1` | | `log_balancedness_verbose` | When set with `log_balancedness`, emits a per-layer / per-rank token table to stderr instead of the one-line summary | `false` | | `expert_load_dump_dir` | Directory to which JSONL records of per-step expert load are appended (one file per MoE model). Only EP rank 0 writes. Independent of `log_balancedness` and `log_balancedness_verbose` — setting it alone is enough | `null` | | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | From fe20aaaabaca800719b2b5c4619e73fc2fb4b600 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Tue, 5 May 2026 14:14:34 +0300 Subject: [PATCH 07/10] Replace print by logger.info Signed-off-by: Artem Perevedentsev --- vllm/distributed/eplb/eplb_state.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index eb5f72b240e0..af555c6e8613 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -27,7 +27,6 @@ """ import json -import sys import threading from collections.abc import Sequence from dataclasses import dataclass @@ -524,7 +523,7 @@ def _log_balancedness( model_name: str, verbose: bool = False, ) -> None: - """Print EPLB balancedness stats to stderr.""" + """Log EPLB balancedness stats.""" 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) @@ -585,7 +584,7 @@ def _log_balancedness( 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) + logger.info("\n%s", "\n".join(lines)) return # Per-layer / per-rank token table @@ -640,7 +639,7 @@ def _log_balancedness( ) lines.append("=== end EPLB dump ===") - print("\n".join(lines), file=sys.stderr, flush=True) + logger.info("\n%s", "\n".join(lines)) def step( self, From 25ce08bd163eb7374aeba80b04e454d52b0191a7 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Mon, 25 May 2026 17:42:36 +0300 Subject: [PATCH 08/10] Simplify patch; remove verbose logging to stderr with heat-map Signed-off-by: Artem Perevedentsev --- docs/serving/expert_parallel_deployment.md | 5 +- vllm/config/parallel.py | 13 +- vllm/distributed/eplb/eplb_state.py | 256 +++++++-------------- vllm/distributed/eplb/eplb_utils.py | 25 -- 4 files changed, 82 insertions(+), 217 deletions(-) diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index 4e94a3f33e7f..891531cd44f2 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -150,9 +150,8 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `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` | -| `log_balancedness_interval` | Cadence (in engine steps) for the balancedness log and the JSONL dump | `1` | -| `log_balancedness_verbose` | When set with `log_balancedness`, emits a per-layer / per-rank token table to stderr instead of the one-line summary | `false` | -| `expert_load_dump_dir` | Directory to which JSONL records of per-step expert load are appended (one file per MoE model). Only EP rank 0 writes. Independent of `log_balancedness` and `log_balancedness_verbose` — setting it alone is enough | `null` | +| `log_balancedness_interval` | How often (in engine steps) to emit the balancedness log and the JSONL dump | `1` | +| `expert_load_dump_dir` | Directory to which JSONL records of per-step expert load are appended (one file per MoE model). Only EP rank 0 writes. Has effect only when `log_balancedness` is `true`; shares the same `log_balancedness_interval` | `null` | | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | | `use_async` | Use non-blocking EPLB for reduced latency overhead | `false` | | `policy` | The policy type for expert parallel load balancing | `"default"` | diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index e2cbfd9f6808..019448fe1b26 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -77,25 +77,16 @@ 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. - """ expert_load_dump_dir: str | None = None """ Directory to which JSONL records of per-step expert load are appended on every ``log_balancedness_interval`` step (one file per MoE model, named ``{model_name}_expert_load.jsonl``). Only EP rank 0 writes. - Independent of ``log_balancedness`` and ``log_balancedness_verbose``; - setting it alone is enough to start writing — the on-screen log and - the JSONL dump can be toggled separately or together. + Has effect only when ``log_balancedness`` is ``True``. Example:: - --eplb-config '{"expert_load_dump_dir":"./eplb_stats"}' + --eplb-config '{"log_balancedness":true,"expert_load_dump_dir":"./eplb_stats"}' """ use_async: bool = False """ diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index af555c6e8613..da8177666bb8 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -49,7 +49,7 @@ from .async_worker import start_async_worker from .eplb_communicator import EplbCommunicator, create_eplb_communicator -from .eplb_utils import CpuGpuEvent, heat_cell +from .eplb_utils import CpuGpuEvent from .policy import EPLB_POLICIES, AbstractEplbPolicy, DefaultEplbPolicy from .rebalance_execute import ( AsyncEplbLayerResult, @@ -244,11 +244,6 @@ def __init__(self, parallel_config: ParallelConfig, device: torch.device): Otherwise, the rearrangement will hang at collective communication calls. """ - self.global_step: int = 0 - """ - Monotonic step counter that is **not** reset on rearrangement. - Used to distinguish rearrangement cycles in EPLB logs. - """ self.expert_rearrangement_step_interval: int = 0 """ Interval for expert rearrangement steps. @@ -479,49 +474,42 @@ def add_model( def _dump_expert_load( self, - window_load_list: list[torch.Tensor], - latest_load_list: list[torch.Tensor], + window_load: torch.Tensor, + latest_load: torch.Tensor, + eplb_model_state: "EplbModelState", ep_size: int, ) -> None: - """Append expert-load snapshots to a per-model JSONL file.""" + """Append one model's expert-load snapshot to its JSONL file.""" dump_dir = self.parallel_config.eplb_config.expert_load_dump_dir - if dump_dir is None: - return - dump_path = Path(dump_dir) dump_path.mkdir(parents=True, exist_ok=True) - for window_load, latest_load, eplb_model_state in zip( - window_load_list, latest_load_list, self.model_states.values() - ): - model = eplb_model_state.model - safe_name = eplb_model_state.model_name.replace("/", "_") - file_path = dump_path / f"{safe_name}_expert_load.jsonl" - - record = { - "model_name": eplb_model_state.model_name, - "world_size": ep_size, - "num_moe_layers": model.num_moe_layers, - "num_physical_experts": model.num_physical_experts, - "num_logical_experts": model.num_logical_experts, - "num_redundant_experts": model.num_redundant_experts, - "window_size": self.expert_load_window_size, - "step": self.global_step, - "window_expert_load": window_load.cpu().tolist(), - "latest_expert_load": latest_load.cpu().tolist(), - "physical_to_logical_map": ( - eplb_model_state.physical_to_logical_map.cpu().tolist() - ), - } - - with open(file_path, "a") as f: - f.write(json.dumps(record, separators=(",", ":")) + "\n") + model = eplb_model_state.model + safe_name = eplb_model_state.model_name.replace("/", "_") + file_path = dump_path / f"{safe_name}_expert_load.jsonl" + + record = { + "model_name": eplb_model_state.model_name, + "world_size": ep_size, + "num_moe_layers": model.num_moe_layers, + "num_physical_experts": model.num_physical_experts, + "num_logical_experts": model.num_logical_experts, + "num_redundant_experts": model.num_redundant_experts, + "window_size": self.expert_load_window_size, + "step": self.expert_rearrangement_step, + "window_expert_load": window_load.cpu().tolist(), + "latest_expert_load": latest_load.cpu().tolist(), + "physical_to_logical_map": ( + eplb_model_state.physical_to_logical_map.cpu().tolist() + ), + } + with open(file_path, "a") as f: + f.write(json.dumps(record, separators=(",", ":")) + "\n") def _log_balancedness( self, num_tokens_per_rank: torch.Tensor, model_name: str, - verbose: bool = False, ) -> None: """Log EPLB balancedness stats.""" layer_means = num_tokens_per_rank.mean(dim=1, dtype=torch.float64) @@ -530,19 +518,12 @@ def _log_balancedness( # Compute balancedness ratio: # for each layer: # (mean load across ranks) / (max load across ranks) - avg_tokens = layer_means.sum().item() - max_tokens = layer_maxes.sum().item() - balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 - - # Per-layer (max - mean) / mean averaged across active layers - # (matches TRT-LLM's imbalance metric). - active = layer_means > 0 - if active.any(): - active_means = layer_means[active] - active_maxes = layer_maxes[active] - imbalance = ((active_maxes - active_means) / active_means).mean().item() - else: - imbalance = 0.0 + active = layer_maxes > 0 + balancedness = ( + (layer_means[active] / layer_maxes[active]).mean().item() + if active.any() + else 0.0 + ) rearrangement_step = self.expert_rearrangement_step steps_left = self.expert_rearrangement_step_interval - rearrangement_step @@ -555,91 +536,13 @@ def _log_balancedness( "(async worker is slow or dead)" ) - if not verbose: - logger.info( - "EPLB stats: model=%s, global_step=%d\n" - " rearrangement_step=%d, %s\n" - " avg_tokens=%.2f, max_tokens=%d\n" - " balancedness=%.4f, imbalance=%.4f", - model_name, - self.global_step, - rearrangement_step, - schedule_line, - avg_tokens, - max_tokens, - balancedness, - imbalance, - ) - return - - # ---- verbose dump ---- - lines: list[str] = [ - "=== EPLB dump ===", - f"model={model_name}, global_step={self.global_step}", - f"rearrangement_step={rearrangement_step}, {schedule_line}", - f"avg_tokens={avg_tokens:.2f}, max_tokens={int(max_tokens)}", - f"balancedness={balancedness:.4f}, imbalance={imbalance:.4f}", - ] - - num_ranks = num_tokens_per_rank.shape[1] - if num_ranks == 0: - lines += ["No EP ranks.", "=== end EPLB dump ==="] - logger.info("\n%s", "\n".join(lines)) - 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 ===") - logger.info("\n%s", "\n".join(lines)) + logger.info( + "EPLB stats: model=%s\n rearrangement_step=%d\n %s\n balancedness=%.4f", + model_name, + rearrangement_step, + schedule_line, + balancedness, + ) def step( self, @@ -661,7 +564,7 @@ def step( log_stats (bool): If `True`, log the expert load metrics. # Stats - The metrics are summed across layers. + 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. @@ -676,14 +579,8 @@ def step( for eplb_model_state in self.model_states.values(): eplb_model_state.expert_load_pass.zero_() - # Terminal log and JSONL dump are independent — each can be enabled - # on its own; both share the same `log_balancedness_interval` cadence - # and the same one synced expert-load pass. - dump_dir_set = self.parallel_config.eplb_config.expert_load_dump_dir is not None - need_periodic_stats = log_stats or dump_dir_set - if ( - need_periodic_stats + log_stats and self.expert_rearrangement_step % self.parallel_config.eplb_config.log_balancedness_interval == 0 @@ -693,50 +590,43 @@ def step( expert_load_pass_list = self._sync_load_pass() ep_group = get_ep_group().device_group - if log_stats: - for expert_load_pass, eplb_model_state in zip( - expert_load_pass_list, self.model_states.values() - ): - # num_tokens_per_rank: (num_moe_layers, num_ranks) - num_tokens_per_rank = ( - expert_load_pass.reshape( - expert_load_pass.shape[0], ep_group.size(), -1 - ) - .sum(dim=-1) - .float() - ) - - if ep_group.rank() == 0: - self._log_balancedness( - num_tokens_per_rank, - model_name=eplb_model_state.model_name, - verbose=( - self.parallel_config.eplb_config.log_balancedness_verbose - ), - ) - - # All-reduce the per-model window aggregates so the dumped window - # is the global view (consistent with the synced latest_load); - # only rank 0 actually writes the files. - if dump_dir_set: + # Collective allreduce - ALL ranks, one time before loop. + window_load_list: list[torch.Tensor] | None = None + if self.parallel_config.eplb_config.expert_load_dump_dir is not None: window_load_list = self._allreduce_list( [ s.expert_load_window.sum(dim=0) for s in self.model_states.values() ] ) + + # Per-model: лог + JSONL, оба под if rank == 0. + for i, (expert_load_pass, eplb_model_state) in enumerate( + zip(expert_load_pass_list, self.model_states.values()) + ): if ep_group.rank() == 0: - self._dump_expert_load( - window_load_list, - expert_load_pass_list, - ep_group.size(), + num_tokens_per_rank = ( + expert_load_pass.reshape( + expert_load_pass.shape[0], ep_group.size(), -1 + ) + .sum(dim=-1) + .float() ) + self._log_balancedness( + num_tokens_per_rank, + eplb_model_state.model_name, + ) + if window_load_list is not None: + self._dump_expert_load( + window_load_list[i], + expert_load_pass, + eplb_model_state, + ep_group.size(), + ) # Update the expert load sliding window if not is_dummy: - should_record = self._should_record_current_step( - log_stats=need_periodic_stats - ) + should_record = self._should_record_current_step(log_stats=log_stats) for eplb_model_state in self.model_states.values(): if should_record: eplb_model_state.expert_load_window[ @@ -754,7 +644,6 @@ def step( # rearrangement step and perform rearrangement to ensure all ranks are # performing collective communication. self.expert_rearrangement_step += 1 - self.global_step += 1 if self.is_async: # Run _move_to_workspace if all ranks have finished transferring the @@ -1121,6 +1010,17 @@ class EplbLayerState: GPU work. """ + def set_layer_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.expert_load_view = expert_load_view[moe_layer_idx] + self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx] + self.logical_replica_count = logical_replica_count[moe_layer_idx] + def _node_count_with_rank_mapping( pg: ProcessGroup | StatelessProcessGroup, diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index ea4b280033ea..92fffd229771 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -2,9 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility functions for EPLB (Expert Parallel Load Balancing).""" -import functools import os -import sys import threading import torch @@ -116,26 +114,3 @@ 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" From af8c63ce045c6bfaf983c8ee6f9e4d1d826ec70f Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Mon, 25 May 2026 18:58:53 +0300 Subject: [PATCH 09/10] Fix mypy issue Signed-off-by: Artem Perevedentsev --- vllm/distributed/eplb/eplb_state.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 3be8220f6399..412f7a86d1d0 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -481,6 +481,9 @@ def _dump_expert_load( ) -> None: """Append one model's expert-load snapshot to its JSONL file.""" dump_dir = self.parallel_config.eplb_config.expert_load_dump_dir + assert dump_dir is not None, ( + "_dump_expert_load called without expert_load_dump_dir set" + ) dump_path = Path(dump_dir) dump_path.mkdir(parents=True, exist_ok=True) From 94f56ff94f89f09c6f1ba9c9b8191e92cedeeea1 Mon Sep 17 00:00:00 2001 From: Artem Perevedentsev Date: Fri, 29 May 2026 23:37:00 +0300 Subject: [PATCH 10/10] Enable moe_report.py Signed-off-by: Artem Perevedentsev --- vllm/distributed/eplb/eplb_state.py | 30 ++++++++++++++++++++++++++--- 1 file changed, 27 insertions(+), 3 deletions(-) diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 412f7a86d1d0..66be8277951c 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -491,17 +491,41 @@ def _dump_expert_load( safe_name = eplb_model_state.model_name.replace("/", "_") file_path = dump_path / f"{safe_name}_expert_load.jsonl" + # tokens: per-layer per-rank token counts; consumed by + # ``OFFLINE_EPLB/eplb_static/moe_report.py`` for per-step imbalance. + tokens = ( + latest_load.reshape(latest_load.shape[0], ep_size, -1) + .sum(dim=-1) + .cpu() + .tolist() + ) + # experts_per_rank: [start, end] physical-expert range per EP rank; + # used by the HTML viewer for chart sublabels. + experts_per_rank_count = model.num_physical_experts // ep_size + experts_per_rank = [ + [r * experts_per_rank_count, (r + 1) * experts_per_rank_count - 1] + for r in range(ep_size) + ] + next_rearrange = ( + self.expert_rearrangement_step_interval + - self.expert_rearrangement_step + ) record = { + "record_type": "eplb_load_stats", + "model": eplb_model_state.model_name, "model_name": eplb_model_state.model_name, - "world_size": ep_size, - "num_moe_layers": model.num_moe_layers, + "num_ranks": ep_size, + "num_layers": model.num_moe_layers, "num_physical_experts": model.num_physical_experts, "num_logical_experts": model.num_logical_experts, "num_redundant_experts": model.num_redundant_experts, + "experts_per_rank": experts_per_rank, + "next_rearrange": next_rearrange, "window_size": self.expert_load_window_size, "step": self.expert_rearrangement_step, + "tokens": tokens, + "expert_load": latest_load.cpu().tolist(), "window_expert_load": window_load.cpu().tolist(), - "latest_expert_load": latest_load.cpu().tolist(), "physical_to_logical_map": ( eplb_model_state.physical_to_logical_map.cpu().tolist() ),