高效训练超大规模专家混合模型,提速近两倍。
MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production
- 定制通信优化策略,实现注意力与前馈层并行计算
- 在1440张H100上训练3520亿参数模型,吞吐达141万词元/秒
- 适合大规模AI系统研发人员参考生产级优化方案
我们提出MegaScale-MoE,一个面向生产环境的大型专家混合(MoE)模型高效训练系统。尽管MoE架构能显著提升大语言模型规模与性能,但现有训练系统在模型与硬件持续扩大的背景下,训练效率持续下降。MegaScale-MoE针对通信瓶颈,为每个MoE层的注意力和前馈网络定制通信高效的并行策略,并在跨操作符与内操作符层面实现通信与计算重叠。此外,通过调整通信模式并应用降精度压缩,进一步提升训练效率。在1,440张NVIDIA Hopper GPU上训练3520亿参数的MoE模型时,实现1.41M tokens/s的训练吞吐量,相比Megatron-LM提升1.88倍。我们分享了实际部署经验,期望推动未来MoE系统研究。
原文摘要 · Abstract (English)
We present MegaScale-MoE, a production system tailored for the efficient training of large-scale mixture-of-experts (MoE) models. MoE emerges as a promising architecture to scale large language models (LLMs) to unprecedented sizes, thereby enhancing model performance. However, existing MoE training systems experience a degradation in training efficiency, exacerbated by the escalating scale of MoE models and the continuous evolution of hardware. Recognizing the pivotal role of efficient communication in enhancing MoE training, MegaScale-MoE customizes communication-efficient parallelism strategies for attention and FFNs in each MoE layer and adopts a holistic approach to overlap communication with computation at both inter- and intra-operator levels. Additionally, MegaScale-MoE applies communication compression with adjusted communication patterns to lower precision, further improving training efficiency. When training a 352B MoE model on 1,440 NVIDIA Hopper GPUs, MegaScale-MoE achieves a training throughput of 1.41M tokens/s, improving the efficiency by 1.88$\times$ compared to Megatron-LM. We share our operational experience in accelerating MoE training and hope that by offering our insights in system design, this work will motivate future research in MoE systems.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。