arXiv:2504.14960cs.LGcs.DC2025-04被引 18

提出新型混合并行策略,让大模型训练更快更高效。

MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core

  • 将注意力与专家层分开并行,各自选最优配置。
  • 在H100上实现49.3%模型浮点利用率,支持1024卡扩展。
  • 适合大规模MoE模型训练者,尤其关注效率的团队。

混合专家(MoE)模型通过动态选择相关专家,提升神经网络可扩展性,可在保持计算成本可控的前提下实现更大规模模型。然而,在数千张GPU上高效训练大规模MoE模型仍面临现有并行策略的限制。本文提出一个端到端的大规模MoE模型训练框架,采用五维混合并行:张量并行、专家并行、上下文并行、数据并行和流水线并行。核心是提出的MoE并行折叠策略,将Transformer中注意力层与MoE层的并行解耦,使每类层可独立选择最优并行配置。同时,设计灵活的基于令牌的调度器,支持跨五维并行的丢弃与不丢弃令牌训练,适应动态张量形状,并协调不同并行方案。实验表明,该框架显著提升训练效率与可扩展性。在H100上,Mixtral 8x22B模型达到最高49.3%模型浮点利用率(MFU),Qwen2-57B-A14B模型达39.0% MFU,优于现有方法。框架可高效扩展至1024张GPU,支持长达128K tokens的序列,验证其在大规模MoE模型训练中的有效性。代码已开源于Megatron-Core。

原文摘要 · Abstract (English)

Mixture of Experts (MoE) models enhance neural network scalability by dynamically selecting relevant experts per input token, enabling larger model sizes while maintaining manageable computation costs. However, efficient training of large-scale MoE models across thousands of GPUs presents significant challenges due to limitations in existing parallelism strategies. We introduce an end-to-end training framework for large-scale MoE models that utilizes five-dimensional hybrid parallelism: Tensor Parallelism, Expert Parallelism, Context Parallelism, Data Parallelism, and Pipeline Parallelism. Central to our approach is MoE Parallel Folding, a novel strategy that decouples the parallelization of attention and MoE layers in Transformer models, allowing each layer type to adopt optimal parallel configurations. Additionally, we develop a flexible token-level dispatcher that supports both token-dropping and token-dropless MoE training across all five dimensions of parallelism. This dispatcher accommodates dynamic tensor shapes and coordinates different parallelism schemes for Attention and MoE layers, facilitating complex parallelism implementations. Our experiments demonstrate significant improvements in training efficiency and scalability. We achieve up to 49.3% Model Flops Utilization (MFU) for the Mixtral 8x22B model and 39.0% MFU for the Qwen2-57B-A14B model on H100 GPUs, outperforming existing methods. The framework scales efficiently up to 1,024 GPUs and maintains high performance with sequence lengths up to 128K tokens, validating its effectiveness for large-scale MoE model training. The code is available in Megatron-Core.

MoE模型并行训练大模型加速

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