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

发表时间: 2026-03 · Tech report by NVIDIA (arXiv:2603.07685)

NVIDIA

速读

一句话结论 本文提出了开源全栈优化框架 Megatron-Core MoE，通过协同设计显存、通信和计算，在 NVIDIA GB300 上训练 685B 参数的 DeepSeek-V3 模型时实现了 1233 TFLOPS/GPU 的极高吞吐量。

要解决什么问题 混合专家（MoE）模型的稀疏性导致了根本的“参数-计算不匹配”：总参数量随专家数庞大，但每个 Token 仅激活极少参数。这种机制在训练时撞上了三堵相互耦合的限制墙。首先是显存墙，所有专家的参数、梯度和状态必须常驻显存，带来远超同等计算量密集模型的显存压力；其次是通信墙，为分散显存压力必须引入专家并行（EP），这要求全互联通信，随着并行度增加，通信会从节点内高带宽链路退化为跨节点低带宽互连；最后是计算效率墙，细粒度专家产生大量小型矩阵乘法，无法喂饱 GPU，且动态路由带来的张量形状变化会产生极高的主机端调度开销。此外，注意力层依赖张量并行，MoE 层依赖专家并行，传统框架强制两者共享并行配置，导致显存浪费和次优切分。

怎么做的 核心思路是通过解耦和全栈协同设计，打破显存、通信与计算的耦合约束。针对密集层与稀疏层的并行冲突，框架引入并行折叠（Parallel Folding）技术，将注意力层和 MoE 层的并行映射彻底解耦，允许专家并行跨越节点内的高带宽组进行“折叠”，打破了专家并行度必须小于等于数据并行度的限制，避免了跨节点通信。为推倒显存墙，框架设计了内存高效排列技术，将路由器权重 $p_i$ 吸收进激活函数中：$$ \mathbf{y} = \sum_{i \in \mathcal{T}(\mathbf{x})} \mathbf{W}_2^{(i)} (p_i \cdot \phi(\mathbf{W}_1^{(i)} \mathbf{x})) $$这一等价变换消除了反向传播时保存路由器张量的需求，实现零计算开销的显存缩减；同时配合双 DeviceMesh FSDP 架构，为主干密集模块和专家模块分别建立独立的切分网格。在打破通信墙方面，框架集成 DeepEP 和 HybridEP 分派器，摒弃全局排列步骤，直接基于 Token 路由信息通过队列进行点对点传输；并设计 1F1B 前后向重叠机制，将反向传播的权重梯度计算与前向计算重叠，隐藏全互联通信延迟。针对计算效率墙，框架采用 Grouped GEMM 将多个专家的零碎矩阵乘法融合成单次调用；针对动态路由导致无法使用 CUDA Graphs 的卡点，引入分层 CUDA Graphs 仅捕获静态组件；配合热点专家弹性克隆（ECHO）机制，在分派阶段将高负载专家的权重复制到空闲 GPU 上，抹平动态路由带来的计算不平衡与显存碎片化。

效果如何 实验在 NVIDIA GB300、GB200（NVL72 拓扑）及扩展至 1024 张卡的 H100（NVL8 拓扑）集群上展开。评估目标为两个最先进的细粒度 MoE 模型：685B 参数的 DeepSeek-V3 和 235B 参数的 Qwen3。由于本文是全栈系统设计，对比基线体现为框架内不同优化路线的消融，如在通信后端上对比 DeepEP 与 HybridEP 路线，在执行模式上对比常规执行与分层 CUDA Graphs 路线。在 256 张 GB300 和 GB200 上，使用 MXFP8 精度训练 DeepSeek-V3 分别达到 1233 和 1048 TFLOPS/GPU 的极高吞吐量；在 1024 张 H100 上使用 FP8-BLK 精度达到 368 TFLOPS/GPU。长上下文压力测试中，128 张 GB300 以 131072 的超长序列训练 Qwen3，维持了 1150 TFLOPS/GPU 的吞吐量。作者也揭示了系统局限：相同模型在不同硬件上需要截然不同的优化栈。在 GB200 上，通信不是瓶颈，系统受限于 CPU 开销，必须采用 HybridEP 配合 CUDA Graphs；而在 H100 上，跨节点通信成为短板，必须切换为 DeepEP 并开启 1F1B 通信重叠隐藏延迟。此外，为满足 CUDA Graphs 对静态形状的要求，系统在某些模式下必须引入 Token 丢弃并填充至最大长度，带来额外开销。

主要贡献

扩展混合专家（MoE）模型的训练引入了密集模型中不存在的系统挑战。由于每个Token仅激活专家的一个子集，这种稀疏性使得总参数量的增长速度远快于每个Token的计算量，从而在内存、通信和计算之间产生了耦合约束。优化其中一个维度通常会将压力转移到另一个维度，这要求跨整个系统栈进行协同设计。

本文通过涵盖内存（细粒度重计算、卸载等）、通信（优化的分派器、重叠等）和计算（Grouped GEMM、内核融合、CUDA Graphs等）的集成优化来解决MoE训练的这些挑战。该框架还提供了用于灵活多维并行的并行折叠（Parallel Folding）技术、对FP8和NVFP4的低精度训练支持，以及高效的长上下文训练。在NVIDIA GB300和GB200上，它在训练DeepSeek-V3-685B时分别达到了1,233和1,048 TFLOPS/GPU，在训练Qwen3-235B时达到了974和919 TFLOPS/GPU。作为一个高性能、可扩展且可用于生产的开源解决方案，它已被学术界和工业界广泛用于在扩展至数千个GPU的集群上训练从数十亿到数万亿参数的MoE模型。

背景知识与关键观察

混合专家（MoE）架构基础 混合专家（MoE）模型通过引入称为专家的专用子网络集合以及一个动态选择哪些专家处理每个输入的轻量级路由器（或门控网络），来增强标准神经网络【3，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017】。在基于Transformer的语言模型中，MoE层通常取代密集的逐位前馈网络（FFN）块：MoE层包含多个FFN专家，每个Token根据学习到的路由权重被路由到这些专家的一小部分（例如top-$k$）。给定输入Token表示$\mathbf{x}$，路由器计算$E$个专家的概率分布：

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

MoE层的输出随后计算为所选专家输出的加权组合：

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

其中$E_i$表示第$i$个专家网络。这种架构提供了三个关键优势：模型容量可以通过添加更多专家来独立于计算成本进行扩展；每个Token只有一小部分参数被激活，相对于同等规模的密集模型减少了FLOPs；不同的专家可以专注于不同类型的输入。

大规模MoE模型训练的挑战 随着模型向具有数百个较小个体容量的专家发展，MoE训练的系统挑战成比例增长。这些挑战源于一个根本原因：MoE的稀疏性，这表现为参数-计算不匹配（总参数远超活跃计算）和密集-稀疏不匹配（注意力和MoE层需要冲突的并行配置）。在密集Transformer中，每个参数都参与每个训练步骤，参数和每个Token的计算量同步扩展。而对于MoE，由于稀疏性导致这种耦合破裂：每个Token只有$E$个专家中的$K$个被激活，因此每个Token的计算量大约是$6N_{\mathrm{active}}$而不是$6N_{\mathrm{total}}$，其中$N_{\mathrm{active}}$随$K$扩展而$N_{\mathrm{total}}$随$E$扩展，且$K \ll E$。这创造了根本的参数-计算不匹配，例如DeepSeek-V3有685B总参数，但每个Token仅激活37B参数，存在$18\times$的差距。由于每个Token的计算量如此之少，模型划分需要比密集模型更加谨慎。因为MoE专家是独立的网络，自然的策略是专家并行（EP），但这引入了全互联（all-to-all）通信。

三大相互耦合的限制墙 参数-计算不匹配和EP的通信需求共同创造了限制每个MoE训练步骤的三个紧密耦合的挑战。首先是内存墙，所有$E$个专家的参数、梯度和优化器状态必须驻留在内存中，即使每个Token只激活$K$个，这创造了远超具有可比每Token计算量的密集模型的内存压力【4，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022】。其次是通信墙，EP需要全互联集合通信以将Token分派给指定的专家并收集结果【12，Gshard: Scaling giant models with conditional computation and automatic sharding，2020】，每个GPU的发送量约为$T \cdot K \cdot h \cdot \frac{EP - 1}{EP}$，且随着EP的增长，通信逐渐从高带宽的节点内链路转移到带宽下降一个数量级的节点间互连。最后是计算效率墙，MoE引入了密集模型中不存在的计算低效性，包括细粒度专家产生的小型GEMM无法充分利用GPU计算单元【24，MegaBlocks: Efficient Sparse Training with Mixture-of-Experts，2023】、路由和排列开销、动态路由导致的不平衡，以及稀疏性和动态张量形状带来的大量主机开销。这三堵墙是紧密耦合的，优化其中一个通常会将压力转移到另一个。

方法细节

Megatron-Core MoE架构

MoE层数据流的四个连续阶段 MoE层通过四个连续阶段处理输入Token。首先是路由（Route）阶段，TopKRouter确定哪些专家处理每个Token，它通过学习到的线性投影将每个Token的隐藏状态映射为$E$个对数几率，通过得分函数转换为概率，并选择得分最高的top-$k$专家。路由器输出包含路由权重的probs张量和指示Token-专家分配的布尔掩码routing_map。接着是分派（Dispatch）阶段，Token分派器准备用于跨GPU通信的Token，它先对Token进行排列，使发往同一专家的所有Token连续，然后使用通信后端将Token移动到托管其分配专家的GPU。随后是专家计算（Expert Computation）阶段，每个GPU对其接收到的Token执行其本地专家，所有本地专家通过TEGroupedMLP在单个Grouped GEMM调用中运行以最大化GPU利用率。最后是组合（Combine）阶段，逆向通信将处理后的Token返回到其原始GPU，随后进行反排列以恢复原始序列顺序，如果配置了共享专家，其输出将在此阶段相加。
MoE层的数据流：路由、分派、计算和组合阶段

路由器架构与Token分配 路由器通过两个操作将全局Token批次转换为特定于专家的工作负载。第一个是门控（Gating），线性投影$\mathbf{W}_r \in \mathbb{R}^{h \times E}$将每个Token的隐藏状态$\mathbf{x} \in \mathbb{R}^h$映射到对数几率$\mathbf{l} = \mathbf{W}_r^\top \mathbf{x} \in \mathbb{R}^E$。第二个是Top-$k$选择，得分函数将对数几率转换为概率：softmax$(p_i = e^{l_i} / \sum_j e^{l_j})$或sigmoid$(p_i = \sigma(l_i) / \sum_j \sigma(l_j))$，为每个Token选择得分最高的top-$k$专家。
路由器架构：线性投影、得分函数、top-k选择和负载均衡

Token分派器的通信抽象 分派器通过六个阶段的流水线管理GPU之间的Token移动：前向的dispatch_preprocess、token_dispatch、dispatch_postprocess，以及反向的combine_preprocess、token_combine、combine_postprocess。框架提供三种分派器类型：AllGather（每个GPU收集所有Token并为本地专家过滤，简单但内存密集）、All-To-All（基于标准NCCL的点对点通信，扩展性好但有同步开销），以及Flex（支持如DeepEP和HybridEP等优化后端的统一设计）。

专家计算模块与共享专家 每个专家是一个带有可选门控的两层MLP。Grouped GEMM使得多个专家计算的批处理变得高效。提供两种实现：TEGroupedMLP（Transformer Engine优化的实现，支持FP8和FP4量化）和SequentialMLP（在循环中一次执行一个专家，用于调试但速度明显较慢）。某些架构（如DeepSeek-V2/V3）包含一个处理所有Token的共享专家，其计算可与分派-计算-组合流水线并行运行以隐藏延迟。

并行组管理与优化器处理 MoE层需要与密集层不同的进程组，因为不同组件有不同的通信模式。Megatron-Core通过ProcessGroupCollection组织这些组：路由器使用tp, cp, tp-cp组；Token分派器使用ep, tp-ep组；专家使用ep, expt_tp, expt_dp组；共享专家使用tp组。这种关注点分离使得并行折叠成为可能。专家参数在分布式优化中需要特殊处理，它们被标记为allreduce := False以区别于密集参数。密集层在dp_cp_group中规约梯度，而专家在expt_dp_group中规约，确保梯度在正确数量的副本上平均。专家梯度按edp_size / dp_size进行缩放，以解释专家与密集层看到的有效批大小的差异。

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

传统并行策略在MoE中的失效 传统的并行策略是为密集Transformer设计的，而MoE模型具有根本不同的计算模式。将张量并行（TP）应用于MoE专家会产生更小的分片，降低计算效率并增加通信占比。由于MoE模型参数量巨大，仅使用流水线并行（PP）需要许多流水线阶段，这会产生过多的流水线气泡并降低吞吐量。仅使用数据并行（DP）需要每个GPU保存完整模型，这对于万亿参数模型是不可能的。

专家并行（EP）的引入与优势 专家并行（EP）将专家分布在各个GPU上【12，Gshard: Scaling giant models with conditional computation and automatic sharding，2020】。当EP度等于$E$时，每个GPU持有$E / \mathrm{EP}$个专家。前向传递依次进行路由、all-to-all分派、本地专家计算和all-to-all组合。EP提供了两个关键优势：首先，将来自不同GPU的Token分组到单个专家增加了计算强度，提高了GEMM效率；其次，随着专家数量的增加，all-to-all通信量保持不变，只有GPU的数量发生变化。
专家并行(EP)将专家分布在各个GPU上

密集-稀疏不匹配与传统框架的局限性 单个Transformer块包含两种根本不同的计算模式：注意力层受益于高TP和高CP，而MoE层需要高EP且高TP会适得其反。传统框架强制注意力和MoE层共享相同的并行配置，将EP视为DP的子维度（$\mathrm{EP \subseteq DP}$）。这带来了三个关键挑战：乘法性的GPU需求（例如请求$\mathrm{EP}=8$强制$\mathrm{DP} \ge 8$，结合$\mathrm{CP}=8$导致最小需要64个GPU）、强制次优的并行性（必须在分片注意力矩阵和保留专家计算效率之间妥协），以及跨节点通信（高EP通常迫使all-to-all通信跨越节点边界）。

MoE并行折叠（Parallel Folding）的解耦映射 Parallel Folding通过解耦注意力和MoE层的并行映射来解决密集-稀疏不匹配问题【26，Moe parallel folding: Heterogeneous parallelism mappings for efficient large-scale moe model training with megatron core，2025】。注意力层在$\mathrm{TP \times CP \times DP \times PP}$上形成组，而MoE层在$\mathrm{ETP \times EP \times EDP \times PP}$上形成组。唯一的约束是流水线并行（PP）必须在两种布局中保持一致以确保正确的梯度流。
并行映射：传统约束与MoE并行折叠解耦对比

并行折叠的四大核心优势 这种解耦设计提供了四个关键优势：首先，它打破了$\mathrm{EP \le DP}$的限制，允许EP通过跨越$\mathrm{TP \times CP}$组进行“折叠”来超过DP；其次，它降低了最小GPU需求，因为CP和EP可以共享同一个GPU组；第三，它实现了独立优化，注意力可以使用高TP而MoE使用$\mathrm{ETP}=1$以保持完整的专家宽度；第四，它将高带宽通信保持在NVLink域内，避免了较慢的跨节点传输。此外，Megatron-Core还将分布式优化器和FSDP与EP支持集成，进一步减少内存足迹。
并行折叠：解耦注意力和MoE并行映射

打破内存墙

内存高效排列（Memory-Efficient Permutation）实现零开销激活减少 传统的路由器权重在专家计算后应用：$\mathbf{y} = \sum_{i \in \mathcal{T}(\mathbf{x})} p_i \cdot \mathbf{W}_2^{(i)} \phi(\mathbf{W}_1^{(i)} \mathbf{x})$。内存高效排列将$p_i$吸收到激活中，在第二个线性层之前应用它：$\mathbf{y} = \sum__i$都必须为SwiGLU的自身反向传播保存，因此没有引入额外的缓冲区，实现了零计算开销的净峰值内存减少。}(\mathbf{x})} \mathbf{W}_2^{(i)} (p_i \cdot \phi(\mathbf{W}_1^{(i)} \mathbf{x}))$。当专家没有偏置项时，这两个公式在数学上是等价的。这种重新排列通过消除路由器反向传播所需的保存张量来降低峰值内存。在内存高效公式中，$p_i$直接乘以$\phi(\mathbf{z}_i)$，因此$\partial \mathcal{L} / \partial p_i$仅依赖于$\phi(\mathbf{z}_i)$，融合的反向内核可以在运行中从$\mathbf{z}_i$重新计算它。由于无论是否使用该技术，$\mathbf{z
内存高效排列

细粒度激活重计算（Fine-grained Recomputation） 与对大型单片部分应用激活检查点不同，Megatron-Core引入了细粒度重计算，仅针对内存密集且计算成本低的操作【27，Reducing activation recomputation in large transformer models，2022】。它结合了颗粒状重计算（用户指定确切重计算哪些计算，如专家MLP内的激活函数或LayerNorm）和输出丢弃重计算（在下游层消耗后立即释放检查点模块的输出，并在反向传播时从重计算结果中恢复）。这确保了内存不会不必要地保留给可以廉价恢复的激活。
选择性重计算

细粒度激活卸载（Fine-grained Activation Offloading） 当GPU内存仍然不足时，细粒度激活卸载通过专用的D2H/H2D CUDA流将激活传输与后续计算并发运行。前向传播期间，输入激活在模块计算后立即异步卸载到CPU。反向传播期间，系统采用层交错重载（Layer-Staggered Reload）模式，在计算当前层梯度的同时，从下一层重载相同模块的激活。由于重载发生在每个模块的反向传播完成后，任何时候GPU内存中只有一个模块类型的激活驻留。对于PP/VPP场景，ChunkOffloadHandler管理每个（微批次，VPP阶段）组合的卸载/重载逻辑，通过逆序（FILO）的VPP阶段和顺序（FIFO）的微批次队列，正确地将卸载的激活与其对应的反向传播配对。
细粒度激活卸载：前向和反向传递的流重叠
细粒度卸载与重计算：互补的内存优化策略

精度感知优化器与状态卸载 精度感知优化器将存储精度与计算精度解耦【57，8-bit optimizers via block-wise quantization，2022】。第一和第二动量可以存储在BF16或FP8中，在每个优化器步骤期间，这些较低精度的状态在FusedAdam内核中动态转换为FP32，并使用全精度算术计算更新以保持数值稳定性。这使优化器状态内存减少了约50%。此外，状态卸载在optimizer.step()之后将优化器状态和主权重传输到CPU，并在下一步之前重载它们，从而在不处于活动状态时回收GPU内存。

用于MoE的双DeviceMesh FSDP设计 Megatron-FSDP通过双DeviceMesh架构解决了密集模块和专家模块需要不同分片范围的核心挑战【53，Pytorch fsdp: Experiences on scaling fully sharded data parallel，2023】。主DeviceMesh管理密集模块的全局DP分片，而辅助的Expert DeviceMesh管理EP模块，将FSDP范围限制在EDP维度内。为了消除通信缓冲区管理和数据拷贝的开销，Megatron-FSDP在模块内展平并连接所有参数，然后跨设备应用非均匀分片，使分片边界与通信缓冲区布局对齐。此外，它在训练开始时预分配两个持久缓冲区并在它们之间循环，通过用户缓冲区注册（UBR）将这些缓冲区注册到NCCL，实现真正的零拷贝通信。
分片策略比较：FSDP2与Megatron-FSDP非均匀分片
持久双缓冲区设计

打破通信墙

DeepEP和HybridEP最大化EP带宽 标准的EP实现依赖于all-to-all通信，需要排列阶段并发送冗余流量。基于Token的分派后端（DeepEP和HybridEP）消除了排列步骤并避免发送冗余Token【23，Deepep: an efficient expert-parallel communication library，2025】。对于分派，HybridEP根据路由信息将数据从全局内存读入共享内存，然后通过FIFO队列将Token写入目标。在节点间情况下，HybridEP首先使用RDMA warp group在跨节点具有相同本地索引的GPU之间交换数据，然后在每个节点内转发。对于组合，HybridEP将规约融合到通信内核中，通过FIFO队列读取数据，执行规约，并将结果直接写入目标位置。
4个GPU和4个专家的专家并行通信模式
HybridEP的分派内核设计
HybridEP的组合内核设计

EP通信重叠（1F1B FWD-BWD）隐藏延迟 为了将all-to-all延迟隐藏在计算背后，Megatron-Core使用了专用的1F1B all-to-all重叠方案【61，1f1b-based ep a2a overlapping for moe models，2025】。首选策略是合并的FWD-BWD，它将一个微批次的前向传递与另一个微批次的反向传递合并。为了最大化这种重叠，使用了流分离（在计算流和通信流之间交替任务）和W/D拆分（将反向MLP拆分为独立于反向分派的权重梯度计算W/mlp，以及馈入反向分派的数据梯度计算D/mlp）。这使得W/mlp可以与前向MLP（F/mlp）重叠，以隐藏反向分派延迟。为了将重叠扩展到大型模型，合并的FWD-BWD与交错流水线并行（Interleaved PP）结合，在进入1F1B阶段之前运行一个额外的微批次，确保相邻的FWD/BWD对没有依赖关系。
合并的FWD-FWD时间线与all-to-all重叠
合并的FWD-BWD时间线与all-to-all重叠
EP all-to-all通信重叠策略：基线与带有W/D拆分的1F1B对比
带有all-to-all重叠的交错PP时间线

打破计算效率墙

Grouped GEMM与内核融合提升内核效率 Grouped GEMM通过重叠内核的波尾效应，将专家的独立GEMM批处理以提高硬件利用率。提供了多流cuBLASLt GEMM和将各个GEMM融合为单个内核的CUTLASS Grouped GEMM。排列融合管道包含预处理（生成偏移映射）、排列（根据映射移动Token）和反排列（累加复制的Token）三个阶段。路由器和辅助损失融合将得分计算（包括top-$k$和softmax/sigmoid）、辅助损失的得分计算以及MoE辅助损失计算分解为三个部分，每个部分融合为一个内核。
训练过程中排列融合的管道
路由器融合的工作流

分层CUDA Graphs消除主机开销 由于dropless MoE中动态的每专家Token数会导致图捕获失败，Megatron-Core引入了分层（部分）CUDA Graphs。它仅捕获每层的静态组件（注意力、共享专家、路由器计算、EP预处理），而将动态的分派和专家GEMM操作留在图外。为了最小化CUDA Graphs带来的内存开销，实施了池共享（按照执行顺序捕获时共享一个内存池）和缓冲区重用（根据PP执行顺序跨图重用静态输入/输出缓冲区）。在不使用PP时，通过引入is_first_microbatch标志使微批次可以共享同一个图，将图的数量从$L \times M \times 2$减少到$L \times 2$。
传统执行与CUDA Graph执行对比
一个训练迭代中的完整与分层CUDA Graphs对比

Transformer层前向传递：未启用与启用部分CUDA Graphs的分析
为什么流水线并行阻止CUDA Graphs跨微批次共享

无同步内核、ECHO与分页暂存实现全图覆盖 为了使dropless MoE实现全CUDA Graphs覆盖，必须解决动态形状需要主机-设备同步以及预分配最坏情况缓冲区导致内存碎片化的问题。首先，开发了设备发起的GPU内核（如cuBLASLt Grouped GEMM和带HybridEP的无同步分派），直接从GPU内存读取形状信息，解耦了工作量与静态启动配置，消除了强制同步。其次，ECHO（热点专家弹性克隆）通过规划器生成热点专家映射，在分派阶段将热点专家权重克隆到利用率不足的等级上的空闲槽位，并在反向阶段规约梯度，从而减少负载不平衡和最坏情况下的缓冲区大小。最后，分页暂存（Paged Stashing）解耦了用于计算的单个最坏情况共享tmp缓冲区和以页面形式组织的暂存缓冲区。每层前向计算后，激活从tmp缓冲区复制到暂存缓冲区，仅存储实际Token数，将内存从$O(\mathrm{layers} \times \mathrm{worst_case})$减少到$O(\mathrm{worst_case} + \mathrm{actual_total})$，并通过专用流重叠内存传输。
前向和反向传递的ECHO工作流
三种执行模式下的内存布局比较
分页暂存流重叠

FP8/FP4低精度训练

选择性精度策略与训练配方 Megatron-Core采用选择性精度策略：路由器、嵌入、输出层、主梯度、主权重和优化器状态保持在高精度，而专家GEMM等大批量计算被量化。提供了三种经验证的FP8配方：每张量缩放（使用实时当前缩放而非延迟缩放以保证精度）、Hopper上的块级FP8（激活/梯度为$1 \times 128$块，权重为$128 \times 128$块，格式为E4M3），以及Blackwell上的MXFP8（$1 \times 32$粒度，E8M0缩放因子）。在Blackwell上还支持NVFP4【75，Pretraining large language models with nvfp4，2025】，采用两级微缩放（每张量FP32缩放和每块E4M3缩放），并结合随机Hadamard变换（RHT）、2D权重缩放和梯度随机舍入以保持训练稳定性。

具有各种FP8配方的线性层计算

FP8/FP4主权重消除冗余存储 传统的低精度训练维护FP32主权重、BF16模型权重和FP8/FP4计算权重。原生FP8/FP4通过建立从FP32主权重到FP8/FP4计算权重的直接转换路径，完全绕过了BF16中间层。在分布式优化器更新FP32参数分片后，直接计算绝对最大值，跨数据并行等级更新缩放因子，并将量化的FP8/FP4值写入模型参数缓冲区。这不仅减少了内存足迹，还将参数AllGather通信量减半。对于块级配方，实现了特定内核以在感知权重2D布局的情况下计算绝对最大值。
用于块级缩放的FP8主权重量化方案
用于延迟缩放和每张量当前缩放的FP8主权重量化方案

MoE特定的量化挑战与融合 FP8/FP4 GEMM要求张量维度对齐到特定倍数（如16或32）。为了避免显式填充和反填充内核的开销，系统通过填充路由映射（代价是发送少量额外Token）或将填充融合到排列内核中来实现动态形状对齐。为了减少量化多个输入张量带来的CPU开销，实现了分组量化内核。对于NVFP4，量化内核必须吸收训练配方逻辑：RHT与NVFP4量化融合以避免额外的BF16全局内存读写；在Wgrad路径中，内核同时吸收了转置操作；并且通过在Token排列内核中融合每专家128倍数的零填充，确保了CUDA Graph安全的分组NVFP4量化和GEMM操作。

长上下文MoE训练

计算重心的转移与内存管理 在长上下文（如16K到64K+ Token）中，缩放点积注意力（SDPA）展现出二次方$O(s^2)$复杂性，主导了计算，而MoE层仍呈线性$O(s)$增长。优化重点因此转向内存和通信。为了管理随序列长度增长的激活内存，结合使用上下文并行（CP）和张量并行（TP）来保持每个设备的子序列长度恒定。此外，采用优化器CPU卸载，并通过选择性重计算优先重计算非SDPA组件（如MLP相关模块），因为重计算SDPA成本过高。在CP和TP的选择上，通常建议在节点内使用TP和all-to-all CP以减少参数内存并提高SDPA效率，在节点间使用点对点（P2P）CP以保留通信-计算重叠。
SDPA表现出O(s^2)复杂性，而MoE表现出O(s)复杂性
TP和两种类型CP的通信和计算模式

打包序列与动态上下文并行（Dynamic-CP） 针对可变长度训练，支持打包序列（Packed Sequences），使用THD（总Token数×头数×维度）张量格式将多个序列连接在单个批次中，通过跟踪累积序列长度消除填充浪费。为了解决打包序列引起的DP不平衡和CP低效问题，引入了动态上下文并行（Dynamic-CP）【90，Speeding up variable-length training with dynamic context parallelism and nvidia megatron-core，2025】。它在每个微批次的基础上自适应选择有效的CP大小。框架在初始化时预构建多个CP组候选，调度器在运行时根据估计的工作量和内存约束，为长序列分配较大的CP大小，为短序列分配较小的CP大小，从而在不重新分配模型状态的情况下减少不必要的CP通信。
未打包与打包序列对比
打包序列上因果注意力的计算不平衡
用于打包序列的动态上下文并行

生产特性

负载均衡、Token丢弃与潜在MoE 提供辅助损失、专家选择路由【93，Mixture-of-experts with expert choice routing，2022】和无辅助损失均衡（专家偏置）【22，Auxiliary-loss-free load balancing strategy for mixture-of-experts，2024】等负载均衡策略。支持无丢弃（dropless）和可丢弃（droppable）模式，后者支持Pad-to-Max功能以满足CUDA Graphs的静态形状需求。潜在MoE（Latent MoE）【19，Latentmoe: Toward optimal accuracy per flop and parameter in mixture of experts，2026】通过在分派前插入共享向下投影并在组合后插入向上投影，在压缩的潜在空间中运行路由专家，将all2all通信量和专家权重大小减少压缩比$\alpha$倍。

共享专家架构

灵活的非对称VPP与分布式检查点 引入了灵活的非对称虚拟流水线并行（Flexible Asymmetric VPP），允许每个虚拟阶段包含不同数量和类型的层，从而实现真正的计算负载均衡并为特殊层（如MTP）提供灵活的放置策略。分布式检查点库通过ShardedTensor描述符实现了与并行性无关的检查点保存与加载，支持跨不同并行配置的自动重新分片。
灵活的流水线并行放置

升级（Upcycling）、多Token预测与Muon优化器 支持通过虚拟组初始化和专家权重缩放将预训练的密集模型转换为MoE架构。集成了多Token预测（MTP）【6，Deepseek-v3 technical report，2025】，并允许将MTP层战略性地放置在VPP布局中。集成了Muon优化器【99，Muon is scalable for llm training，2025】，提供全矩阵正交化支持、与分布式优化器的无缝集成以及CPU卸载功能。

补充细节

性能优化的三阶段工作流
性能优化遵循一个系统化的三阶段迭代工作流：
1. 建立内存可行的并行性：使用分布式优化器跨DP等级分片状态，优先增加数据并行度。将EP和TP通信保持在NVLink域内，并在跨节点扩展时优先使用PP。对于专家层，优先使用EP而非TP以获得更好的GEMM效率。长序列（$\ge 8\mathrm{K}$）启用CP。
2. 选择最佳并行策略：在内存可行的基础上，选择最小化通信开销的策略。
3. 分析和优化瓶颈：通过分析识别主导瓶颈。对于内存瓶颈，应用FP8、细粒度重计算和卸载；对于通信瓶颈，应用EP分派器和通信重叠；对于CPU开销瓶颈，启用CUDA Graphs和减少内核启动；对于计算瓶颈，应用Grouped GEMM和内核融合。

强化学习（RL）后训练支持
RL后训练带来了可变长度序列、内存卸载、在线权重导出和训练稳定性等挑战。Megatron-Bridge提供了Hugging Face与Megatron格式之间的双向检查点转换。为了优化RL训练，除了打包序列和Dynamic-CP外，还引入了按“注意力成本”（序列长度平方和）排序微批次的蛇形调度（小-大-小模式），以平衡计算时间并减少同步气泡。支持CPU优化器卸载和FP16训练以提高稳定性。最后，支持路由器重放（Router Replay）【104，Router replay: Improving sample efficiency in mixture-of-experts reinforcement learning，2025】，在训练中强制执行推理时记录的Token专家分配，以解耦路由可变性与权重更新，提高收敛一致性。

实验环境

数据集与模型架构：评估使用了两个最先进的细粒度MoE架构：DeepSeek-V3-685B（685B总参数，每个Token激活37B，256个专家，包含MTP和MLA）和Qwen3-235B（235B总参数，细粒度专家设计）。标准序列长度为4096，长上下文测试序列长度达到131,072。
硬件配置：NVIDIA GB300、GB200（NVL72拓扑）和H100（扩展至1024个GPU，NVL8拓扑）。
软件配置：Megatron-LM dev分支，最新的TransformerEngine。依赖于cuBLAS、cuDNN和NCCL等库。

实验结果

大规模MoE训练吞吐量：在256个GB300和GB200 GPU上，使用MXFP8精度训练DeepSeek-V3分别达到了1,233和1,048 TFLOPS/GPU。在1,024个H100 GPU上，使用FP8-BLK精度训练DeepSeek-V3达到了368 TFLOPS/GPU。对于Qwen3-235B，GB300和GB200分别达到了974和919 TFLOPS/GPU。GB300和GB200平台的Token吞吐量比H100高约3倍。
长上下文压力测试：在128个GB300 GPU上，以131,072的序列长度训练Qwen3-235B，系统维持了1,150 TFLOPS/GPU的高吞吐量，证明了长上下文优化的有效性。
DeepSeek-V3优化案例分析：相同的模型在不同硬件上需要截然不同的优化栈。在GB200（NVL72）上，EP64保留在NVLink域内，通信不再是瓶颈，CPU开销成为主要限制。因此，配置使用了HybridEP和CUDA Graphs，并减少了重计算。在H100（NVL8）上，EP64跨越了8个节点，跨节点通信成为瓶颈。因此，配置使用了DeepEP和EP通信重叠来隐藏延迟，这得益于FP8释放的内存。

结论

本文介绍了Megatron-Core MoE，这是一个用于训练大规模混合专家模型的开源全栈解决方案。通过将多维并行（Parallel Folding）、内存优化（细粒度重计算、卸载、双DeviceMesh FSDP）、通信优化（DeepEP/HybridEP、1F1B重叠）、计算效率提升（Grouped GEMM、CUDA Graphs、无同步执行）、FP8/FP4低精度训练以及长上下文和RL支持集成到一个连贯的系统中，Megatron-Core MoE成功打破了由MoE稀疏性引起的内存、通信和计算效率三大限制墙。该框架在最新的NVIDIA GPU上实现了卓越的训练吞吐量，为学术界和工业界在数千个GPU上训练万亿参数级别的MoE模型提供了生产级的工具和实践指南。