From 101100e25bb7c5884478cc7097dff0f6341f4641 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 25 Mar 2026 21:21:34 +0800 Subject: [PATCH] feat(eplb): add eplb warmup-only mode Freeze expert layout after the initial EPLB warmup rebalances so heavy workloads avoid recurring rebalance memory spikes and OOMs. --- docs/advanced_features/expert_parallelism.md | 2 +- docs/advanced_features/server_arguments.md | 1 + python/sglang/srt/eplb/eplb_manager.py | 23 +++++++++++++++++++- python/sglang/srt/server_args.py | 9 ++++++++ 4 files changed, 33 insertions(+), 2 deletions(-) diff --git a/docs/advanced_features/expert_parallelism.md b/docs/advanced_features/expert_parallelism.md index 5c052114b..77dfb9b7f 100644 --- a/docs/advanced_features/expert_parallelism.md +++ b/docs/advanced_features/expert_parallelism.md @@ -141,7 +141,7 @@ SGLang introduces a dispatcher-hook system for Single-Batch Overlap (SBO), enabl SGLang integrates the [Expert Parallelism Load Balancer (EPLB)](https://github.com/deepseek-ai/EPLB) from DeepSeek to address routing imbalances in MoE models. By analyzing expert activation statistics, EPLB computes an optimal expert arrangement, strategically placing or replicating experts to minimize GPU utilization variance, reduce idle cycles, and enhance scalability. -To enable EPLB, use the flags `--enable-eplb`. For optimal performance, increase batch sizes to stabilize activation statistics and configure periodic rebalancing (e.g., every 1000 requests) to adapt to evolving workloads. Simulations demonstrate significant improvements in load balancedness (ratio of mean to max computation time), correlating strongly with throughput gains. +To enable EPLB, use `--enable-eplb`. For optimal performance, increase batch sizes to stabilize activation statistics and configure periodic rebalancing (e.g., every 1000 requests) to adapt to evolving workloads. If you only want EPLB to act during its initial warmup rebalances and then freeze the resulting expert placement, add `--eplb-warmup-only`. Simulations demonstrate significant improvements in load balancedness (ratio of mean to max computation time), correlating strongly with throughput gains. For more details, refer to the [EPLB Section in the Large-Scale EP Blog](https://lmsys.org/blog/2025-05-05-large-scale-ep/#expert-parallelism-load-balancer) and the [EPLB Repository](https://github.com/deepseek-ai/eplb). diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 61cfe91e0..6229f5d36 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -321,6 +321,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--ep-dispatch-algorithm` | The algorithm to choose ranks for redundant experts in expert parallel. | `None` | Type: str | | `--init-expert-location` | Initial location of EP experts. | `trivial` | Type: str | | `--enable-eplb` | Enable EPLB algorithm | `False` | bool flag (set to enable) | +| `--eplb-warmup-only` | Run EPLB only during its warmup rebalances, then keep expert layout static. | `False` | bool flag (set to enable) | | `--eplb-algorithm` | Chosen EPLB algorithm | `auto` | Type: str | | `--eplb-rebalance-num-iterations` | Number of iterations to automatically trigger a EPLB re-balance. | `1000` | Type: int | | `--eplb-rebalance-layers-per-chunk` | Number of layers to rebalance per forward pass. | `None` | Type: int | diff --git a/python/sglang/srt/eplb/eplb_manager.py b/python/sglang/srt/eplb/eplb_manager.py index c9ff6d925..bb77fea9a 100644 --- a/python/sglang/srt/eplb/eplb_manager.py +++ b/python/sglang/srt/eplb/eplb_manager.py @@ -18,6 +18,7 @@ class EPLBManager: super().__init__() self._model_runner = model_runner self._server_args = model_runner.server_args + self._warmup_only = self._server_args.eplb_warmup_only self._rebalance_layers_per_chunk = ( self._server_args.eplb_rebalance_layers_per_chunk ) @@ -27,7 +28,9 @@ class EPLBManager: assert ( self._server_args.eplb_rebalance_num_iterations >= self._server_args.expert_distribution_recorder_buffer_size - ), "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size" + ), ( + "eplb_rebalance_num_iterations must be greater than expert_distribution_recorder_buffer_size" + ) if not get_global_expert_distribution_recorder().recording: get_global_expert_distribution_recorder().start_record() @@ -35,6 +38,10 @@ class EPLBManager: logger.info( f"[EPLBManager] system started, will rebalance per {self._rebalance_num_iterations} iterations." ) + if self._warmup_only: + logger.info( + "[EPLBManager] warmup-only mode enabled; expert layout will become static after warmup rebalances." + ) self._main_generator = self._entrypoint() @@ -55,12 +62,26 @@ class EPLBManager: yield yield from self.rebalance() + if self._warmup_only: + self._freeze_after_warmup() + while True: + yield + while True: for _ in range(self._rebalance_num_iterations): yield yield from self.rebalance() + def _freeze_after_warmup(self): + recorder = get_global_expert_distribution_recorder() + if recorder.recording: + recorder.stop_record() + torch.get_device_module().empty_cache() + logger.info( + "[EPLBManager] warmup-only mode completed; expert layout is now static and EPLB rebalancing is disabled." + ) + def rebalance(self): logger.info("[EPLBManager] rebalance start") diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d9f1b5614..cc067aa37 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -525,6 +525,7 @@ class ServerArgs: ep_dispatch_algorithm: Optional[Literal["static", "dynamic", "fake"]] = None init_expert_location: str = "trivial" enable_eplb: bool = False + eplb_warmup_only: bool = False eplb_algorithm: str = "auto" eplb_rebalance_num_iterations: int = 1000 eplb_rebalance_layers_per_chunk: Optional[int] = None @@ -2754,6 +2755,9 @@ class ServerArgs: ), "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK (default 4096) must be larger or equal to chunked_prefill_size" def _handle_eplb_and_dispatch(self): + if self.eplb_warmup_only: + assert self.enable_eplb, "--eplb-warmup-only requires --enable-eplb." + if self.enable_eplb and (self.expert_distribution_recorder_mode is None): self.expert_distribution_recorder_mode = "stat" logger.warning( @@ -4862,6 +4866,11 @@ class ServerArgs: action="store_true", help="Enable EPLB algorithm", ) + parser.add_argument( + "--eplb-warmup-only", + action="store_true", + help="Run EPLB only during its warmup rebalances, then keep expert layout static.", + ) parser.add_argument( "--eplb-algorithm", type=str,