Forced load imbalance#2917
Conversation
| 'The default value 1e-3 is same as that used in DeepSeekV3.') | ||
| group.add_argument('--moe-router-force-load-balancing', action='store_true', | ||
| help='[Experimental] Force override routing to balance token distribution using random logits for MoE routers, supporting naive top-k and group-limited top-k. This experimental feature is for benchmarking purposes only!') | ||
| group.add_argument('--moe-router-force-biased', type=float, default=None, |
There was a problem hiding this comment.
Overall LGTM. I think the argument name and help message could be refined to make it more clear to users.
Suggested changes:
group.add_argument('--moe-router-force-expert-bias-std', type=float, default=None,
help='[Experimental] Apply random expert bias in normal distribution with specified std to router logits. The random seeds are shared across all ranks.'
'If the specified value is positive, generates new random bias each forward pass. '
'If the specified value is negative, generates bias once per layer and reuses it (abs value is std).'
'This experimental feature is for benchmarking purposes only!')
There was a problem hiding this comment.
Thanks. The help message has been updated.
|
/ok to test 9c46233 |
|
Thank you, this is very useful for MLSys debugging. |
|
The lint issue needs to be fixed. |
Head branch was pushed to by a user without write access
|
/ok to test 1689973 |
|
/ok to test 38a11bb |
|
Thank you for your contribution! NVIDIA Megatron-LM is currently transitioning to development on Github. We will aim to review your PR after we complete our transition and stabilize our Github development process. Thank you for your understanding. |
| """[Experimental] Force load balancing with random logits for MoE router, supports naive topk | ||
| and group-limited topk. This is an experimental feature and only for benchmark.""" | ||
|
|
||
| moe_router_force_biased: Optional[float] = None |
There was a problem hiding this comment.
One API issue is that can the be fused together as an enum string parsing with the current Boolean value for force load balancing? Or at least have some though going into other force routing strategies for benchmarking, for instance, you may want to try the worst case of all experts being biased to the same GPU group to test worst case memory scenarios.
|
Communicated with Nan. Nan will find some time to file the main PR in the near future. |
Co-authored-by: Dennis(Zhenhuan) Liu <denliu@nvidia.com>
Co-authored-by: Dennis(Zhenhuan) Liu <denliu@nvidia.com>
What does this PR do ?
This PR tries to add a utility to force load imbalance among experts.
Contribution process
flowchart LR A[Pre-checks] --> B[PR Tests] subgraph Code Review/Approval C1[Expert Review] --> C2[Final Review] end B --> C1 C2 --> D[Merge]Pre-checks
Core 0.8)Code review
The following process is enforced via the CODEOWNERS file for changes into
megatron/core. For changes outside ofmegatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.For MRs into `main` branch
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
(Step 1): Add PR label
Expert Review(Step 2): Collect the expert reviewers reviews
Expert Reviewlabel when your PR is ready for review.Final Review might get declined if these requirements are not fulfilled.
(Step 3): Final Review
Final Reviewlabel(Optional Step 4): Cherry-pick into release branch
If this PR also needs to be merged into
core_r*release branches, after this PR has been merged, selectCherry-pickto open a new PR into the release branch.For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.Merging your PR
Any member of core-adlr and
core-nemowill be able to merge your PR.