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

可扩展的专家混合模型训练与 Megatron Core
作者/机构: NVIDIA

---

A1 主要贡献

1. 引言

与计算量随模型规模线性增长的密集型 Transformer 模型【1, J. Kaplan 等人, "Scaling laws for neural language models", 2020; 2, J. Hoffmann 等人, "Training compute-optimal large language models", 2022】不同，专家混合（MoE）模型遵循一种不同的扩展模式：通过将每个令牌（token）路由到一部分专家网络而非激活所有参数，其单位令牌的计算量随模型规模呈次线性增长【3, N. Shazeer 等人, "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer", arXiv preprint arXiv:1701.06538, 2017; 4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022】。近期的 MoE 模型在同等质量下，相较于密集型模型，训练成本降低了几个数量级【5, A. Q. Jiang 等人, "Mixtral of experts", 2024; 6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025; 7, W. Cai 等人, "A survey on mixture of experts in large language models", IEEE Transactions on Knowledge and Data Engineering, 2025】。

然而，大规模训练 MoE 模型引入了密集模型框架未曾涉及的系统性挑战。本文介绍了 Megatron-Core MoE，这是 Megatron-Core【8, M. Shoeybi 等人, "Megatron-lm: Training multi-billion parameter language models using model parallelism", 2020】中的 MoE 训练技术栈，涵盖了训练万亿参数级别 MoE 模型所需的高吞吐量架构、并行策略和系统优化。

1.1. 专家混合模型

MoE 模型通过一组称为“专家”的专业化子网络来增强标准神经网络，并配备一个轻量级的路由器（或门控网络），用于动态选择处理每个输入的专家【3, N. Shazeer 等人, "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer", arXiv preprint arXiv:1701.06538, 2017; 9, R. A. Jacobs 等人, "Adaptive mixtures of local experts", Neural Computation, 1991; 10, D. Eigen 等人, "Learning factored representations in a deep mixture of experts", arXiv preprint arXiv:1312.4314, 2013】。在基于 Transformer 的语言模型中【11, A. Vaswani 等人, "Attention is all you need", Advances in Neural Information Processing Systems, 2017】，MoE 层通常取代密集的（FFN）模块。

形式上，给定输入令牌表示 x，路由器计算 E 个专家上的概率分布：

$$\mathbf{p}(\mathbf{x}) = \operatorname{Softmax}(\mathbf{W}_r \mathbf{x})$$

MoE 层的输出是所选专家输出的加权组合：

$$\mathrm{MoE}(\mathbf{x}) = \sum_{i \in \mathrm{TopK}(\mathbf{p}(\mathbf{x}))} p_i(\mathbf{x}) \cdot E_i(\mathbf{x})$$

该架构有三大优势：可扩展容量、计算效率和专业化。尽管 MoE 概念可追溯至 20 世纪 90 年代初【9, R. A. Jacobs 等人, "Adaptive mixtures of local experts", Neural Computation, 1991; 10, D. Eigen 等人, "Learning factored representations in a deep mixture of experts", arXiv preprint arXiv:1312.4314, 2013】，但其与现代 Transformer 的结合推动了新一轮的研究热潮。GShard【12, D. Lepikhin 等人, "Gshard: Scaling giant models with conditional computation and automatic sharding", arXiv preprint arXiv:2006.16668, 2020】开创了大规模分布式 MoE 训练，引入了专家并行（Expert Parallelism）和负载均衡辅助损失。Switch Transformer【4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022】证明了 MoE 可以扩展到万亿参数并保持训练稳定。GLaM【13, N. Du 等人, "Glam: Efficient scaling of language models with mixture-of-experts", 2022】表明 MoE 能以更低的训练成本达到与密集模型相当的质量。Tutel【14, C. Hwang 等人, "Tutel: Adaptive mixture-of-experts at scale", Proceedings of Machine Learning and Systems, 2023】和 DeepSpeed-MoE【15, S. Rajbhandari 等人, "Deepspeedmoe: Advancing mixture-of-experts inference and training to power next-generation ai scale", International Conference on Machine Learning, 2022】等框架进一步推动了 MoE 训练系统的发展。

MoE 的应用在学术界和工业界加速普及。Mixtral-8x7B【5, A. Q. Jiang 等人, "Mixtral of experts", 2024】展示了开源 MoE 模型能媲美专有密集模型，同时降低推理成本。DeepSeek-V2 和 DeepSeek-V3【16, D. Dai 等人, "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models", arXiv preprint arXiv:2401.06066, 2024; 17, DeepSeek-AI 等人, "Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model", 2024; 6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】通过使用数百个小型专家来最大化容量与计算的比率，进一步扩展了这一理念。NVIDIA 的 Nemotron-3 系列【18, NVIDIA, "Nvidia nemotron 3: Efficient and open intelligence", 2025】采用了混合 Mamba-Transformer MoE 架构和 LatentMoE【19, V. Elango 等人, "Latentmoe: Toward optimal accuracy per flop and parameter in mixture of experts", 2026】，并使用 Megatron-Core 进行大规模训练。扩展定律研究证实，细粒度 MoE 实现了有利的计算最优权衡【20, J. Krajewski 等人, "Scaling laws for fine-grained mixture of experts", arXiv preprint arXiv:2402.07871, 2024; 21, J. Ludziejewski 等人, "Joint moe scaling laws: Mixture of experts can be memory efficient", 2025】，加速了这一趋势，同时也加剧了其带来的系统挑战。

1.2. 训练大规模 MoE 模型的挑战

MoE 训练的系统挑战源于其稀疏性，这表现为参数-计算不匹配（总参数远超活跃计算）和密集-稀疏不匹配（注意力层和 MoE 层需要冲突的并行配置）。

• 参数-计算不匹配: 在密集模型中，参数和单位令牌计算量同步扩展。而在 MoE 模型中，只有 K 个专家被激活，单位令牌计算量与 $N_{active}$（与 K 成正比）相关，而总参数量与 $N_{total}$（与 E 成正比）相关，且 $K \ll E$。例如，DeepSeek-V3 拥有 6850 亿总参数，但每个令牌仅激活 370 亿参数，存在 18 倍的差距。这导致了三个紧密耦合的挑战，即“三座大山”：
  • 内存墙 (Memory Wall): 所有 E 个专家的参数、梯度和优化器状态都必须驻留在内存中，即使每个令牌只激活 K 个。这导致内存压力远超同等计算量的密集模型【4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022; 7, W. Cai 等人, "A survey on mixture of experts in large language models", IEEE Transactions on Knowledge and Data Engineering, 2025】。
  • 通信墙 (Communication Wall): 专家并行（EP）需要 all-to-all 通信来分发和收集令牌。随着 EP 规模扩大，通信越来越依赖带宽较低的节点间互联，而稀疏的激活模式提供的计算量有限，难以掩盖通信开销。在 DeepSeek-V3 这类架构中，未经优化的 all-to-all 通信可占总训练时间的 60%。
  • 计算效率墙 (Compute Efficiency Wall): MoE 引入了密集模型中不存在的计算效率低下问题：
    • 小 GEMM: 细粒度专家产生大量小型矩阵乘法，无法充分利用 GPU 计算单元。
    • 路由和排列开销: 令牌路由和排列增加了额外的层执行时间。
    • 负载不均衡: 动态路由导致专家接收的令牌数量不均，造成计算资源浪费。
    • 主机开销: MoE 启动更多的小核函数（kernel），增加了主机端开销，导致 GPU 在核函数之间空闲。
      这三座大山相互关联：优化一个方面通常会将压力转移到另一个方面。

1.3. Megatron-Core MoE

Megatron-Core MoE 是一个基于 PyTorch 的大规模 Transformer 训练库【8, M. Shoeybi 等人, "Megatron-lm: Training multi-billion parameter language models using model parallelism", 2020; 25, A. Paszke 等人, "Pytorch: An imperative style, high-performance deep learning library", 2019】，它通过以下方式同时解决这三座大山：

• 多维并行: 将专家并行（EP）与张量、流水线、序列和数据并行集成。MoE Parallel Folding【26, D. Liu 等人, "Moe parallel folding: Heterogeneous parallelism mappings for efficient large-scale moe model training with megatron core", 2025】解耦了注意力层和 MoE 层的配置，打破了传统的 EP ≤ DP 约束。
• 内存优化: 通过细粒度激活重计算【27, V. Korthikanti 等人, "Reducing activation recomputation in large transformer models", 2022】、内存高效排列、精度感知优化器和激活卸载【28, A. A. A 等人, "Learning (with) distributed optimization", 2023】来减少内存占用。支持全面的低精度训练（FP8 和 FP4）。
• 通信优化: 高性能令牌分发器（DeepEP, HybridEP）和通信-计算重叠技术。
• 计算优化: Grouped GEMM、核函数融合、CUDA Graphs 和无同步执行。
• 生产级特性: 负载均衡、令牌丢弃、分布式优化器和 FSDP 支持、分布式检查点【29, NVIDIA, "Nvidia megatron core moe user guide", 2024】等。

---

A3 背景知识/关键 Observation/设计原则

2. Megatron-Core MoE 架构

本节分为两部分介绍 Megatron-Core 的 MoE 实现架构。首先，描述 MoE 层自身的内部设计：其模块化组件（路由器、令牌分发器、专家）以及将输入令牌转换为输出表示的四阶段前向传播。然后，探讨其外部设计：如何组织并行进程组以支持分布式执行，以及优化器如何处理专家参数。

2.1. MoE 层架构和前向传播

MoE 层的模块化组件和四阶段前向传播。MoE 层通过三个模块化组件实现：路由器（负责令牌到专家的分配）、令牌分发器（负责跨 GPU 通信）和专家（负责计算）。这三者通过如图 1 所示的四阶段前向传播连接。这种关注点分离的设计使得各部分可以独立优化：路由器可以被融合到 CUDA Graphs 中而不影响分发器逻辑；分发器可以在 all-to-all 和 DeepEP 之间切换而不改变专家计算；专家实现可以使用不同的 GEMM 后端。

2.1.1. 前向传播：路由、分发、计算、合并

MoE 层通过四个顺序阶段处理输入令牌：

1. 路由 (Route)：TopKRouter 决定每个令牌由哪些专家处理。它通过一个学习到的线性投影将每个令牌的隐藏状态映射到 E 个 logits，然后一个评分函数（softmax 或 sigmoid）将 logits 转换为概率，最后 top-k 选择为每个令牌识别出得分最高的专家。路由器输出 probs（路由权重）和 routing_map（令牌-专家分配的布尔掩码）两个张量。
2. 分发 (Dispatch)：令牌分发器为跨 GPU 通信准备令牌。它首先对令牌进行排列（permute），使所有发往同一专家的令牌在内存中是连续的，这对于专家端的高效密集 GEMM 至关重要。然后，分发器使用三种后端之一将令牌移动到托管其指定专家的 GPU 上：AllGather（简单但内存密集）、all-to-all（标准的基于 NCCL 的）或 Flex（支持 DeepEP 和 HybridEP 等优化后端）。
3. 专家计算 (Expert Computation)：每个 GPU 在接收到的令牌上执行其本地专家。所有本地专家通过 TEGroupedMLP 在一个 Grouped GEMM 调用中运行，即使单个专家的工作负载很小，也能最大化 GPU 利用率。
4. 合并 (Combine)：通过反向通信将处理后的令牌返回到其原始 GPU，然后进行反向排列（unpermutation）以恢复原始序列顺序。如果配置了共享专家，其输出（可以选择与路由专家并行计算）在此阶段被添加。

2.1.2. 路由器：令牌到专家的分配

路由器通过两个操作将全局令牌批次转换为专家特定的工作负载【3, N. Shazeer 等人, "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer", arXiv preprint arXiv:1701.06538, 2017】：
1. 门控 (Gating)：一个线性投影 $W_r \in R^{h \times E}$ 将每个令牌的隐藏状态 $x \in R^h$ 映射到 logits $l = W_r^T x \in R^E$。
2. Top-k 选择 (Top-k Selection)：一个评分函数将 logits 转换为概率，例如 softmax 或 sigmoid（DeepSeek-V3【6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】使用）。为每个令牌选择概率最高的 top-k 个专家。

为解决专家利用率不均的问题，Megatron-Core 支持多种负载均衡策略（见图 2），这些策略将在第 7 节讨论。

2.1.3. 令牌分发器：通信抽象

分发器通过一个六阶段的流水线管理 GPU 之间的令牌移动：dispatch_preprocess → token_dispatch → dispatch_postprocess（前向传播），以及 combine_preprocess → token_combine → combine_postprocess（反向传播）。

通信后端。提供了三种分发器类型：
* AllGather (allgather)：每个 GPU 收集所有令牌并筛选出本地专家的令牌。简单但内存密集，适用于小规模 EP。
* All-To-All (all-to-all)：标准的基于 NCCL 的点对点通信【12, D. Lepikhin 等人, "Gshard: Scaling giant models with conditional computation and automatic sharding", arXiv preprint arXiv:2006.16668, 2020】。每个 GPU 只发送每个目标 GPU 需要的令牌。扩展性好但会产生同步开销。
* Flex (flex)：统一设计，支持 DeepEP（仅集成高吞吐量核函数，通过重叠隐藏 NVLink 通信延迟）【23, C. Zhao 等人, "Deepep: an efficient expert-parallel communication library.", 2025】和 HybridEP（高带宽 MoE 通信核函数，在 NVLink 丰富的拓扑（如 NVL72）上性能更佳）。

2.1.4. 专家：计算模块

每个专家是一个两层 MLP，可选择性地带有门控（用于 SwiGLU/GeGLU 激活函数【30, N. Shazeer, "Glu variants improve transformer", arXiv preprint arXiv:2002.05202, 2020】）。Grouped GEMM 能够高效地批处理多个专家的计算【24, T. Gale 等人, "MegaBlocks: Efficient Sparse Training with Mixtureof-Experts", Proceedings of Machine Learning and Systems, 2023; 14, C. Hwang 等人, "Tutel: Adaptive mixture-of-experts at scale", Proceedings of Machine Learning and Systems, 2023】。提供了两种实现：
* TEGroupedMLP：Transformer Engine【31, N. Corporation, "Transformer engine.", 2025】的优化实现，支持 FP8 和 FP4 量化。
* SequentialMLP：在循环中逐个执行专家。对调试有用，但速度明显较慢。

共享专家。一些架构（如 DeepSeek-V2/V3【17, DeepSeek-AI 等人, "Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model", 2024; 6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】、Qwen【32, Qwen 等人, "Qwen2.5 technical report", 2025】）包含一个共享专家，它处理所有令牌，无论路由如何。共享专家的计算可以与分发-计算-合并流水线并行运行，从而隐藏其延迟。详情见第 7 节。

2.2. 系统集成：并行化与优化器

并行进程组管理和优化器处理。本节探讨 MoE 层如何集成到分布式训练系统中，包括并行进程组的组织和优化器对专家参数的处理。

2.2.1. 并行组管理

MoE 层需要与密集层不同的进程组，因为不同组件有不同的通信模式。Megatron-Core 通过 ProcessGroupCollection 来组织这些组：

ProcessGroupCollection
|-- Attention Layer Groups: tp, cp, dp, pp
|-- Expert Layer Groups:    ep, expt_tp, expt_dp, pp

每个 MoE 组件根据其通信需求使用特定的组，如下表 1 所示。这种分离使得 Parallel Folding（第 3.3 节）成为可能：注意力层和 MoE 层可以使用不同的 TP/DP 配置。例如，注意力层可能使用 TP=4，而 MoE 层使用 ETP=1 和更高的 EP，从而独立地优化每种层类型。

2.2.2. 优化器和梯度处理

专家参数在分布式优化中需要特殊处理。Megatron-Core 使用一个 ChainedOptimizer，它包装了用于密集参数和专家参数的独立优化器。三个关键设计决策支持正确的 MoE 优化：
1. 参数识别：专家参数被标记为 allreduce=False，以区别于使用标准数据并行梯度归约的密集参数。
2. 独立的归约组：密集层在 dp_cp_group（完整数据并行）上归约梯度，而专家在 expt_dp_group（专家数据并行）上归约。这确保了梯度在正确数量的副本上进行平均。
3. 梯度缩放：专家梯度按 edp_size / dp_size 进行缩放，以解释专家（处理依赖于路由的令牌子集）与密集层所看到的有效批次大小不同。

这种设计使得 ZeRO 风格的优化器状态分片【33, S. Rajbhandari 等人, "Zero: Memory optimizations toward training trillion parameter models", 2020】能够与 MoE 无缝协作：专家参数的优化器状态在 EP 组内分片，而密集参数的状态则遵循标准的 DP 分片。

3. 扩展 MoE：并行折叠与多维并行

MoE 的稀疏性如何影响并行设计。第 1.2 节指出，MoE 的稀疏性解耦了模型大小与单位令牌计算量，造成了参数-计算不匹配，这使得专家并行（EP）成为必要，但也引入了 all-to-all 通信。本节探讨这种不匹配如何具体影响并行设计。我们首先回顾密集模型的并行基线及其权衡，然后展示 MoE 的稀疏性如何打破这些策略的基础假设。最后，我们介绍 MoE Parallel Folding，这是 Megatron-Core 针对由此产生的密集-稀疏不匹配（即注意力层和 MoE 层具有冲突的最优并行配置）的解决方案。Parallel Folding 解耦了它们的映射，使每一层都能使用其最优拓扑。

3.1. 为何需要并行，以及 MoE 的不同之处

大型模型训练需要并行。在探讨 MoE 特有的挑战之前，我们先阐明大型模型训练为何必须使用并行化，以及它所引入的权衡。

3.1.1. 大型模型训练为何需要并行

内存是根本约束。单个 GPU 的内存有限，但训练大模型需要同时存储模型参数、优化器状态、梯度和激活【34, J. Rasley 等人, "Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters", 2020】。以 BF16 精度和 Adam 优化器训练的 Llama-405B【35, A. Grattafiori 等人, "The llama 3 herd of models", 2024; 36, H. Touvron 等人, "Llama: Open and efficient foundation language models", 2023; 37, H. Touvron 等人, "Llama 2: Open foundation and fine-tuned chat models", 2023】为例，其内存需求远超任何单个 GPU 的容量。因此，多 GPU 并非可选项，而是必需品。

计算吞吐量是效率原因。除了内存，单个 GPU 的计算能力也有限。大型模型训练需要巨大的 FLOPs；聚合多个 GPU 可以提高吞吐量，将训练时间从数年缩短到数周。

3.1.2. 并行化的权衡

并行化并非没有代价。每种并行策略都会引入开销：通信开销、同步、流水线气泡和计算强度降低。结果是，模型 FLOP 利用率（MFU）受并行策略的严重影响。对于密集模型，并行化的“成本”和“收益”成比例扩展。因为计算量随模型大小增长，通信在每一步中所占的份额变小，使得 MFU 在模型扩展时保持相对稳定。

MoE 的稀疏性打破了这种平衡。由于每个令牌只激活 E 个专家中的 K 个，总参数随 E 扩展，而单位令牌计算量仅随 K 扩展。虽然需要更多 GPU 来满足内存需求，但单位令牌计算量并未相应增长，导致通信开销暴露出来。

3.2. MoE 并行化的挑战

MoE 的并行悖论。传统并行策略是为密集 Transformer 设计的。MoE 模型具有根本不同的计算模式，需要一个新的并行维度——专家并行（EP），并在将 EP 与现有策略结合时产生独特的挑战。图 3 通过绘制总参数与每令牌前向 FLOPs 的关系，说明了这一悖论。密集模型（圆形）遵循 ≈2N 的参考线，随着规模扩大形成一个良性循环。而 MoE 模型（三角形）则显著低于 2N 参考线，用更少的 FLOPs 实现了同等能力。

根本的不对称性。以 DeepSeek-V3 为例，其参数量是其活跃计算量所暗示的 18 倍。这产生了复合效应：
1. 内存快速增长，迫使模型分布在许多 GPU 上。
2. 通信量增加，通过 EP 分发专家引入的 all-to-all 流量随 K（每个令牌被分派到 K 个专家）扩展。
3. 但计算量保持低位，单位令牌 FLOPs 仅与 $N_{active}$（∝ K）成正比，不足以与日益增长的通信重叠。

结果是，除非并行策略专门为这种不对称性设计，否则 MoE 训练将是通信密集型的。

3.2.2. 传统并行策略

密集 Transformer 训练通常结合四种并行策略：张量并行（TP）【8, M. Shoeybi 等人, "Megatron-lm: Training multi-billion parameter language models using model parallelism", 2020】、流水线并行（PP）【38, Y. Huang 等人, "Gpipe: Efficient training of giant neural networks using pipeline parallelism", Advances in Neural Information Processing Systems, 2019; 39, D. Narayanan 等人, "Pipedream: Generalized pipeline parallelism for dnn training", Proceedings of the 27th ACM Symposium on Operating Systems Principles, 2019; 8, M. Shoeybi 等人, "Megatron-lm: Training multi-billion parameter language models using model parallelism", 2020; 40, P. Qi 等人, "Zero bubble pipeline parallelism", 2023; 41, P. Qi 等人, "Pipeline parallelism with controllable memory", 2024; 42, S. Li 等人, "Chimera: Efficiently training large-scale neural networks with bidirectional pipelines", 2021; 43, L. Guan 等人, "Pipeoptim: Ensuring effective 1f1b schedule with optimizer-dependent weight prediction", 2025】、数据并行（DP）【44, C. J. Shallue 等人, "Measuring the effects of data parallelism on neural network training", 2019; 45, H. Bai, "Modern distributed data-parallel large-scale pre-training strategies for nlp models", 2022】和上下文并行（CP）【46, V. A. Korthikanti 等人, "Sequence parallelism: Long sequence training from system perspective", Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2023; 47, H. Liu 等人, "Ring attention with blockwise transformers for near-infinite context", 2023; 48, S. A. Jacobs 等人, "Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models", 2023】。

传统并行策略对 MoE 失效的原因：
* 在 MoE 专家上应用 TP：专家的隐藏维度较小，应用 TP 会产生更小的分片，降低计算效率。
* 使用 PP 处理众多专家：MoE 模型参数量巨大，若仅用 PP 需要很多流水线阶段，导致过多的流水线气泡。
* 单独使用 DP：DP 在每个 GPU 上复制整个模型，对于万亿参数模型这是不可能的。

3.2.3. 专家并行：第五个维度

专家并行（EP）跨 GPU 分布专家【12, D. Lepikhin 等人, "Gshard: Scaling giant models with conditional computation and automatic sharding", arXiv preprint arXiv:2006.16668, 2020】。EP 的独特权衡在于，它引入 all-to-all 通信，但能保持每个专家的计算规模，并减少每 GPU 的内存压力。EP 有两个关键好处（图 4）：一是通过将不同 GPU 的令牌分组到同一个专家，增加了计算强度，提高了 GEMM 效率；二是 all-to-all 通信量随专家数量增加保持不变，只改变参与的 GPU 数量。

然而，仅有 EP 是不够的。EP 仅适用于 MoE 层；注意力层仍然需要 TP、CP 等策略。这种异构性造成了“密集-稀疏不匹配”。

3.2.4. 结合 EP 与传统并行的挑战

一个 Transformer 块内包含两种根本不同的计算模式（见表 2），它们的最佳并行配置直接冲突：高 TP 有利于注意力层，但会分割专家分片；高 CP 有助于长上下文注意力，但对 MoE 层无益；高 EP 对 MoE 至关重要，但与注意力无关。

密集-稀疏不匹配。先前的 MoE 框架将 EP 视为 DP 的一个子维度【12, D. Lepikhin 等人, "Gshard: Scaling giant models with conditional computation and automatic sharding", arXiv preprint arXiv:2006.16668, 2020; 4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022】，即：

$$\text{World Size} = \text{TP} \times \text{CP} \times \text{PP} \times \text{DP}, \quad \text{where EP} \subseteq \text{DP}$$

这一约束带来了三个关键挑战：
1. GPU 需求的乘法效应：传统框架至少需要 TP × CP × PP × DP 个 GPU。由于 EP ⊆ DP，请求 EP=8 会强制 DP ≥ 8，这增加了 MoE 训练的入门门槛。
2. 被迫的次优并行：从业者必须在两种次优配置中选择：使用高 TP 来高效分割大的注意力矩阵，但这会分割小的专家；或者使用低 TP 来保持专家计算效率，但这会使注意力层并行化不足。
3. 跨节点通信：EP 约束在 DP 内，高 EP 常常迫使 all-to-all 通信跨越节点边界，那里的带宽比 NVLink 低 5-10 倍。

3.3. Megatron-Core 的解决方案：并行折叠与多维框架

并行折叠 (Parallel Folding) 是 Megatron-Core 对密集-稀疏不匹配的回应。它解耦了注意力层和 MoE 层的并行映射，允许每种层类型使用其最优拓扑【26, D. Liu 等人, "Moe parallel folding: Heterogeneous parallelism mappings for efficient large-scale moe model training with megatron core", 2025】。

3.3.1. 解耦的并行映射

核心思想是：不强迫注意力层和 MoE 层共享并行配置。并行折叠为注意力层和 MoE 层引入了独立的并行组：
* 注意力层在 TP × CP × DP × PP 上形成组，为序列级密集计算优化。
* MoE 层在 ETP × EP × EDP × PP 上形成组，其中 ETP（专家张量并行）和 EDP（专家数据并行）是 MoE 特有的维度。
唯一的约束是流水线并行（PP）必须在两种布局中保持一致。

3.3.2. 完整的多维并行技术栈

Megatron-Core 协调了五个并行维度，如图 5 所示。关键配置原则是：注意力层优化大矩阵（高 TP）和长序列（高 CP）；MoE 层优化众多小专家（高 EP，通常 ETP=1）；PP 在两者间保持一致。

此外，Megatron-Core 的并行折叠框架还集成了带有 EP 支持的分布式优化器和 FSDP，以进一步减少内存占用：
* 分布式优化器 + EP：每个 rank 只驻留本地专家的权重和梯度；优化器状态在同一专家的副本间（通过 EDP）进行分片。
* FSDP + EP：Megatron-Core 的定制 FSDP (Megatron-FSDP) 通过双 DeviceMesh 架构，在数据/专家组上完全分片参数、梯度和优化器状态，减少内存占用，同时将 AllGather 和 ReduceScatter 与计算重叠。

3.3.3. 并行折叠的优势

如图 5 和图 6 所示，MoE 并行折叠消除了 EP ≤ DP 的限制，允许 EP “折叠”到注意力并行配置的任意子组上。这带来了四个关键优势：
1. 打破 EP ≤ DP 约束：EP 现在可以通过在 TP×CP 组上折叠来超过 DP。
2. 降低最小 GPU 需求：传统配置中 CP=8, EP=8 至少需要 64 个 GPU。通过折叠，CP 和 EP 共享同一 GPU 组，仅需 8 个 GPU。
3. 实现独立优化：注意力层可以使用高 TP，而 MoE 层可以使用 ETP=1。
4. 将高带宽通信保持在 NVLink 域内：CP（用于注意力）和 EP（用于 MoE）的 all-to-all 通信都可以保留在 NVLink 连接的 GPU 组内。

总之，MoE 并行折叠通过解耦并行映射解决了密集-稀疏不匹配问题，与分布式优化器和 FSDP 一起，将所有并行维度整合到一个统一的框架中，可将 MoE 训练扩展到数千个 GPU。

---

A2 方法细节

4. 扩展 MoE：攻克内存、通信和计算效率三座大山

引言。第 3 节确立了如何在数千个 GPU 上分布 MoE 训练。本节探讨如何使其高效。挑战源于 MoE 的稀疏性驱动的参数-计算不匹配，这导致了三个基本障碍：内存墙、通信墙和计算效率墙。这三座大山紧密耦合：优化一个通常会恶化另一个。例如，启用激活重计算以解决内存墙问题，可能会暴露出通信墙；实现通信-计算重叠以解决通信墙，可能会因细粒度专家计算太快而效果不佳。Megatron-Core 通过集成的系统设计来解决这些挑战。

4.1. 攻克内存墙

内存是硬约束。内存是 MoE 训练的第一个硬约束。如果参数、优化器状态和激活的总占用超过 GPU 容量，训练就无法进行。

4.1.1. 内存剖析：内存消耗的来源

MoE 模型通常比等效的密集模型消耗更多的内存。以 BF16 精度、PP4 × VPP4 × EP64 配置在 256 个 GPU 上训练 DeepSeek-V3 为例，每个 GPU 的内存明细如表 3 所示，总需求为 199.5 GB，远超 H100 GPU 的 80 GB 容量。

激活是主要目标。激活是最大的内存消耗者，超过了权重和优化器状态的总和，因此是优化的主要目标。四种互补的策略可以解决内存约束：降低存储精度、用计算换存储、卸载到主机内存、跨设备分布。

4.1.2. 内存高效的排列：零开销激活减少

零开销的内存优化。内存高效排列通过简单的代数重排来消除冗余的中间张量，从而实现零计算开销的内存优化。它改变了应用路由器概率的时间，而不是是否应用它们。

标准公式与内存高效公式。考虑一个令牌 x 被路由到其 top-k 专家。标准公式中，路由权重在专家计算后应用：

$$ \mathbf{y}=\sum_{i \in \mathcal{T}(\mathbf{x})} p_{i} \cdot \mathbf{W}_{2}^{(i)} \phi(\mathbf{W}_{1}^{(i)} \mathbf{x}). $$

其中 $E_i(x) = W_2^{(i)} \phi(W_1^{(i)} x)$。

内存高效排列将 $p_i$ 吸收到激活中，在第二个线性层之前应用：

$$\mathbf{y}=\sum_{i \in \mathcal{T}(\mathbf{x})} \mathbf{W}_{2}^{(i)}\left(p_{i} \cdot \phi\left(\mathbf{W}_{1}^{(i)} \mathbf{x}\right)\right).$$

当专家没有偏置项时，这两个公式在数学上是等价的。这种重排通过消除为路由器反向传播所需的保存张量来降低峰值内存。对于 DeepSeek-V3（表 3），内存高效排列为每个 GPU 节省了约 26.3 GB 的激活内存。

4.1.3. 降低精度的训练：使用 FP8/FP4 减少激活内存

用低精度存储激活。在前向传播期间，每个线性层的输入必须被保留，以便在反向传播中计算权重梯度。通过将这些输入张量以 FP8/FP4 而非 BF16 存储，每个张量的内存占用可减少 50%/75%。对于表 3 中的 DeepSeek-V3 配置，启用 FP8 训练可减少约 16 GB 的激活内存。此减少与其他激活优化（如重计算和卸载）是正交的，可以组合使用。

4.1.4. 重计算：用计算换取内存

细粒度重计算。激活重计算（或激活检查点）是一种通过在反向传播期间重新计算中间激活来丢弃它们的技术【50, T. Chen 等人, "Training deep nets with sublinear memory cost", 2016】。然而，简单的全层重计算会增加约 33% 的计算开销，对于 MoE 层甚至更高。Megatron-Core 引入了细粒度重计算，只针对内存密集但计算廉价的操作，以最小的开销实现显著的内存节省【27, V. Korthikanti 等人, "Reducing activation recomputation in large transformer models", 2022】。

两种技术。该策略由两种技术组成（图 8）：
1. 粒度化重计算：用户可以精确指定在反向传播中要重计算哪些计算，例如只重计算专家 MLP 内的激活函数或 LayerNorm 模块。
2. 输出丢弃重计算：Megatron-Core MoE 在检查点模块的输出被后续层消耗后立即释放它们，在反向传播期间从重计算结果中恢复输出，从而避免了冗余存储。

表 4 总结了对表 3 中 DeepSeek-V3 配置的不同重计算目标所实现的内存减少。

4.1.5. 细粒度的激活卸载

用 PCIe 带宽换取 GPU 内存。当 GPU 内存即使在精度和重计算优化后仍然不足时，将激活卸载到 CPU 内存可提供额外容量。挑战在于将传输延迟隐藏在计算之后。细粒度 MoE 模型具有极端的参数膨胀，使得激活内存成为一个严重问题，而 EP 或 PP 策略并不能减少激活内存。

高层思想：重叠和预取。GPU 的复制引擎和计算引擎可以独立运行。当一个模块的计算时间超过其激活传输时间时，D2H 复制可以与后续计算并行运行，成本为零。图 9 说明了前向和反向传播的流重叠机制。

关键技术特性。
* 模块级粒度：用户通过 --offload-modules 指定要卸载的模块。
* 异步传输：专用的 D2H/H2D CUDA 流与计算并行运行传输。
* 重计算集成：与细粒度重计算（4.1.4 节）结合。
* CUDA Graphs 兼容性：使用外部事件而非流同步。

性能。细粒度卸载和重计算作为互补策略协同工作（图 10）。表 5 显示了多种配置下的结果。细粒度卸载可将内存减少 10-18%，吞吐量开销仅为 1.6-2%。

4.1.6. 权重和优化器优化：低精度存储和卸载

精度感知优化器。Adam 优化器【56, D. P. Kingma 等人, "Adam: A method for stochastic optimization", 2014】为每个参数维护两个状态张量。传统实现将它们存储在 FP32 中。精度感知优化器【57, T. Dettmers 等人, "8-bit optimizers via block-wise quantization", International Conference on Learning Representations, 2022】解耦了存储精度和计算精度，允许将动量存储在 BF16 甚至 FP8 中，从而将每个参数的存储从 8 字节减少到 4 或 2 字节。更新计算仍在 FP32 中进行以保持数值稳定性。这可将优化器状态内存减少约 50%。

状态卸载。优化器状态卸载在 optimizer.step() 之后将优化器状态和主权重传输到 CPU，并在下一次步骤之前重新加载它们。这种方法在保留 GPU 计算的同时，回收了前向和反向传播期间的内存。在 GB200 上，状态卸载可节省 15-20 GB 的 GPU 内存，每次迭代的开销仅为 0.1-0.2 秒。

4.1.7. 用于 MoE 的 FSDP

FSDP 与 EP 的结合。完全分片数据并行（FSDP）【53, Y. Zhao 等人, "Pytorch fsdp: Experiences on scaling fully sharded data parallel", 2023; 33, S. Rajbhandari 等人, "Zero: Memory optimizations toward training trillion parameter models", 2020】将模型参数、梯度和优化器状态分片到数据并行 rank 上。对于 MoE 模型，将 FSDP 与 EP 结合具有具体的好处：
* 减少内存和通信：EP 为每个 GPU 分配一部分专家，FSDP 在专家数据并行（EDP）组内对这些本地专家进行分片，而不是在整个 DP 组内。
* 无 PP 的灵活性：FSDP+EP 避免了流水线并行的几个工程难题。
* 广泛的模型兼容性：Megatron-FSDP 支持 Megatron-Core 模块和 PyTorch 原生模块。

双 DeviceMesh 设计。核心设计挑战是密集层和专家层需要不同的分片范围。Megatron-FSDP 通过双 DeviceMesh 架构解决了这个问题。一个主 DeviceMesh 管理密集模块的 DPShard、DP-Outer、TP 和 CP。一个辅助的 Expert DeviceMesh 管理 EP 模块，FSDP 的范围限定在 EDP 维度。

零拷贝通信。Megatron-FSDP 通过两种优化消除了通信开销：
1. 非均匀分片：Megatron-FSDP 将模块内的所有参数展平并连接，然后跨设备进行非均匀分片（图 11b），使分片边界与通信缓冲区布局对齐。
2. 带 NCCL 用户缓冲区注册的持久双缓冲区：Megatron-FSDP 在训练开始时预分配两个持久缓冲区并循环使用它们（图 12），消除了分配开销。然后通过用户缓冲区注册（UBR）将这些缓冲区注册到 NCCL，实现真正的零拷贝通信。

计算-通信重叠。Megatron-FSDP 使用专用的 CUDA 流将 AllGather 和 ReduceScatter 与计算流水线化，从而隐藏通信延迟。

4.1.8. 总结

表 6 总结了本节描述的内存优化技术及其主要目标。这些优化互为补充，共同将内存从一个阻塞性障碍转变为一个可管理的约束。内存优化不仅是训练的先决条件，还为其他优化（如增加批处理大小以隐藏通信延迟）释放了空间。

4.2. 攻克通信墙

通信开销直接降低 GPU 利用率。在优化之前，EP all-to-all 通信通常消耗 20-60% 的训练时间。对于像 DeepSeek-V3 和 Kimi-K2 这样的细粒度 MoE 架构，高频率的 dispatch/combine 操作、跨节点瓶颈和低算术强度加剧了这一挑战。

4.2.1. 通信剖析：专家并行模式

图 13 说明了专家并行的通信模式。在标准的 EP 实现中，每个 MoE 层需要两次集合通信操作：dispatch 发送令牌，combine 返回处理后的令牌。优化必须针对通信的执行和调度方式。两种策略协同工作：最大化带宽利用率和将延迟隐藏在计算之后。

4.2.2. DeepEP 和 HybridEP：最大化 EP 带宽

基于令牌的分发后端。为了解决传统 all-to-all 的局限性，Megatron-Core 提供了两个基于令牌的分发后端，HybridEP 和 DeepEP【23, C. Zhao 等人, "Deepep: an efficient expert-parallel communication library.", 2025】。基于令牌的分发消除了排列步骤，避免了发送冗余令牌，从而减少了总通信量并提高了有效带宽。

HybridEP 设计。HybridEP 由 NVIDIA 开发。对于分发（图 14），HybridEP 根据路由信息将数据从全局内存读入共享内存，然后通过 FIFO 队列将令牌写入目的地。在跨节点情况下，它首先在节点间交换数据，然后在每个节点内转发，减少了跨节点流量。对于合并（图 15），HybridEP 将归约操作融合到通信核函数中。

性能评估。如表 7 所示，在所有测试设置中，HybridEP 的性能始终优于 all-to-all，在跨节点场景中增益更大。

4.2.3. EP 通信重叠：隐藏 EP 通信延迟

1F1B all-to-all 重叠方案。优化后的分发器提高了带宽利用率，但 all-to-all 仍位于关键路径上。Megatron-Core 使用专用的 1F1B all-to-all 重叠方案【61, Z. Liu 等人, "1f1b-based ep a2a overlapping for moe models.", 2025】，将相邻微批次的前向和反向传播合并，并在 CUDA 流上交错计算和 all-to-all 核函数【39, D. Narayanan 等人, "Pipedream: Generalized pipeline parallelism for dnn training", Proceedings of the 27th ACM Symposium on Operating Systems Principles, 2019; 38, Y. Huang 等人, "Gpipe: Efficient training of giant neural networks using pipeline parallelism", Advances in Neural Information Processing Systems, 2019】。

两种模式。我们评估了两种模式：合并的 FWD-FWD / BWD-BWD（图 16）和合并的 FWD-BWD（图 17）。合并的 FWD-BWD 是首选策略，因为它没有额外的内存开销。

最大化重叠的优化。
1. 流分离：我们将工作负载分成计算流和通信流，通过在流之间交替任务，实现 all-to-all 与计算的并行运行。
2. W/D 拆分（权重梯度/数据梯度拆分）【6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】：我们将反向 MLP 拆分为权重梯度计算（W/mlp）和数据梯度计算（D/mlp），以打破依赖关系并增加重叠机会。

图 18 总结了这些优化。这些方法将 H100 上 DeepSeek-V3 训练中的 EP 通信开销从 30-40% 降低到 5% 以下。对于大型模型，合并的 FWD-BWD 可以与交错流水线并行（Interleaved PP）结合（图 19）。

4.2.4. 总结

通信墙优化协同工作：DeepEP/HybridEP 最大化带宽利用率；FWD-BWD 重叠隐藏延迟；Interleaved PP 扩展重叠机会。它们共同将 all-to-all 对训练时间的贡献降低到 10% 以下。

4.3. 攻克计算效率墙

计算效率是最终约束。即使有足够的内存和隐藏的通信，如果核函数效率低下或主机分派工作不够快，GPU 资源仍可能未被充分利用。

4.3.1. 计算剖析：效率低下的来源

MoE 训练中的计算效率低下源于两个方面：
* 核函数效率：细粒度 MoE 架构产生的 GEMM 规模很小，无法充分利用 GPU 资源。
* 主机开销：大量的小操作产生核函数启动开销，CPU 无法足够快地分派工作以保持 GPU 饱和。

五种策略协同工作：融合相关操作、用低精度加速、消除每次迭代的 CPU 逻辑、移除主设备同步、平衡专家负载。

4.3.2. Grouped GEMM 和核函数融合：提高核函数效率

Grouped GEMM。Grouped GEMM 通过批处理专家计算来提高硬件利用率【24, T. Gale 等人, "MegaBlocks: Efficient Sparse Training with Mixtureof-Experts", Proceedings of Machine Learning and Systems, 2023】。Megatron-Core 提供了多种 Grouped GEMM 实现，包括多流 cuBLASLt GEMM、CUTLASS Grouped GEMM，以及正在开发的 cuBLASLt Grouped GEMM（通过 CUBLASLT_BATCH_MODE_GROUPED）和带有融合的 cuteDSL Grouped GEMM。

核函数融合。核函数融合将路由和排列操作合并到更少的核函数中【63, T. Chen 等人, "TVM: An automated end-to-end optimizing compiler for deep learning", Proceedings of the 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2018】。Megatron-Core 在三个层面实现这些优化：专家计算（Grouped GEMM）、令牌路由（排列融合）和路由器逻辑（路由器和辅助损失融合）。

4.3.3. 排列融合

目前的 Grouped GEMM 要求分配给同一专家的令牌连续存储。排列融合流水线（图 20）包括三个阶段：
* 预处理：生成一个偏移图，指示每个令牌在输入和输出缓冲区中的偏移。
* 排列：根据偏移图将令牌从输入缓冲区移动到输出缓冲区。
* 反排列：是排列操作的逆过程，在合并阶段将令牌加在一起。

4.3.4. 路由器和辅助损失融合

路由器和辅助损失【4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022; 65, B. Zoph 等人, "St-moe: Designing stable and transferable sparse expert models", 2022; 14, C. Hwang 等人, "Tutel: Adaptive mixture-of-experts at scale", Proceedings of Machine Learning and Systems, 2023; 15, S. Rajbhandari 等人, "Deepspeedmoe: Advancing mixture-of-experts inference and training to power next-generation ai scale", International Conference on Machine Learning, 2022】计算会产生许多小核函数。我们将剩余操作分解为三个部分，每个部分都融合成一个单独的核函数（图 21）：
* 分数计算，包括 top-k 和 softmax/sigmoid。
* 辅助损失的分数计算。
* MoE 辅助损失的计算。

4.3.5. 低精度训练：低精度加速

用 FP8/FP4 最大化 Tensor Core 利用率。低精度训练通过减少数据移动和使用更快的 FP8/FP4 Tensor Core 操作来解决内存带宽瓶颈问题。对于大规模 MoE 训练，FP8 训练可实现约 10% 到 25% 的端到端性能提升。

4.3.6. CUDA Graphs：消除主机开销

用可重放的图消除启动开销。核函数融合减少了核函数启动的数量；CUDA Graphs 消除了这些启动的每次迭代成本。CUDA Graphs 将一系列 GPU 核函数捕获到一个可重放的图中【66, N. Corporation, "CUDA graphs.", 2024; 67, J. Ansel 等人, "PyTorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation", Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024】。后续迭代发出单个 Graph Launch 调用，从而减少 CPU 端延迟（图 22）。

完整 CUDA Graphs vs. 逐层 CUDA Graphs。然而，CUDA Graphs 要求所有捕获的操作都具有静态、预定的形状，这与 MoE 中令牌计数动态变化的情况存在根本矛盾。Megatron-Core 提供了两种 CUDA Graph 模式（图 23）：
* 完整 CUDA Graphs：捕获整个前向-反向传播过程。仅支持使用令牌丢弃和填充的 MoE 模型。
* 逐层 CUDA Graphs（部分 CUDA Graphs）：对于像 DeepSeek-V3 这样的无丢弃 MoE 模型，我们引入了部分 CUDA Graphs。它只捕获每层中的静态组件（如注意力、路由器、EP 预处理），而将动态的专家计算部分留在图外（图 24）。

CUDA Graphs 的内存优化。CUDA Graphs 会增加内存开销。Megatron-Core 和 Transformer Engine 实现了一些优化来最小化这种开销：
* 减少图的数量：在没有流水线并行的情况下，微批次可以共享同一个图（图 26）。
* 内存池共享：所有图可以共享一个内存池。
* 缓冲区重用：静态输入/输出缓冲区可以根据 PP 执行顺序在图之间重用。

4.3.7. 通过无同步核函数、ECHO 和分页存储实现无丢弃 MoE 的完整 CUDA Graphs 覆盖

无同步 MoE 的挑战。为了使 CUDA Graphs 能够覆盖动态部分，必须解决两个挑战：
1. 内核启动：传统上，许多 GPU 操作符假定（动态）形状信息在主机上可用。
2. 内存分配：无同步执行意味着缓冲区大小需要在主机上预先确定，最坏情况下的缓冲区分配会浪费内存。

设备启动的 GPU 核函数。为了消除主机-设备同步，核函数必须是设备启动的。TE 引入了设备启动的 Grouped GEMM，它直接从设备内存中读取形状信息。同样，HybridEP 也可以预分配输出缓冲区以消除同步。

ECHO：热门专家的弹性克隆。负载不平衡在 MoE 中是固有的。ECHO 通过将热门专家动态克隆到未充分利用的 rank 上的空闲槽位来解决这个问题（图 27）。ECHO 规划器使用 bin-packing 算法来最小化克隆数量，同时实现充分的负载均衡。这减少了 CUDA Graphs 的内存碎片，并改善了计算效率。

分页存储 (Paged Stashing)。虽然 ECHO 减少了负载不平衡，但 Paged Stashing 解决了互补问题：在 CUDA Graph 内实现细粒度的内存管理以减少内部内存碎片。Paged Stashing 解耦了计算所需的临时缓冲区和存储激活的存储缓冲区（图 28）。它维护一个在所有层之间共享的、大小为最坏情况的 tmp 缓冲区，以及一个以页形式组织的 stashing 缓冲区，只存储每层实际使用的令牌。这显著减少了内存消耗。为了最小化开销，Paged Stashing 使用专用的 CUDA 流将内存传输与计算重叠（图 29）。

整合。这三种技术共同实现了无丢弃 MoE 的完整 CUDA Graphs 覆盖，消除了 CPU 开销，同时保留了动态路由的灵活性。

5. 用于 MoE 的 FP8/FP4 低精度训练

低精度训练是跨领域的优化。低精度训练同时影响内存、通信和计算效率。对于 MoE 架构，它既放大了收益，也放大了风险。

策略：选择性精度。解决方案是在关键部分保持精度，在其他地方追求效率。
1. 保护路由决策：路由器保持在 FP32。
2. 保留关键组件的精度：嵌入、输出层、主梯度、主权重和优化器状态保持其原始精度。
3. 量化大块计算：专家 GEMM 使用低精度训练。

5.2. 低精度训练对三座大山的影响

低精度训练在所有三个性能墙上都提供了好处，使其成为一个统一的优化。表 8 总结了这些影响。

5.3. 低精度训练方案：Per-Tensor FP8, Blockwise FP8, MXFP8, 和 NVFP4

一个低精度训练方案包括数据格式、缩放粒度、量化部分和附加机制。Megatron-Core 提供了三种 FP8 方案和一种 NVFP4 方案（图 30）。

• Per-Tensor FP8 方案：每个张量使用一个缩放因子。简单但精度有限。
• Hopper 上的 Blockwise FP8：采用 E4M3 格式，对激活和梯度进行 1×128 tile 量化，权重进行 128×128 block 量化。在 Hopper 平台上是推荐的 FP8 方案。
• Blackwell 上的 MXFP8：在 Blackwell 平台上，我们采用 MXFP8，这是一种更细粒度的量化方案。激活和权重都以 1×32 的粒度进行量化。由于 Tensor Core 的原生支持，MXFP8 是 Blackwell 平台上的默认 FP8 方案。
• Blackwell 上的 NVFP4：NVFP4【75, NVIDIA 等人, "Pretraining large language models with nvfp4", 2025】是一种用于 LLM 训练的 4 位微缩放格式。它使用两级微缩放：一个 per-tensor FP32 缩放和一个 per-block 8 位缩放。稳定的 NVFP4 训练还依赖于随机哈达玛变换（RHT）、2D 缩放和随机舍入等技术。

5.3.5. FP8/FP4 主权重：消除冗余存储

传统的低精度训练维护三层参数：FP32 主权重、BF16 模型权重和 FP8/FP4 计算权重。原生 FP8/FP4 通过建立从 FP32 主权重到 FP8/FP4 计算权重的直接转换路径，完全绕过了 BF16 中间层，减少了内存占用并加速了参数 AllGather。图 33 和图 32 分别展示了 delayed/current scaling 和 blockwise scaling 的量化方案。

5.4. MoE 特有的低精度训练挑战和解决方案

5.4.1. 填充和去填充的融合：动态形状对齐

FP8 和 FP4 GEMM 要求张量维度对齐到特定倍数。在权重梯度 GEMM 中，令牌维度 M 是动态变化的，通常不满足对齐要求，因此需要沿令牌维度进行零填充。我们提出了两种解决方案：路由图填充和将填充融合到排列中。

5.4.2. 分组量化和分组 GEMM

分组量化。为避免因对不同专家的输入张量逐一应用量化核函数而产生大量小核函数，我们实现了分组量化核函数，将多个张量的量化融合成一个单一的核函数。
分组 GEMM。在 4.3.7 节中讨论了优化和 CUDA-Graphable 的分组 GEMM。

5.4.3. NVFP4 量化融合

NVFP4 量化因其数值技术而特别复杂，因此需要积极的核函数融合。我们的 NVFP4 量化核函数不仅是简单的“缩放+转换”，还吸收了训练方案逻辑，包括随机哈达玛变换、2D 缩放和随机舍入。对于 MoE 层，我们必须实现用于完整 NVFP4 流水线的分组量化核函数，以保持性能和图安全。

6. 长上下文 MoE 训练

当注意力占主导地位时的计算转变。随着序列长度的增加，MoE 训练的性能特征发生根本性转变：注意力计算成为主导，三座大山的相对重要性也随之改变。

6.1. 当注意力占主导地位：计算的转变

MoE 和 MLP 组件的计算成本随序列长度呈线性增长（$O(s)$），而缩放点积注意力（SDPA）的计算成本则呈二次方增长（$O(s^2)$）（图 34）。在 64K 令牌时，SDPA 消耗 69% 的 FLOPs。幸运的是，SDPA 在大多数 GPU 加速库中得到了高度优化（表 9），因此它不会成为性能瓶颈。优化的重点因此转移到内存和通信上。

6.2. 管理激活内存的增长

组合多种技术。为应对长上下文训练中加剧的内存墙，我们应用了一系列技术：
* 上下文并行和张量并行：结合 CP 和 TP，通过跨设备分布激活内存，使子序列长度保持近似恒定。
* 优化器 CPU 卸载：卸载优化器状态可回收大量 GPU 内存。
* 选择性重计算：基于内存与计算的权衡进行模块级选择。在长上下文设置中，重计算 SDPA 通常代价过高。我们建议禁用核心注意力重计算，优先重计算其他模块。

6.3. 上下文并行 vs. 张量并行

CP 和 TP 都能减少长上下文训练中的激活内存，但它们的通信模式和内存效应不同（图 35）。TP 通常在节点内更受青睐，因为通信速度快且内存效益强。跨节点时，P2P CP 通常更可取，因为 TP 通信开销增加，而 P2P 重叠仍然有效。一个实用的起点是在节点内使用 all-to-all CP 和 TP，在节点间使用 P2P CP。

6.4. 用于可变长度训练的打包序列

打包序列以消除填充浪费。实际训练场景，特别是强化学习和监督微调，通常涉及可变长度序列。传统的批处理需要将所有序列填充到批次内的最大长度，导致大量浪费。Megatron-LM 引入了打包序列，可以在单个批次内处理可变长度序列而无需填充浪费。

6.4.1. 打包序列支持

Megatron-LM 支持打包序列，通过 THD（总令牌数×头数×维度）张量格式实现，与传统的 SBHD 格式相对。核心机制依赖于累积序列长度跟踪，这使得 SDPA 和 RoPE 等操作能够在尊重序列边界的情况下高效运行。

6.4.2. 用于打包序列的动态上下文并行

动态 CP 以解决不平衡问题。虽然打包序列消除了填充浪费，但并未消除数据并行（DP） rank 间的计算不平衡（图 36、37）。我们引入了动态上下文并行（Dynamic-CP）【90, K. Li 等人, "Speeding up variable-length training with dynamic context parallelism and nvidia megatron-core.", 2025】，它根据每个微批次的情况自适应地选择有效的 CP 大小。Dynamic-CP 将 CP 大小调整视为一种轻量级的运行时选择，避免了任何昂贵的模型状态重新分布（图 38）。

对于可变长度序列，损失在逐令牌的基础上计算，以避免填充带来的偏差：

$$\mathcal{L}=\frac{\sum_{t \in \mathcal{V}} \ell_{t}}{|\mathcal{V}|},$$

根据基准测试，Dynamic-CP 在序列长度分布高度不平衡的真实场景中，可带来 35-60% 的端到端性能提升。

---

A7 补充细节

7. 生产特性

负载均衡、令牌丢弃与共享专家。本节描述了满足生产需求的特性，包括负载均衡、共享专家架构、分布式检查点和多令牌预测等。

7.1. 负载均衡和令牌丢弃

两种机制协同工作。动态路由会导致工作负载不平衡。我们通过负载均衡和令牌丢弃来解决这些挑战。
* 负载均衡：图 39 说明了 Megatron-Core MoE 中可用的负载均衡策略。主要方法是使用辅助损失【12, D. Lepikhin 等人, "Gshard: Scaling giant models with conditional computation and automatic sharding", arXiv preprint arXiv:2006.16668, 2020; 4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022】，我们还支持专家选择路由【93, Y. Zhou 等人, "Mixture-of-experts with expert choice routing", 2022】和通过可学习的专家偏置项实现的无辅助损失均衡【22, L. Wang 等人, "Auxiliary-loss-free load balancing strategy for mixture-ofexperts", 2024】。
* 令牌丢弃：我们支持两种分派策略。在无丢弃模式（默认）下，所有路由的令牌都被处理。在可丢弃模式下，强制执行显式的专家容量限制【12, D. Lepikhin 等人, "Gshard: Scaling giant models with conditional computation and automatic sharding", arXiv preprint arXiv:2006.16668, 2020; 4, W. Fedus 等人, "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity", 2022】，超过容量的令牌被丢弃。可丢弃模式还启用了 pad-to-max 功能，将所有专家输入填充到相同的容量限制，从而将 MoE 固有的动态令牌计数转换为静态形状，以启用 CUDA Graphs 等优化。

7.2. 共享专家

并行计算以隐藏延迟。一些 MoE 架构（DeepSeek-V2/V3【17, DeepSeek-AI 等人, "Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model", 2024; 6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】、Qwen【32, Qwen 等人, "Qwen2.5 technical report", 2025】）包含一个共享专家，它处理所有令牌（图 40）。当启用重叠（--moe-shared-expert-overlap）时，共享专家的计算与 all-to-all 通信和路由专家的计算并行运行，从而隐藏其延迟。

7.3. 隐空间 MoE (Latent MoE)

降低通信和参数成本。LatentMoE【19, V. Elango 等人, "Latentmoe: Toward optimal accuracy per flop and parameter in mixture of experts", 2026】通过在分派前插入一个共享的下投影 $W_{\downarrow} \in R^{\ell \times d}$ 和在合并后插入一个上投影 $W_{\uparrow} \in R^{d \times \ell}$ 来降低成本，其中 $\ell < d$ 是隐空间维度。路由仍在完整隐藏维度上操作，但每个路由专家完全在压缩的隐空间中操作。层输出变为：

$$\text{output}(\mathbf{x}) = W_{\uparrow} \cdot \left( \sum_{i \in \mathcal{T}_{K,E}} p_{i} E_{i}(W_{\downarrow} \cdot \mathbf{x}; \ell) \right) + \sum_{j} E_{j}^{\text{shared}}(\mathbf{x}; d).$$

7.4. 分布式检查点

与并行配置无关的检查点。我们的分布式检查点库通过与并行配置无关的检查点和自动重分片解决了传统检查点的局限性。核心抽象是 ShardedTensor 描述符，它编码了每个本地张量的全局形状、偏移和分片模式。这使得可以进行任意到任意的并行重配置。

7.5. 灵活的非对称虚拟流水线并行

处理工作负载异构性。传统的 VPP 要求均匀的层分布。然而，MoE 模型表现出显著的工作负载异构性。我们引入了灵活的非对称 VPP，允许每个虚拟阶段有不同数量和类型的层。表 10 展示了 DeepSeek-V3 的层分布。这种方法（图 41）可以实现真正的负载均衡，并为特殊层提供灵活的放置策略。

7.6. 升级改造 (Upcycling)

从密集模型到稀疏 MoE。升级改造将预训练的密集模型转换为稀疏的 MoE 架构，无需从头开始重新训练【95, A. Komatsuzaki 等人, "Sparse upcycling: Training mixture-of-experts from dense checkpoints", arXiv preprint arXiv:2212.05055, 2022; 96, E. He 等人, "Upcycling large language models into mixture of experts", 2024; 97, A. Vavre 等人, "Llama 3 meets moe: Efficient upcycling", arXiv preprint arXiv:2412.09952, 2024】。我们支持虚拟组初始化和专家权重缩放，以实现从密集检查点到细粒度 MoE 模型的无缝适应（图 42）。

7.7. 多令牌预测

密集化监督信号。多令牌预测（MTP）【98, F. Gloeckle 等人, "Better & faster large language models via multi-token prediction", Proceedings of the 41st International Conference on Machine Learning (ICML), 2024; 6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】优化模型以在每个位置预测多个连续的未来令牌，从而加速收敛并提高生成质量。我们将 MTP 与灵活的流水线并行集成，允许 MTP 层在 VPP 布局中进行策略性放置。

7.8. Muon 优化器

矩阵感知的更新。Muon 优化器【99, J. Liu 等人, "Muon is scalable for llm training", 2025】通过对整个权重矩阵进行正交化，引入了一种矩阵感知的更新方法。我们的集成提供了对大规模分布式训练的生产级支持，包括对拆分 QKV 权重布局的完全支持、与分布式优化器的无缝集成以及 Muon 正交化缓冲区的 CPU 卸载。

---

A4 实验环境

8.1. 实验设置

• 基准模型: 我们在两种先进的细粒度 MoE 架构上评估 Megatron-Core MoE：DeepSeek-V3-685B【6, DeepSeek-AI 等人, "Deepseek-v3 technical report", 2025】和 Qwen3-235B【32, Qwen 等人, "Qwen2.5 technical report", 2025; 87, A. Yang 等人, "Qwen3 technical report", 2025】。这两种模型都采用了细粒度专家设计，放大了“三座大山”的挑战。

• 硬件和软件栈:

  • 硬件: NVIDIA GB200 和 H100 GPU。
  • 软件: Megatron-LM 的 dev 分支和最新的 TransformerEngine。
  • 启用的优化: 启用了第 4 节中描述的完整优化栈，包括 FP8 训练（GB200/GB300 上为 MXFP8，H100 上为 blockwise FP8）、内存优化（选择性重计算、激活和优化器状态卸载）、优化的令牌分发器（GB200/GB300 NVL 系统上为 HybridEP，H100 系统上为 DeepEP）、通信重叠（1F1B all-to-all 重叠）以及核函数优化（Grouped GEMM、路由器融合、排列融合和 CUDA Graphs）。

• 指标: 我们报告了每个 GPU 的吞吐量（以 TFLOPS 和每秒处理的令牌数/GPU 为单位）。

A4 实验结果

8.2. 关键性能结果

• 概述: 表 11 总结了在 NVIDIA GB300、GB200 和 H100 系统上，启用完整优化栈后，代表性细粒度 MoE 训练工作负载的每 GPU 吞吐量。所有配置都使用强制负载均衡路由。

• DeepSeek-V3 和 Qwen3-235B (大规模训练): 表 11 报告了在标准序列长度 4096 下的训练结果，包括在 H100 上进行的 1024-GPU DeepSeek-V3 运行。在所有三个平台上，测量结果都显示出强大的每 GPU 效率和高持续吞吐量。

• FP8 训练有效性: 如表 11 所示，在 H100 上，DeepSeek-V3 的 FP8 训练达到了每 GPU 368 TFLOPS，证明了 blockwise FP8 方案能够在 Hopper 架构上高效训练大规模 MoE 模型。

• GB200 和 GB300 性能: 与 H100 相比，GB200 和 GB300 平台在两种 MoE 模型上，以相当或更少的 GPU 数量，提供了大约 3 倍的令牌吞吐量。这得益于卓越的内存带宽、更高的计算能力和原生的 MXFP8 Tensor Core 支持。在 GB200 上，DeepSeek-V3 训练达到了每 GPU 超过 1048 TFLOPS，在 GB300 上进一步推高到 1233 TFLOPS。

• 可扩展性: 这些结果表明，Megatron-Core MoE 能够高效地扩展到数千个 GPU 和数千亿参数的生产工作负载。

• 长上下文压力测试: 表 11 还包括在 GB300 上以 131,072 序列长度运行的长上下文 Qwen3-235B 测试。即使在这种内存和通信密集型设置下，系统仍能维持每 GPU 1150 TFLOPS 的性能。

• 配置细节: 每个基准测试条目的并行配置、批处理设置和其他优化细节见附录 B。

---

A7 补充细节 (续)

9. 性能最佳实践

系统化的优化方法。本节将优化 MoE 模型的经验提炼为可操作的指南，分为一个三阶段的工作流程，并通过 DeepSeek-V3 的案例研究进行说明。

9.1. 系统化的优化工作流程

我们提出了一个三阶段的 MoE 性能优化工作流程。这个过程本质上是迭代的。

9.1.1. 阶段 1：建立内存可行的并行配置

内存可行性是首要约束。在优化吞吐量之前，配置必须能装入 GPU 内存。表 12 总结了各种并行策略对每 GPU 内存和通信开销的影响。

9.1.2. 阶段 2：选择最优并行策略

最小化通信开销。一旦确定了内存可行的配置，就选择能最小化通信开销同时保持吞吐量的策略。
- 指南 1：最小化模型并行，最大化数据并行：目标是保持 TP/EP/PP/CP 尽可能小，同时避免 OOM。
- 指南 2：将 EP 和 TP 通信保持在 NVLink 域内：确保 EP×TP 适合 NVLink 域（通常是单个节点中的 8 个 GPU）。
- 指南 3：使用流水线并行（PP）进行多节点扩展：使用 PP 跨节点分布层，同时保持 EP×TP 在 NVLink 内。
- 指南 4：对于专家层，优先选择 EP 而非 TP：EP 为 MoE 层提供了更好的 GEMM 效率和更低的通信开销。
- 指南 5：为长序列启用上下文并行（CP）：当序列长度 ≥ 8K 令牌时使用。

9.1.3. 阶段 3：剖析和优化瓶颈

识别主导瓶颈。建立可行的并行配置后，对训练运行进行剖析，以确定是哪个瓶颈（内存、通信或计算）占主导地位，并应用有针对性的优化。
- 内存瓶颈：症状是为避免 OOM 而被迫使用完全重计算或过大的并行度。解决方案见表 13。
- 通信瓶颈：症状是剖析显示大量时间花在集合操作上。解决方案见表 14。
- CPU 开销瓶颈（计算效率墙）：症状是 Nsight Systems 时间线显示 GPU 内核之间存在间隙。解决方案见表 15。
- 计算瓶颈（计算效率墙）：症状是尽管没有通信或 CPU 瓶颈，但 GPU SM 利用率仍然很低。解决方案见表 16。

9.2. 案例研究：在 GB200 和 H100 上调优 DeepSeek-V3

DeepSeek-V3 是一个极限压力测试。本案例研究探讨了优化决策如何相互作用，形成一个连贯的优化栈，以及为什么同一个模型在 GB200 和 H100 上需要根本不同的策略。

9.2.1. 最终优化配置和性能

表 17 总结了 GB200 和 H100 平台上的最终优化配置和性能。表 18 总结了在每个平台上应用的关键优化。

9.2.2. 优化配置剖析

并行布局：两个平台都使用并行折叠。关键的平台差异在于内存容量。GB200 的 192 GB 内存允许使用 TP1/PP4，而不是 H100 上的 TP2/PP8，从而减少了流水线深度和气泡。
内存墙：FP8 训练是第一个杠杆。在 H100 上，释放的内存预算对于启用 EP 通信重叠至关重要。在 GB200 上，更高的 C2C 带宽使得优化器状态卸载非常有效。
通信墙：在 H100 上，EP64 跨越 8 个节点，EP 通信重叠对于隐藏跨节点延迟至关重要。在 GB200 (NVL72) 上，EP64 完全在 NVLink 域内，通信墙由硬件拓扑单独解决。
计算效率墙：FP8 加速了 GEMM，但也暴露了 CPU 开销。在 GB200 上，CPU 开销成为主要限制。部分 CUDA Graphs 和核函数融合解决了这个问题。

9.2.3. 经验教训

1. 平台特性驱动策略：GB200 的更大内存和更高带宽支持更激进的选择。
2. 并行折叠释放了灵活性：解耦注意力 TP 和专家 EP 允许独立优化。
3. FP8 改变了瓶颈：FP8 训练加速了 GEMM 并减少了内存，但放大了 CPU 开销。
4. 迭代优化：优化过程本质上是循环的，需要持续的剖析和调整。

10. Megatron-Core MoE 在强化学习中的应用

RL 后训练的兴起。强化学习（RL）后训练已成为一种关键范式。许多领先的 RL 训练模型都是 MoE 架构，而 RL 工作负载引入了新的挑战。

10.1. RL 后训练的挑战

1. 可变长度序列：RL 工作负载产生高度可变的序列长度，使得平衡计算效率和峰值内存消耗变得困难。
2. 内存卸载：RL 框架通常在同一 GPU 上共存训练和推理引擎，需要快速、完整地释放和恢复内存。
3. 在线权重导出：推理引擎必须在每个训练步骤后用最新的参数进行更新。
4. 训练稳定性：推理和训练引擎使用不同的优化内核，可能会导致路由决策不一致，从而引入非策略偏差。

10.2. Megatron-Bridge

检查点互操作性。Megatron-Bridge 解决了检查点互操作性层的问题，实现了 HF 到 Megatron 检查点格式的快速转换，用于初始化、训练和导出。

10.3. Megatron-Core 的强化学习优化

• 打包序列支持：如 6.4.1 节所述，我们可以打包序列以消除可变序列长度批次中的填充。
• 动态上下文并行：如 6.4.2 节所述，Dynamic-CP 通过在每个微批次上自适应地选择 CP 度来消除“一刀切”的限制。
• CPU 优化器卸载：为进一步缓解 GPU 内存压力，我们将优化器状态卸载到主机 DRAM。
• FP16 训练支持：Megatron-Core MoE 实现提供了功能齐全的 FP16 路径。
• 路由器重放：最近的研究【104, S. Fan 等人, "Router replay: Improving sample efficiency in mixture-of-experts reinforcement learning", 2025】表明，捕获推理期间 MoE 路由器的决策并在后续训练阶段重放它们，可以提高收敛一致性。Megatron-Core MoE 现在支持此功能。

---

A5 结论

11. 结论

本文介绍了 Megatron-Core MoE，一个用于训练大规模专家混合模型的开源技术栈。MoE 的稀疏性引入了两个基本挑战：表现为“三座大山”（内存、通信和计算效率）的参数-计算不匹配，以及需要为注意力和 MoE 层解耦并行的密集-稀疏不匹配。通过结合针对这些挑战的集成解决方案，Megatron-Core MoE 实现了高效的万亿参数级训练。

关键技术贡献包括：
- 多维并行与并行折叠：打破了 EP ≤ DP 的限制，实现了灵活的并行映射。
- 内存优化：将 DeepSeek-V3 的每 GPU 内存需求从 199.5 GB 降低到 80 GB 以下。
- 通信优化：将 all-to-all 从瓶颈转变为隐藏在专家计算背后的后台操作。
- 计算效率：通过 Grouped GEMM、核函数融合、CUDA Graphs 和无同步执行来解决。
- FP8 训练：在所有三座大山上同时提供收益。
- 长上下文训练：通过上下文和张量并行扩展，支持 16K 到 64K+ 令牌序列的训练。
- 生产特性：负载均衡、分布式检查点、升级改造和多令牌预测集成。
- 强化学习支持：与流行的 RL 框架无缝集成，并通过打包序列、动态上下文并行和路由器重放解决 RL 的独特挑战。

这些优化使得 Megatron-Core MoE 实现了强大的训练吞吐量：DeepSeek-V3（685B 参数）在 256 个 GB300/GB200 上分别以 1233/1048 TFLOPS/GPU 的速度训练；Qwen3-235B 在 GB300/GB200 上分别达到 974/919 TFLOPS。通过开源这些能力，Megatron-Core MoE 为研究人员和从业者提供了在规模上进行 MoE 实验和部署的生产级工具。

---

A6 附录

A. 符号参考

表 19 总结了本报告中使用的符号。

B. 详细基准配置

重现性能数据。本节列出了用于重现表 11 中报告的性能数据的并行配置和详细设置。

B.1. 配置细节

表 20 列出了与表 11 中吞吐量结果相对应的基准配置。这些设置是撰写本文时通过经验调优得到的最佳配置，可能并非全局最优。

B.2. 关键优化

性能关键特性。本节总结了我们基准测试中使用的对性能最关键的优化特性的配置。

DeepSeek-V3
* 分发器: GB300/GB200 上为 HybridEP；H100 上为 DeepEP。
* 重计算: GB300 上无；GB200 上为 mlp；H100 上为 up_proj, mlp。
* 1F1B 重叠: GB300 和 H100 上开启；GB200 上关闭。
* CUDA Graphs: GB300/GB200 上为 attn, moe_router, moe_preprocess；H100 上关闭。

Qwen3-235B
* 分发器: GB300/GB200 上为 HybridEP；H100 上为 DeepEP。
* 重计算: GB200 和 H100 上为 moe_act, layernorm。
* 1F1B 重叠: GB300 上关闭；GB200 和 H100 上开启。
* CUDA Graphs: attn, moe_router, moe_preprocess。

B.3. 可复现性

两种支持的工作流程。表 11 中的基准数据可以通过两种支持的工作流程来复现。首先，用户可以使用 Megatron-Bridge 运行端到端训练。其次，用户可以使用 Megatron-MoE-ModelZoo 中的模型特定脚本直接用 Megatron-Core（MCore）启动训练。