Elamur Research⌕
☰
home / papers / 2025 / Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models

Paper · 2025 · Zihan Qiu, Zeyu Huang, Bo Zheng et al.

Summary

Mixture-of-Experts LLMs rely on a Load-balancing Loss (LBL) so that the router spreads tokens across experts. The authors show that popular frameworks (Megatron-Core, Megablocks, Tutel and others) compute the LBL per micro-batch and then average the results. Because a micro-batch at billion-parameter scale holds only a few sequences, balance is effectively enforced within each sequence. This forces even domain-specific text, such as code or math, onto all experts equally and blocks specialization.

Their fix computes the loss at the global-batch level. Each parallel group shares its expert selection frequencies, a small vector with one entry per expert. When compute is limited and gradient accumulation is used, a buffer accumulates selection counts across accumulation steps.

They pre-trained MoE models with up to 42.8B parameters on as many as 400B tokens. Global-batch balancing reduced perplexity by roughly 0.1 and raised benchmark scores by about 2 points. Experts also developed clear domain specialization. The technique adds little overhead and was adopted in Qwen3-MoE, so it is directly relevant to anyone training MoE models.

Key points

  • Three model scales were tested: MoE-3.4A0.6B, MoE-15A2.54B and MoE-43A6.6B. They were evaluated on Hellaswag, MMLU, GSM8k, C-eval and perplexity.
  • For MoE-15A2.54B trained on 400B tokens, synchronized global balancing changed two results:
    • GSM8k rose from 48.07 to 54.28.
    • Perplexity fell from 5.778 to 5.603.
  • The auxiliary-loss-free balancing method also improved when moved from micro-batch to global-batch scope. For the 3.4A0.6B model, perplexity dropped from 8.521 to 8.080.
  • Ablations show the gains come mainly from the greater diversity of sequences in the balancing scope, not from the larger number of tokens. With suitable settings, the extra communication adds under 3% latency.
  • Under global balancing, experts show distinct domain preferences and more concentrated routing scores. For example, Chinese-language domains share a set of high-frequency experts that barely overlap with code experts. The study covers pre-training only and does not address fine-tuning or multimodal models.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Stanford UniversityStanford · United States434 works · 314,419 citations🎓University of EdinburghEdinburgh · United Kingdom359 works · 77,352 citations🏢Alibaba GroupHangzhou · China168 works · 15,386 citations

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

Cite

Cite as Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou and Junyang Lin. 2025. Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models. Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2025.acl-long.249

@inproceedings{qiu2025demons,
  title = {{Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models}},
  author = {Zihan Qiu, Zeyu Huang, Bo Zheng, Kaiyue Wen, Zekun Wang, Rui Men, Ivan Titov, Dayiheng Liu, Jingren Zhou and Junyang Lin},
  booktitle = {Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
  year = {2025},
  doi = {10.18653/v1/2025.acl-long.249},
  url = {https://doi.org/10.18653/v1/2025.acl-long.249}
}

References

37 references, 10 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  2. Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, et al. 2024. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. arXiv preprint arXiv:2401.06066. find
  3. Hantian Ding, Zijian Wang, Giovanni Paolini, Varun Kumar, Anoop Deoras, Dan Roth, and Stefano Soatto. 2024. Fewer truncations improve language modeling. arXiv preprint arXiv:2404.10830. — [arXiv:2404.10830](https://arxiv.org/abs/2404.10830) find
  4. Wenyu Du, Shuang Cheng, Tongxu Luo, Zihan Qiu, Zeyu Huang, Ka Chun Cheung, Reynold Cheng, and Jie Fu. 2024. Unlocking continual learning abilities in language models. arXiv preprint arXiv:2406.17245. find
  5. William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1–120:39. find
  6. Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. 2023. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems, 5:288–304. find
  7. Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2020. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300.
  8. Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Yao Fu, et al. 2024. C-eval: A multi-level multi-discipline chinese evaluation suite for foundation models. Advances in Neural Information Processing Systems, 36.
  9. Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, et al. 2023. Tutel: Adaptive mixture-of-experts at scale. Proceedings of Machine Learning and Systems, 5:269–287. find
  10. Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. 2024. Mixtral of experts. arXiv preprint arXiv:2401.04088. find
  11. Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. 2025. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313.
  12. Bin Lin, Zhenyu Tang, Yang Ye, Jinfa Huang, Junwu Zhang, Yatian Pang, Peng Jin, Munan Ning, Jiebo Luo, and Li Yuan. 2024. Moe-llava: Mixture of experts for large vision-language models. Preprint, arXiv:2401.15947. — [arXiv:2401.15947](https://arxiv.org/abs/2401.15947) find
  13. Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. 2024a. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437. find
  14. Liyuan Liu, Young Jin Kim, Shuohang Wang, Chen Liang, Yelong Shen, Hao Cheng, Xiaodong Liu, Masahiro Tanaka, Xiaoxia Wu, Wenxiang Hu, Vishrav Chaudhary, Zeqi Lin, Chenruidong Zhang, Jilong Xue, Hany Awadalla, Jianfeng Gao, and Weizhu Chen. 2024b. Grin: Gradient-informed moe. Preprint, arXiv:2409.121…
  15. Ka Man Lo, Zeyu Huang, Zihan Qiu, Zili Wang, and Jie Fu. 2024. A closer look into mixture-of-experts in large language models. arXiv preprint arXiv:2406.18219. find
  16. Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith,… find
  17. Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. 2024b. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.02060. find
  18. Zihan Qiu, Zeyu Huang, Shuang Cheng, Yizhi Zhou, Zili Wang, Ivan Titov, and Jie Fu. 2024a. Layerwise recurrent router for mixture-of-experts. Preprint, arXiv:2408.06793. — [arXiv:2408.06793](https://arxiv.org/abs/2408.06793) find
  19. Zihan Qiu, Zeyu Huang, and Jie Fu. 2024b. Unlocking emergent modularity in large language models. Preprint, arXiv:2310.10908. find
  20. Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. 2022. Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale. In International conference on machine learning, pages 18332–1…
  21. Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang, Xinyu Zhang, Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yuqiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. 2025. Qwen3 technical report. Preprint, arXiv:2505.09388…
  22. David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. 2024. Mixture-of-depths: Dynamically allocating compute in transformer-based language models. arXiv preprint arXiv:2404.02258. find
  23. An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. 2024. Qwen2. 5 technical report. arXiv preprint arXiv:2412.15115. find
  24. Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In 5th International Conference on Learning Representations, ICLR 2017. OpenReview.net. find
  25. Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053. find
  26. Peter T. Szymanski and Michael D. Lemmon. 1993. Adaptive mixtures of local experts are source coding solutions. In Proceedings of International Conference on Neural Networks (ICNN’88), pages 1391– 1396. IEEE. find
  27. Chameleon Team. 2024. Chameleon: Mixed-modal early-fusion foundation models. Preprint, arXiv:2405.09818.
  28. A Vaswani. 2017. Attention is all you need. Advances in Neural Information Processing Systems.
  29. Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. 2024. Auxiliary-loss-free load balancing strategy for mixture-of-experts. arXiv preprint arXiv:2408.15664. — [arXiv:2408.15664](https://arxiv.org/abs/2408.15664) find
  30. Haoze Wu, Zihan Qiu, Zili Wang, Hang Zhao, and Jie Fu. 2024. Gw-moe: Resolving uncertainty in moe router with global workspace theory. arXiv preprint arXiv:2406.12375. — [arXiv:2406.12375](https://arxiv.org/abs/2406.12375) find
  31. Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. 2024. Openmoe: An early effort on open mixture-of-experts language models. arXiv preprint arXiv:2402.01739.
  32. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang, Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang, Jing Zhou, Jin… find
  33. Chenggang Zhao, Shangyan Zhou, Liyue Zhang, Chengqi Deng, Zhean Xu, Yuxuan Liu, Kuai Yu, Jiashi Li, and Liang Zhao. 2025. Deepep: an efficient expert-parallel communication library. https://github.com/deepseek-ai/DeepEP. ↗
  34. Hao Zhao, Zihan Qiu, Huijia Wu, Zili Wang, Zhaofeng He, and Jie Fu. 2024. Hypermoe: Towards better mixture of experts via transferring among experts. arXiv preprint arXiv:2402.12656. find
  35. Zexuan Zhong, Mengzhou Xia, Danqi Chen, and Mike Lewis. 2024. Lory: Fully differentiable mixture-of-experts for autoregressive language model pre-training. Preprint, arXiv:2405.03133. find
  36. Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. 2022. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103–7114. find
  37. Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/qiu-2025-demons-detail-implementing-load-balancing/