Add EP fake balance #961
Add EP fake balance #961
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @STwangyingrui, I'm Gemini Code Assist1! 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!
This pull request introduces an experimental feature to improve the performance of Expert Parallelism (EP) mode by implementing a "fake balance" mechanism. This new functionality, enabled via a command-line argument, aims to artificially distribute tokens more evenly across experts, leading to significant reported gains in both prefill and decoding throughput. The changes involve adding a new utility for generating balanced expert assignments and integrating this logic into the existing expert selection process.
Highlights
- New Feature: EP Fake Balance: Introduces an experimental "fake balance" mechanism for Expert Parallelism (EP) mode, controlled by the
--enable_ep_fake_balancecommand-line option. - Performance Improvement: This feature is reported to significantly increase prefill throughput by 35% and decoding throughput by 15%.
- Expert Load Balancing Logic: A new
BalancedTensorutility is added to programmatically generate and cache balanced expert assignments, ensuring tokens are distributed more evenly across experts rather than relying solely on router logits. - Integration with MoE Selection: The
select_expertsfunction in the Mixture-of-Experts (MoE) pipeline is modified to utilize the newBalancedTensorfor expert selection when the fake balance feature is enabled, overriding the default selection. - Diagnostic Logging: Conditional logging has been added to the
fused_experts_implfunction to monitor token distribution during both prefill and decoding stages when the fake balance is active. Note that decoding stage logging requires--disable_cudagraph.
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.
Footnotes
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a "fake balance" mode for Expert Parallelism to improve throughput, controlled by the --enable_ep_fake_balance flag. The core logic involves overwriting the expert selection with a balanced distribution. A critical performance issue exists in the balanced tensor generation, where an iterative approach can be replaced by a faster vectorized one. Additionally, there are opportunities for code cleanup.
lightllm/utils/balance_utils.py
Outdated
| def generate_balanced_tensor(self, length): | ||
| tensor = torch.empty((length, self.num_selected), dtype=torch.int, device="cuda") | ||
| expert_load = torch.zeros(self.num_experts, dtype=torch.int, device="cuda") | ||
|
|
||
| expert_indices = torch.arange(self.num_experts, device="cuda") | ||
|
|
||
| for i in range(length): | ||
| # To break ties randomly when loads are equal, we can shuffle indices | ||
| # of experts with the same load. A simple way is to shuffle all | ||
| # indices and then sort by load. | ||
| shuffled_indices = expert_indices[torch.randperm(self.num_experts, device="cuda")] | ||
| sorted_shuffled_indices = shuffled_indices[torch.argsort(expert_load[shuffled_indices])] | ||
|
|
||
| # Select the top `num_selected` experts with the lowest load | ||
| selected_experts = sorted_shuffled_indices[: self.num_selected] | ||
|
|
||
| tensor[i] = selected_experts | ||
|
|
||
| # Update loads for the selected experts using an efficient scatter_add | ||
| expert_load.scatter_add_(0, selected_experts, torch.ones_like(selected_experts, dtype=torch.int)) | ||
|
|
||
| return tensor |
There was a problem hiding this comment.
The current implementation of generate_balanced_tensor uses a Python for loop that iterates length times, performing CUDA operations (randperm, argsort, scatter_add_) within the loop. This sequential kernel launching can be a performance bottleneck, especially for large length during prefill. Consider a vectorized approach to generate a balanced tensor without loops for better efficiency.
def generate_balanced_tensor(self, length):
total_assignments = length * self.num_selected
num_repeats = (total_assignments + self.num_experts - 1) // self.num_experts
expert_indices = torch.arange(self.num_experts, device="cuda")
balanced_pool = expert_indices.repeat(num_repeats)[:total_assignments]
shuffled_pool = balanced_pool[torch.randperm(total_assignments, device="cuda")]
tensor = shuffled_pool.view(length, self.num_selected).to(torch.int)
return tensor
lightllm/utils/balance_utils.py
Outdated
| if self.balanced_tensors.get(num_tokens) is not None: | ||
| # logger.info(f"find balanced tensor for num_tokens={num_tokens}") | ||
| return self.balanced_tensors[num_tokens] | ||
| else: | ||
| # logger.info(f"generate balanced tensor for num_tokens={num_tokens}") | ||
| tensor = self.generate_balanced_tensor(num_tokens) | ||
| self.balanced_tensors[num_tokens] = tensor | ||
| return tensor |
There was a problem hiding this comment.
For key existence checks, if key in dict: is more idiomatic and efficient than if dict.get(key) is not None:. Also, remove the commented-out logger.info calls to keep the code clean.
if num_tokens in self.balanced_tensors:
return self.balanced_tensors[num_tokens]
tensor = self.generate_balanced_tensor(num_tokens)
self.balanced_tensors[num_tokens] = tensor
return tensor| hidden_states=hidden_states, gating_output=router_logits, topk=top_k, renormalize=renormalize | ||
| ) | ||
|
|
||
| # EP fake负载平衡开关 |
There was a problem hiding this comment.
The comment is in Chinese; translate it to English for consistency.
| # EP fake负载平衡开关 | |
| # Enable EP fake balance |
more robust balance management, from env control to option control, better logger info control, better format
STwangyingrui
force-pushed
the
yr/ep_balance
branch
from
d2141e6 to
66c0f2d
Compare
Add fake balance for EP mode, which is controled by option of --enable_ep_fake_balance.
Cost: EP8 batch128 input64 (40+ different seqlens) totally cost about 5 seconds.
Benefit: prefill throughput increase 35%, decoding throughput increase 15%, and the overheads become stable.