arXiv:2603.07685cs.DCcs.CL2026-03被引 12

Megatron Core让专家模型训练更快更省资源,支持万亿参数规模。

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

  • 通过内存、通信、计算三方面协同优化,解决专家模型训练瓶颈。
  • 在GB300/GB200上实现每GPU超1000 TFLOPS算力,支持千卡集群训练。
  • 开源可落地,适合需要训练超大规模专家模型的研究与工业团队。

扩展专家混合(MoE)模型的训练带来了密集模型中不存在的系统挑战。由于每个标记仅激活部分专家,这种稀疏性使得总参数量可远超每标记计算量,导致内存、通信和计算间形成耦合约束。优化任一维度常会加剧其他维度压力,需全系统栈协同设计。我们通过集成优化解决这些问题:内存方面采用细粒度重计算、卸载等;通信方面优化分发器并支持重叠;计算方面引入分组GEMM、融合操作及CUDA Graphs。框架还提供平行折叠支持多维并行,支持FP8与NVFP4低精度训练,并实现高效长上下文训练。在NVIDIA GB300与GB200上,对DeepSeek-V3-685B达到1,233/1,048 TFLOPS/GPU,Qwen3-235B达974/919 TFLOPS/GPU。该方案为高性能、可扩展、生产就绪的开源工具,已在学术界与工业界用于训练从数十亿到数万亿参数的MoE模型,支持上千GPU集群。本报告详解各项技术原理、权衡与系统级交互,为使用Megatron Core扩展MoE模型提供实用指导。

原文摘要 · Abstract (English)

Scaling Mixture-of-Experts (MoE) training introduces systems challenges absent in dense models. Because each token activates only a subset of experts, this sparsity allows total parameters to grow much faster than per-token computation, creating coupled constraints across memory, communication, and computation. Optimizing one dimension often shifts pressure to another, demanding co-design across the full system stack. We address these challenges for MoE training through integrated optimizations spanning memory (fine-grained recomputation, offloading, etc.), communication (optimized dispatchers, overlapping, etc.), and computation (Grouped GEMM, fusions, CUDA Graphs, etc.). The framework also provides Parallel Folding for flexible multi-dimensional parallelism, low-precision training support for FP8 and NVFP4, and efficient long-context training. On NVIDIA GB300 and GB200, it achieves 1,233/1,048 TFLOPS/GPU for DeepSeek-V3-685B and 974/919 TFLOPS/GPU for Qwen3-235B. As a performant, scalable, and production-ready open-source solution, it has been used across academia and industry for training MoE models ranging from billions to trillions of parameters on clusters scaling up to thousands of GPUs. This report explains how these techniques work, their trade-offs, and their interactions at the systems level, providing practical guidance for scaling MoE models with Megatron Core.

MoE训练系统优化千卡训练高算力

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。