[MoE] More balanced expert sharding #21497
Conversation
Signed-off-by: Woosuk Kwon <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels. Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can either: Add 🚀 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a more balanced expert sharding strategy when the number of experts is not evenly divisible by the expert parallelism size. The previous implementation assigned all remaining experts to the last rank, which could create a performance bottleneck. The new implementation distributes the remainder experts one by one to the initial ranks, resulting in a much more even distribution. The logic is correct, clear, and improves performance by balancing the workload. I have no further comments and approve these changes.
| # Distribute experts as evenly as possible to each rank. | ||
| base_experts = global_num_experts // ep_size | ||
| remainder = global_num_experts % ep_size | ||
| if ep_rank < remainder: | ||
| local_num_experts = base_experts + 1 | ||
| else: | ||
| local_num_experts = base_experts | ||
|
|
||
| # Create a tensor of size num_experts filled with -1 | ||
| expert_map = torch.full((global_num_experts, ), -1, dtype=torch.int32) | ||
| # Create a expert map for the local experts | ||
| if ep_rank < (ep_size - 1): | ||
| # Each non-last rank gets local_num_experts experts. | ||
| expert_map[ep_rank * local_num_experts: | ||
| (ep_rank + 1) * local_num_experts] = \ | ||
| torch.arange(0, local_num_experts, dtype=torch.int32) | ||
| else: | ||
| # All remaining experts are assigned to the last rank. | ||
| local_num_experts = (global_num_experts - ep_rank * local_num_experts) | ||
|
|
||
| expert_map[-local_num_experts:] = \ | ||
| torch.arange(0, local_num_experts, dtype=torch.int32) | ||
| start_idx = ep_rank * base_experts + min(ep_rank, remainder) | ||
| expert_map[start_idx:start_idx + local_num_experts] = torch.arange( | ||
| 0, local_num_experts, dtype=torch.int32) |
There was a problem hiding this comment.
Suggestion for more concise version:
# Distribute experts as evenly as possible to each rank.
base_experts = global_num_experts // ep_size
remainder = global_num_experts % ep_size
start_idx = lambda r: r * base_experts + min(r, remainder)
# Create a tensor of size num_experts filled with -1
expert_map = torch.full((global_num_experts, ), -1, dtype=torch.int32)
# Create a expert map for the local experts
expert_map[start_idx(ep_rank):start_idx(ep_rank + 1)] = torch.arange(
0, local_num_experts, dtype=torch.int32)There was a problem hiding this comment.
Actually I'm not sure if the suggested code is simpler. 😅 We need to compute local_num_experts anyways because it's one of the return values.
There was a problem hiding this comment.
@WoosukKwon my bad, I missed that we also return that value.
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: x22x22 <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Jinzhen Lin <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Paul Pak <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]> Signed-off-by: Diego-Castan <[email protected]>
Signed-off-by: Woosuk Kwon <[email protected]>
Currently, when
num_expertsis not divisible evenly byep_size, the experts are distributed unevenly among ranks. For instance, if there are 30 experts and anep_sizeof 8, the existing logic results in an unbalanced distribution:[3, 3, 3, 3, 3, 3, 3, 9], assigning all remainder experts to the last rank.This PR addresses this issue by distributing the experts in a more balanced manner. After applying this fix, the experts are distributed as
[4, 4, 4, 4, 4, 4, 3, 3].