Scalable Training of Mixture-of-Experts Models with Megatron Core

扩展混合专家（MoE）模型的训练会带来密集型模型所不具备的一系列系统级挑战。由于每个输入词元（token）仅激活全部专家中的一部分，这种稀疏性使得模型总参数量的增长速度远超单个词元所需的计算量，从而在内存、通信与计算三者之间形成紧密耦合的约束关系。优化其中任一维度，往往会导致压力向其他维度转移，因此亟需对整个系统栈进行协同设计（co-design）。 我们通过覆盖内存（如细粒度重计算、显存卸载等）、通信（如优化的专家调度器、计算与通信重叠等）以及计算（如分组矩阵乘法 GEMM、算子融合、CUDA Graph 等）三大层面的集成化优化，系统性地应对 MoE 训练中的上述挑战。该框架还支持“并行折叠”（Parallel Folding），实现灵活的多维并行；提供对 FP8 和 NVFP4 低精度训练的原生支持；并具备高效的长上下文训练能力。在 NVIDIA GB300 和 GB200 硬件平台上，其针对 DeepSeek-V3-685B 模型分别达到 1,233 / 1,048 TFLOPS/卡的实测计算吞吐，针对 Qwen3-235B 模型则分别达到 974 / 919 TFLOPS/卡。作为一款高性能、高可扩展性且面向生产环境就绪的开源解决方案，该框架目前已广泛应用于学术界与工业界，在规模达数千张 GPU 的集群上成功训练了参数量从数十亿到数万亿不等的 MoE 模型。 本报告深入阐释了上述各项技术的工作原理、各自权衡取舍（trade-offs）及其在系统层面的相互作用机制，为基于 Megatron Core 构建和扩展 MoE 模型提供了切实可行的工程实践指南。

大规模MoE模型训练面临内存、通信与计算三者强耦合的系统性瓶颈：专家稀疏激活虽降低单token计算量，却加剧显存分配不均、All-to-All通信拥塞、GPU利用率波动等问题；传统密集模型优化技术（如标准梯度检查点、静态数据并行）无法直接迁移，亟需面向MoE特性的全栈协同设计。

提出首个面向生产级MoE训练的全栈协同优化框架——以Parallel Folding实现专家并行、数据并行、序列并行与张量并行的动态解耦与融合；通过细粒度重计算+专家卸载+Dispatch-aware通信重叠打破内存-通信-计算三角约束；首创Grouped GEMM与NVFP4/FP8原生支持的混合精度MoE内核，并深度集成CUDA Graphs消除内核启动开销。

在GB300/GB200集群上实现DeepSeek-V3-685B达1233 TFLOPS/GPU（92%硬件峰值）、Qwen3-235B达974 TFLOPS/GPU；开源Megatron-Core MoE扩展模块（已集成至NVIDIA官方仓库）；支持万卡级训练，实测千卡集群线性扩展效率>94%；提供长上下文（>128K）MoE专属优化；代码、配置模板、性能分析工具链全部开源（GitHub: megatron-core/moe）。

1. 'Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity' (Google, 2021); 2. 'GLaM: Efficient Scaling of Language Models with Mixture of Experts' (Google, 2022); 3. 'FastMoE: A Fast Mixture-of-Experts Training System' (Tsinghua, 2022); 4. 'DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training' (Microsoft, 2023); 5. 'Colossal-AI MoE: Memory-Efficient Large-Scale MoE Training' (PKU, 2023)