MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core
MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core
Section titled “MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core”⚠️ AI 生成 · 建议对照原文 本页为自动整理的学习笔记；关键数据与引用如需引用，请回查 PDF / 官方版本。
学习档位 中文笔记
类型 文献 · 更新 2026-07-19
中文学习笔记（自动生成，需核验）
Section titled “中文学习笔记（自动生成，需核验）”Topic: training-systems · Tier: recent · Year: 2025 · Venue: —
Evidence level: partial · 本地全文: 是 · 建议阅读: ~45 分钟
Paper: https://arxiv.org/abs/2504.14960
Code: —
Generator: grok
为什么值得读
Section titled “为什么值得读”本文提出面向大规模MoE的5D混合并行框架与MoE Parallel Folding，解耦Attention与MoE层并行映射，配合灵活token dispatcher，显著提升训练效率与可扩展性（H100上高MFU），代码开源至Megatron-Core，对具身/自动驾驶相关大模型训练工程实践极具参考价值。
通过MoE Parallel Folding实现Attention与MoE异构并行映射，结合5D混合并行与高效token dispatcher，在大规模MoE训练中达到更高MFU与更好扩展性。
大规模MoE模型在数千GPU上高效训练困难：现有单一或简单混合并行可扩展性不足；MoE稀疏导致计算/参数比低，小模型并行易OOM、大并行通信开销大；Attention与MoE层计算/通信特性不同，统一并行映射次优；动态token路由与可变张量形状进一步增加实现复杂度。
Transformer架构与MoE基本原理（Top-K gating、capacity factor）；分布式训练并行策略（TP、EP、CP、DP、PP及ZeRO）；All-to-All/AllGather等集合通信；Megatron-LM/Megatron-Core框架基础。
- 提出MoE Parallel Folding：首次解耦Transformer中Attention与MoE层的并行化策略，使各自可采用最优配置，并将通信密集型维度折叠以更好利用高带宽节点内网络。
- 设计灵活高效的token-level dispatcher：支持token-dropping与token-dropless训练，兼容5D混合并行（TP/EP/CP/DP/PP），处理动态张量形状并协调Attention与MoE不同并行方案。
- 实验验证性能提升：在H100上Mixtral 8x22B达49.3% MFU、Qwen2-57B-A14B达39.0% MFU，扩展至1024 GPU且支持长至128K序列。
构建基于TP×EP×CP×DP×PP的5D混合并行端到端框架。对Attention层形成TP×CP×DP×PP组；对MoE层形成ETP×EP×EDP×PP组（PP组数量与成员一致）。MoE Parallel Folding允许独立最优映射并将通信折叠到更紧凑组。Token dispatcher负责路由：router按Top-K分配→permutation使同expert token连续→All-to-All-V跨EP交换→AllGather-V在ETP共享激活→专家FFN计算→ReduceScatter-V→All-to-All-V返回→un-permutation恢复顺序；支持sub-sequence dropping作为默认token-drop策略；反向对应通信替换。
关键模块和设计取舍
Section titled “关键模块和设计取舍”- MoE Parallel Folding：解耦Attention（密集序列级）与MoE（token级稀疏）并行，允许用更通信高效的EP替代ETP并与Attention的TP折叠，限制仅PP一致；取舍是牺牲完全统一映射换取灵活性与节点内通信优化。2) Token dispatcher：统一处理ETP+EP，支持drop/dropless与动态shape；选择sub-sequence dropping而非full-sequence以降低跨rank logits通信与负载不均，经验上不影响收敛。3) 专家并行流程：dispatch→compute→restore，依赖All-to-All-V等。设计优先通信效率与正确性而非最简实现。
数据集、实验设置与指标
Section titled “数据集、实验设置与指标”实验在Eos集群（DGX H100节点，NVLink+IB，PyTorch 2.5/CUDA 12.6，BF16）进行，最多1024 GPU。模型：粗粒度Mixtral 8x22B与upcycle的Llama3-8x70B；细粒度Qwen2-57B-A14B与reparam的Mixtral-8x22B-G8T8。主要指标MFU（Model Flops Utilization，对比BF16理论峰值）与TFLOPS；benchmark用token-drop（capacity factor=1）减轻负载不均抖动；全局batch size=256。基线包括FSDP、FSDP+EP、TP+EP+DP、MCore 5D。
主要结果（需原文证据）
Section titled “主要结果（需原文证据）”Abstract与Table 1：MCore w/ Folding在Mixtral-8x22B（128 GPU）达49.3% MFU、Llama3-8x70B（256 GPU）41.6%、Qwen2-57B-A14B（64 GPU）39.0%、Mixtral-8x22B-G8T8（128 GPU）28.8%，均优于FSDP（低至<10%或OOM）、FSDP+EP、TP+EP+DP与无Folding的MCore。支持扩展至1024 GPU与序列长度至128K tokens。强扩展结果见图3（具体曲线细节待来源核验）。
局限、失败场景与适用边界
Section titled “局限、失败场景与适用边界”PP组数量与成员在Attention/MoE间必须一致；token-drop采用sub-sequence可能与full-sequence有细微差异（虽经验无显著收敛影响）；FSDP类基线在大模型易OOM或通信难重叠；细粒度MoE（如G8T8）MFU相对较低。完整失败场景、更长序列或更多模型细节及消融待来源核验（摘录截断于实验部分）。适用边界：主要针对Transformer-MoE在H100类集群的5D混合并行训练。
与前序 / 同期 / 后续方法的关系
Section titled “与前序 / 同期 / 后续方法的关系”MoE基础：GShard、Switch Transformer、GLaM等稀疏扩展；token-drop vs Megablocks dropless；细粒度如DeepSeek-MoE。分布式：传统TP/DP/CP/PP/ZeRO；MoE专用EP；混合如FSDP+EP、TP+EP+DP等。本文在MCore 5D基础上引入Folding解耦与统一dispatcher，超越先前统一映射或受限EP⊂DP的做法。
官方代码与复现建议
Section titled “官方代码与复现建议”官方代码在Megatron-Core（https://github.com/NVIDIA/Megatron-LM）。复现建议：使用H100集群、BF16、指定模型配置与5D并行+Folding参数；benchmark时设capacity=1、GBS=256；调优各维度并行度找最优MFU。具体脚本与完整配置待来源核验。
推荐阅读顺序
Section titled “推荐阅读顺序”先读Abstract与Introduction（问题与贡献）→Section 3.1 Preliminary（MoE/EP基础）→3.2 MoE Parallel Folding与Figure 1→3.3 Token Dispatcher与Figure 2→Section 4.1-4.2实验设置与Table 1/Figure 3→Related Work→结论与代码链接。重点精读Folding与dispatcher设计。
- Q: MoE Parallel Folding的核心思想是什么？它如何设置Attention与MoE的并行组？ A: 解耦Attention与MoE层的并行映射，使各自最优。Attention：TP×CP×DP×PP；MoE：ETP×EP×EDP×PP（仅要求PP组一致）。
- Q: 为什么统一并行映射对MoE次优？Folding带来哪两个主要好处？ A: Attention是序列级密集计算（适合TP/CP），MoE是token级稀疏（更适合EP）。好处：1）独立选最优映射（如用EP替ETP）；2）折叠使通信落入高带宽节点内网络，降低开销。
- Q: Token dispatcher如何支持token-dropping？为何默认sub-sequence dropping？ A: Router按Top-K，drop时可用full-sequence（需gather logits，通信大）或sub-sequence（仅本地sub-seq决策）。默认后者因通信少、缓解负载不均，且经验上不损害收敛。
- Q: Table 1中MCore w/ Folding在主要模型上的MFU结果是什么？相对基线如何？ A: Mixtral-8x22B 49.3%、Llama3-8x70B 41.6%、Qwen2-57B-A14B 39.0%、Mixtral-8x22B-G8T8 28.8%（对应GPU数见文），优于FSDP、FSDP+EP、TP+EP+DP及无Folding MCore。
- Q: 5D混合并行包含哪些维度？dispatcher如何协调Attention与MoE的不同切分？ A: TP、EP、CP、DP、PP。Attention输出按batch(DP)或序列(CP/TP)切分，但MoE按token独立计算，故dispatcher用统一workflow（permutation + All-to-All-V + AllGather-V等）处理任意Attention映射。
- Abstract (page 1): 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… scales efficiently up to 1,024 GPUs and maintains high performance with sequence lengths up to 128K tokens… The code is available in Megatron-Core
- Section 1 Contributions (page 2): 1. MoE Parallel Folding: We introduce MoE Parallel Folding, the first approach that decouples parallelization strategies for attention and MoE layers… 2. Flexible and efficient token-level dispatcher… 3. Performance enhancements… achieve 49.3% MFU for Mixtral 8x22B and 39.0% MFU for Qwen2-57B-A14B on H100 GPUs.
- Section 3.2 MoE Parallel Folding (page 5): for the attention layers, we form a four-dimensional parallel group comprising TP × CP × DP × PP. For the MoE layers, we define another four-dimensional group consisting of TP × EP × DP × PP. … The only restriction is that the number of PP groups and members of each PP group for the Attention and MoE layer must be consistent.
- Section 3.3 Token Dispatcher (page 6): we adopt the sub-sequence dropping approach as the default strategy in this work. … Empirically, we observe that sub-sequence dropping does not adversely affect model convergence compared to full-sequence dropping.
- Table 1 (page 8): MCore w/ Folding 49.3% [Mixtral-8x22B] 41.6% [Llama3-8x70B] 39.0% [Qwen2-57B-A14B] 28.8% [Mixtral-8x22b-G8T8]
Discovery evidence
Section titled “Discovery evidence”- topic:
training-systems - sources:
arxiv - retrieved_at: 2026-07-20
- query: Megatron large scale transformer training
- arxiv:
2504.14960 - score_total: 50
- suggested_tier:
recent
Relevance
Section titled “Relevance”(no prose relevance explanation — numeric score only or HTTP source)
Snippets
Section titled “Snippets”(no snippet evidence in candidate pool)
延伸解读与背景补充
Section titled “延伸解读与背景补充”生成：2026-07-21 · 来源条数 2 · 模型
heuristic· 需人工核验数字
围绕「MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core」的核心问题与动机（待结合全文核验）。
方法要点补强
Section titled “方法要点补强”- 见原文方法章节；以下为基于摘要/摘录的要点提示。
- MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model …
与相近工作的关系
Section titled “与相近工作的关系”与相近工作的关系待核验；请对照 related work。
- 勿仅凭摘要推断未给出的数值指标。
- 这篇工作的输入/输出表示是什么？（MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core）
- 训练目标与评测协议各是什么？
- 主要失败模式或局限是什么？
外部解读索引
Section titled “外部解读索引”- [2504.14960] MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core — 2026-07-21 — 官方摘要/二次页面（自动抓取）
- [2504.14960] MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core — 2026-07-21 — 官方摘要/二次页面（自动抓取）
方法结构（重绘）
Section titled “方法结构（重绘）”flowchart LR A["输入 / 观测"] --> B["表示 / 编码"] B --> C["推理 / 解码"] C --> D["输出 / 动作或检测"] %% method sketch for: MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Sca方法结构示意（重绘；细节以原论文为准，待 PDF 核验）。
论文摘录图/表
Section titled “论文摘录图/表”
来源：原论文约 p.1（arch）；学习用途摘录。
来源：原论文约 p.18（table）；学习用途摘录。
英文自动分析（可折叠）
Section titled “英文自动分析（可折叠）”展开英文 Paper Card / AI deep analysis
Paper Card
Section titled “Paper Card”| Field | Content |
|---|---|
| Year | 2025 |
| Authors | Dennis Liu, Zijie Yan, Xin Yao, Tong Liu, Vijay Korthikanti, Evan Wu, Shiqing Fan, Gao Deng, Hongxiao Bai, Jianbin Chang, Ashwath Aithal, Michael Andersch |
| arXiv | 2504.14960 |
| DOI | — |
| Topics | training-systems, deployment-inference |
| Paper | https://arxiv.org/abs/2504.14960 |
原文摘录与素材（可折叠）
Section titled “原文摘录与素材（可折叠）”展开 Extract / Selections / Local assets
Selections
Section titled “Selections”training-systems: tier=recent rank=3 score=50 — auto refresh 2026-07-19 sources=arxivdeployment-inference: tier=watch rank=5 score=51 — cross-topic assign from registry title match=1 keywords; 2026-07-19
Extract excerpt
Section titled “Extract excerpt”MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core
Dennis Liu∗ Zijie Yan∗ Xin Yao Tong Liu Vijay Korthikanti Evan Wu Shiqing Fan Gao Deng Hongxiao Bai Jianbin Chang Ashwath Aithal Michael Andersch Mohammad Shoeybi
arXiv:2504.14960v3 [cs.LG] 2 Mar 2026 Jiajie Yao Chandler Zhou David Wu Xipeng Li June Yang † NVIDIA {denliu, zijiey, xiny, tongliu, vkorthikanti, evwu, shiqingf, gdeng, hongxiaob, jianbinc, aaithal, mandersch, mshoeybi, jiajiey, chandlerz, davidwu, xipengl, juney}@nvidia.com
Abstract Mixture of Experts (MoE) models enhance neural network scalability by dynami- cally 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 train- ing 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 l