[ENHANCEMENT] Global Batch Load Balancing for MoE Models

[Taishi-N324]

Is your feature request related to a problem? Please describe. A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

Describe the solution you'd like Implement global-batch level load balancing.

[Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models](https://arxiv.org/abs/2501.11873)

Benefits Based on the paper

• Improves pre-training perplexity by ~0.1
• Increases benchmark scores by ~2 points
• Enables interpretable domain specialization of experts

[yanring]

Thanks! Let us take a deeper look @Victarry

[Victarry]

We will add this feature to MCore v0.13. The ETA is end of this month.

[Taishi-N324]

Thank you for reviewing and accepting this feature request! I greatly appreciate your comprehensive support and the ongoing development and maintenance of Megatron-LM.

[github-actions[bot]]

Marking as stale. No activity in 60 days.

[sbhavani]

Finished with 72d23540d0358ae24a41ff289d1461b094a770fa

[Taishi-N324]

Thanks very much, I’ll try it out right away!