Support picking variants of EPLB algorithms

[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 team, gemini-code-assist here to provide a summary of this pull request. This PR introduces support for selecting different variants of the Expert Parallelism Load Balancing (EPLB) algorithms, specifically focusing on the DeepSeek EPLB methods. Previously, the system might have implicitly used a specific variant or selection logic. With these changes, users can now explicitly choose an algorithm via a new command-line argument, or let the system automatically select one based on the configuration. This enhances flexibility and allows for experimentation with different load balancing strategies.

Highlights

• EPLB Algorithm Selection: Introduces the ability to select from multiple EPLB algorithm variants (DeepSeek, DeepSeek Hierarchical, DeepSeek Vec, DeepSeek Vec Hierarchical) via a new configuration option.
• New EPLB Algorithms Package: Adds a new Python package sglang.srt.managers.eplb_algorithms to house the implementations of different EPLB strategies, copied from the DeepSeek EPLB repository.
• Automatic Algorithm Selection: Adds an 'auto' mode for algorithm selection that chooses a default hierarchical or non-hierarchical DeepSeek algorithm based on the relationship between num_groups and num_nodes.
• Configuration Option: Adds a new command-line argument --eplb-algorithm to server_args.py to allow users to specify the desired EPLB algorithm.

Changelog

Click here to see the changelog

• python/sglang/srt/managers/eplb_algorithms/init.py
  • New file: Defines the EplbAlgorithm enum for different variants.
  • New file: Implements a dispatcher function rebalance_experts that calls the appropriate algorithm based on the EplbAlgorithm enum.
  • New file: Implements compute_algorithm to select an algorithm based on a raw string input ('auto' or specific name) and system parameters.
• python/sglang/srt/managers/eplb_algorithms/deepseek.py
  • New file: Contains the implementation for the deepseek and deepseek_hierarchical EPLB algorithms, including helper functions like balanced_packing and replicate_experts.
• python/sglang/srt/managers/eplb_algorithms/deepseek_vec.py
  • New file: Contains the implementation for the deepseek_vec and deepseek_vec_hierarchical EPLB algorithms.
  • Modified the rebalance_experts function signature to accept an enable_hierarchical boolean instead of a phase literal.
• python/sglang/srt/managers/expert_location.py
  • Updated import to use the new eplb_algorithms package.
  • Modified the call to the rebalance_experts function to use the new dispatcher from eplb_algorithms.__init__.py and pass the algorithm determined by compute_algorithm.
• python/sglang/srt/server_args.py
  • Added a new field eplb_algorithm to the ServerArgs dataclass with a default value of 'auto'.
  • Added a corresponding command-line argument --eplb-algorithm to the argument parser.

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.

---

Algorithms abound,
For experts to be found.
Choose the right path,
Avoid load's wrath.

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 selecting different EPLB (Expert Parallelism Load Balancing) algorithm variants, which is a valuable enhancement for optimizing MoE models. The changes involve creating a new eplb_algorithms package to house the logic, including an Enum for algorithm types, a dispatcher function, and a computation function to determine the algorithm. The existing deepseek_vec.py is updated, and a new deepseek.py (copied from an external source) is added.

The overall structure is logical. However, there are a few critical and medium-severity issues related to error handling, type safety with Optional types, and potential performance concerns that should be addressed.

Summary of Findings

• Robustness of compute_algorithm: The compute_algorithm function in python/sglang/srt/managers/eplb_algorithms/__init__.py can raise a TypeError if num_groups is None during 'auto' mode, and a KeyError if an invalid raw_algorithm string is provided. This was commented on with critical severity.
• Handling of num_groups=None for Hierarchical Algorithms: The rebalance_experts dispatcher in python/sglang/srt/managers/eplb_algorithms/__init__.py might pass num_groups=None to underlying hierarchical algorithms if num_groups is None in the model configuration, potentially leading to errors. This was commented on with high severity.
• Missing Return Type Hint: The rebalance_experts function in python/sglang/srt/managers/eplb_algorithms/__init__.py is missing a return type hint. This was commented on with medium severity.
• Potential Performance Bottleneck: The deepseek.py algorithm moves tensors to the CPU (e.g., weight.cpu()). This could be a performance issue for large tensors on GPU. This was commented on with medium severity.

Merge Readiness

The pull request introduces a good architectural change for selecting EPLB algorithms. However, there are critical and high-severity issues related to error handling and type safety, particularly around the num_groups parameter when it's None. These issues could lead to runtime errors.

I recommend that these critical and high-severity issues be addressed before merging. The medium-severity issues (missing type hint, potential performance concern with CPU transfers) should also be considered for improvement.

As an AI reviewer, I am not authorized to approve pull requests. Please ensure further review and approval from other maintainers after addressing the feedback.

[gemini-code-assist]

This function has a couple of potential issues:

1. TypeError with num_groups: If raw_algorithm == "auto" and num_groups is None (which it can be, as model_config_for_expert_location.num_groups is Optional), the expression num_groups % num_nodes on line 60 will raise a TypeError. The type hint for num_groups here is int, but it can receive None from the caller in expert_location.py.
2. KeyError for invalid raw_algorithm: If raw_algorithm is not "auto" and not a valid name of an EplbAlgorithm member, EplbAlgorithm[raw_algorithm] on line 57 will raise a KeyError.

Could we make this function more robust by handling these cases? For example, by checking if num_groups is None before the modulo operation in "auto" mode, and by using a try-except KeyError block or getattr for the explicit algorithm lookup?

[gemini-code-assist]

The num_groups parameter is Optional[int]. If a hierarchical algorithm (e.g., EplbAlgorithm.deepseek_hierarchical) is selected (either explicitly by the user or by the compute_algorithm function) and num_groups happens to be None (e.g., from model config), passing num_groups=None to the underlying deepseek.rebalance_experts (line 29) or deepseek_vec.rebalance_experts (line 43) will likely cause a TypeError or an AssertionError in those functions, as they expect num_groups to be an int for hierarchical variants (e.g., deepseek.rebalance_experts_hierarchical takes num_groups: int).

Should we add a check here to ensure num_groups is not None before calling these functions if a hierarchical algorithm is active? Or perhaps compute_algorithm should guarantee that it doesn't return a hierarchical algorithm if num_groups is unsuitable (e.g. None) even for explicit user choices?

[gemini-code-assist]

The function rebalance_experts is missing a return type hint. Based on the underlying implementations (deepseek.rebalance_experts and deepseek_vec.rebalance_experts), it seems to return a Tuple[torch.Tensor, torch.Tensor, torch.Tensor].

Would it be beneficial to add this type hint for better code clarity and static analysis?

23:    algorithm: EplbAlgorithm,
24:) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

[gemini-code-assist]

The weight tensor is explicitly moved to the CPU using .cpu() here (and similarly in balanced_packing for indices and pack_index). If weight is a large tensor residing on the GPU, this transfer could introduce a performance bottleneck.

Is this CPU processing intentional and necessary for the correctness of the EPLB algorithm, or could these operations be performed on the GPU to avoid data transfer overhead?