MegaScale-MoE:生产环境中大规模混合专家模型的高效通信训练系统
机器学习
2025-10-20 v3 分布式、并行与集群计算
摘要
我们提出了 MegaScale-MoE,一个专为高效训练大规模混合专家 (MoE) 模型而量身定制的生产系统。MoE 作为一种有前景的架构,可将大型语言模型 (LLM) 扩展到前所未有的规模,从而提升模型性能。然而,现有的 MoE 训练系统面临训练效率下降的问题,且该问题因 MoE 模型规模的不断扩大和硬件的持续演进加剧。认识到高效通信在提升 MoE 训练效率中的关键作用,MegaScale-MoE 为每个 MoE 层中的注意力机制和前馈网络 (FFN) 定制了通信高效并行策略,并采用整体方法在算子间和算子内两个层面实现通信与计算的重叠。此外,MegaScale-MoE 应用了通信压缩,并调整通信模式以降低精度,进一步提高了训练效率。在 1,440 个 NVIDIA Hopper GPU 上训练一个 352B 参数的 MoE 模型时,MegaScale-MoE 实现了 1.41M tokens/s 的训练吞吐量,相比 Megatron-LM 将效率提升了 1.88 倍。我们分享了加速 MoE 训练的实践经验,并希望通过提供系统设计方面的见解,激励未来在 MoE 系统方面的研究。
引用
@article{arxiv.2505.11432,
title = {MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production},
author = {Chao Jin and Ziheng Jiang and Zhihao Bai and Zheng Zhong and Juncai Liu and Xiang Li and Ningxin Zheng and Xi Wang and Cong Xie and Qi Huang and Wen Heng and Yiyuan Ma and Wenlei Bao and Size Zheng and Yanghua Peng and Haibin Lin and Xuanzhe Liu and Xin Jin and Xin Liu},
journal= {arXiv preprint arXiv:2505.11432},
year = {2025}
}