diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index fef4df770fa3..891531cd44f2 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_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 6e8f78e4ee30..51da72984ae0 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -77,6 +77,17 @@ class EPLBConfig: """ Interval for logging the balancedness. """ + 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. + Has effect only when ``log_balancedness`` is ``True``. + + Example:: + + --eplb-config '{"log_balancedness":true,"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 319a5f22c922..66be8277951c 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -26,9 +26,11 @@ physical experts. """ +import json import threading from collections.abc import Sequence from dataclasses import dataclass +from pathlib import Path import torch from torch.distributed import ProcessGroup, all_reduce @@ -470,6 +472,105 @@ 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: torch.Tensor, + latest_load: torch.Tensor, + eplb_model_state: "EplbModelState", + ep_size: int, + ) -> 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) + + model = eplb_model_state.model + 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, + "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(), + "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, + ) -> None: + """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) + + # Compute balancedness ratio: + # for each layer: + # (mean load across ranks) / (max load across ranks) + 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 + # 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)" + ) + + 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, is_dummy: bool = False, @@ -515,44 +616,39 @@ 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 - ) - .sum(dim=-1) - .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 + # 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() + ] + ) + for i, (expert_load_pass, eplb_model_state) in enumerate( + zip(expert_load_pass_list, self.model_states.values()) + ): 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, + 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, - avg_tokens, - max_tokens, - balancedness, - self.expert_rearrangement_step_interval - - self.expert_rearrangement_step, ) + 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: