Support layerwise rebalancing experts

[fzyzcjy]

Motivation

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @fzyzcjy, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello! I am gemini-code-assist, an experienced software engineer here to provide a summary of this pull request. This PR introduces support for layerwise rebalancing of experts within the Expert Parallel Load Balancing (EPLB) system. Previously, rebalancing would update the expert locations for all layers simultaneously. With this change, rebalancing can now be performed incrementally, updating only a subset of layers per rebalance iteration, controlled by a new configuration option. This is managed through a generator-based approach in the EPLBManager and requires modifications to the expert location update logic to handle updates for specific layers. New tests have been added to verify this layerwise update behavior.

Highlights

• Layerwise Expert Rebalancing: The core change is enabling the EPLB system to rebalance experts layer by layer, rather than all layers at once. This is controlled by a new argument --eplb-rebalance-layers-per-chunk.
• Generator-based Rebalancing Flow: The EPLBManager now uses a generator (_entrypoint) to manage the rebalancing process. It yields for a configured number of forward passes (eplb_rebalance_num_iterations) and then yields from the rebalance method, which itself yields between processing chunks of layers if layerwise rebalancing is enabled.
• Selective Expert Location Updates: The ExpertLocationMetadata and ExpertLocationUpdater classes have been modified to accept a list of layer IDs (update_layer_ids) and apply updates only to the specified layers. A CPU copy of the physical-to-logical map is added to facilitate this.
• New Configuration Option: A new command-line argument --eplb-rebalance-layers-per-chunk is added to ServerArgs to control how many layers are rebalanced in each step of the layerwise update.
• Updated Tests: The existing EPLB tests are refactored using a base class, and a new test case TestDynamicEPLBMultiChunk is added specifically to verify the functionality of rebalancing experts layer by layer by setting the chunk size to 1.

Changelog

Click here to see the changelog

• python/sglang/srt/managers/eplb_manager.py
  • Added List import.
  • Stored eplb_rebalance_layers_per_chunk and eplb_rebalance_num_iterations from server args.
  • Updated logging message to use stored _rebalance_num_iterations.
  • Replaced simple on_forward_pass_end logic with a generator-based _entrypoint.
  • Modified rebalance method to optionally enable timing based on chunking.
  • Implemented chunking logic for updating expert locations based on _rebalance_layers_per_chunk.
  • Called _model_runner.update_expert_location with update_layer_ids.
  • Added yield between processing layer chunks if multiple chunks exist.
  • Updated rebalance end logging message.
  • Added _compute_update_layer_ids_chunks method to determine layer chunks.
  • Added _chunk_list helper function.
• python/sglang/srt/managers/expert_location.py
  • Added physical_to_logical_map_cpu attribute to ExpertLocationMetadata.
  • Initialized physical_to_logical_map_cpu by copying the GPU tensor to CPU in _init_raw.
  • Modified update method to accept update_layer_ids.
  • Included physical_to_logical_map_cpu in the list of fields to update.
  • Changed update logic to use a mask based on update_layer_ids to update fields selectively per layer.
• python/sglang/srt/model_executor/expert_location_updater.py
  • Imported global_server_args_dict (though it seems unused in the diff).
  • Modified update method to accept update_layer_ids.
  • Passed update_layer_ids to _update_expert_weights and old_expert_location_metadata.update.
  • Modified _update_expert_weights to accept update_layer_ids.
  • Used the first layer ID from update_layer_ids to get a sample tensor for create_temp_buffers.
  • Iterated only over update_layer_ids in the update loop.
  • Used physical_to_logical_map_cpu for converting to list.
• python/sglang/srt/model_executor/model_runner.py
  • Modified update_expert_location to accept update_layer_ids.
  • Passed update_layer_ids to self.expert_location_updater.update.
  • Removed the forward_pass_id argument from the call to self.eplb_manager.on_forward_pass_end.
• python/sglang/srt/server_args.py
  • Added eplb_rebalance_layers_per_chunk attribute to ServerArgs with a default of None.
  • Added --eplb-rebalance-layers-per-chunk command-line argument.
• test/srt/test_eplb.py
  • Removed unused import ExpertDistributionStorage.
  • Renamed TestDynamicEPLB to _BaseTestDynamicEPLB and made it a base class.
  • Added extra_args class attribute to the base class.
  • Passed *cls.extra_args to the server arguments in setUpClass.
  • Added SGLANG_EXPERT_LOCATION_UPDATER_CANARY: "1" to environment variables for the server process.
  • Added TestDynamicEPLBSimple inheriting from _BaseTestDynamicEPLB.
  • Added TestDynamicEPLBMultiChunk inheriting from _BaseTestDynamicEPLB and setting extra_args to test layerwise rebalancing with a chunk size of 1.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

Layers shift and experts roam,
Rebalance finds their new home.
Chunk by chunk they settle down,
Efficiency across the town.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for layer-wise rebalancing of experts in the EPLB (Expert Placement Load Balancing) manager. This is a valuable feature as it allows the rebalancing process to be spread across multiple forward passes, potentially reducing the latency impact of a full rebalance operation.

The core changes involve modifying the EPLBManager to use generators for managing the rebalancing schedule and state, allowing it to pause and resume. The ExpertLocationMetadata and ExpertLocationUpdater have been updated to support partial updates based on specified layer IDs. Tests have also been added to cover this new chunked rebalancing functionality.

The implementation appears largely correct and follows the intended design. I have one main concern regarding a potential runtime error if a specific server argument is set to zero.

Summary of Findings

• Potential ValueError from zero chunk size: If eplb_rebalance_layers_per_chunk is set to 0, the _chunk_list helper function will receive chunk_size=0, leading to a ValueError from range(). This needs to be addressed to prevent runtime crashes.
• Accessing update_layer_ids[0] without explicit empty check: In _update_expert_weights, update_layer_ids[0] is accessed to create temporary buffers. While the current calling code in eplb_manager.py seems to prevent update_layer_ids from being empty in this context, adding a defensive check or ensuring create_temp_buffers can handle an empty list scenario could improve robustness.
• Layer-wise Rebalancing Implementation: The introduction of generator-based incremental rebalancing is a good approach to distribute the rebalancing workload over time. The changes are well-integrated across the relevant modules.

Merge Readiness

The pull request introduces a significant and useful feature for layer-wise expert rebalancing. The implementation is generally solid. However, there is one critical issue identified concerning a potential ValueError if eplb_rebalance_layers_per_chunk is set to zero, which must be addressed before merging. There's also a medium-severity point about robustness in expert_location_updater.py that would be good to consider.

Once the critical issue is resolved, and the medium severity point is considered, this PR should be in a much better state for merging. As an AI assistant, I am not authorized to approve pull requests; please ensure further review and approval from team members.

[gemini-code-assist]

The chunk_size for _chunk_list is determined by self._rebalance_layers_per_chunk or 1000000. If self._server_args.eplb_rebalance_layers_per_chunk is explicitly set to 0 (e.g., via the --eplb-rebalance-layers-per-chunk CLI argument), chunk_size will become 0.

Passing chunk_size=0 to range(0, len(items), chunk_size) in the _chunk_list function will raise a ValueError: range() arg 3 must not be zero.

Could we add a check to ensure eplb_rebalance_layers_per_chunk is positive if provided, or handle the chunk_size=0 case gracefully in _chunk_list or here to prevent this runtime error?

[gemini-code-assist]

Accessing update_layer_ids[0] assumes that update_layer_ids will never be an empty list when _update_expert_weights is called. While current logic in eplb_manager.py seems to prevent this by not iterating over empty update_layer_ids_chunks, it might be safer to add a check here or ensure that _compute_update_layer_ids_chunks never produces empty lists within its output list of lists if all_layer_ids was non-empty.

For instance, if update_layer_ids could somehow be empty, this would lead to an IndexError. Could we add a guard, for example, checking if update_layer_ids is non-empty before proceeding, or ensure that create_temp_buffers can handle being called with an empty list if that's a valid state (though it seems sample_tensors is expected to be non-empty by create_temp_buffers)?