Megatron Core 的 MoE 训练:打破内存、通信与计算三堵墙
本文译自 NVIDIA 的技术报告 Scalable Training of Mixture-of-Experts Models with Megatron Core1,2026 年 3 月 8 日提交于 arXiv,3 月 10 日修订到 v2。原文 88 页、42 张图,这里是全文翻译,覆盖摘要、§1 引言、§2 架构、§3 Parallel Folding 与多维并行、§4 内存/通信/计算三堵墙、§5 FP8/FP4 低精度训练、§6 长上下文训练、§7 Production Features、§8 性能评测、§9 调优最佳实践、§10 RL 场景、§11 结论,以及附录 A 符号表与附录 B benchmark 配置。参考文献一节不译,正文里的引用走脚注。
这篇和 garden 里其他技术报告的定位不同:它不是某个模型的报告,而是 Megatron-Core 这套训练栈自己的报告。全文的主 benchmark 模型是 DeepSeek-V3,所有内存数字、通信占比、调优案例都是围绕它展开的。
摘要
把 Mixture-of-Experts(MoE)训练做大,会带来 dense 模型上不存在的系统性挑战。因为每个 token 只激活一部分专家,这种稀疏性让总参数量的增长可以远快于每 token 的计算量,从而在内存、通信、计算三个维度上制造出彼此耦合的约束。优化其中一维往往会把压力转移到另一维,这要求在整个系统栈上做协同设计。我们通过一整套贯穿内存(细粒度重计算、offload 等)、通信(优化的 dispatcher、overlap 等)、计算(Grouped GEMM、fusion、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。作为一个高性能、可扩展、生产可用的开源方案,它已被学术界和工业界用于在最多数千卡的集群上训练从十亿到万亿参数规模的 MoE 模型。这份报告解释这些技术如何工作、它们的取舍是什么、以及它们在系统层面如何相互作用,为用 Megatron Core 扩展 MoE 模型提供实践指导。
1 引言
大规模训练 dense transformer,所需计算量随模型规模线性增长23。Mixture of Experts(MoE)模型遵循另一种 scaling 模式:把每个 token 路由到选定的一部分专家网络,而不是激活全部参数,于是每 token 的计算量随模型规模亚线性增长45。近期的 MoE 模型已经证明,相对于质量对等的 dense 模型,训练成本可以下降一个数量级678。
然而,把 MoE 训练做大会带来 dense 模型框架当初没有考虑的系统性挑战。这份报告介绍 Megatron-Core MoE,也就是 Megatron-Core9 之内的 MoE 训练栈,覆盖以高吞吐训练万亿参数级 MoE 模型所需的架构、并行策略和系统优化。
1.1 Mixture of Experts
一个 Mixture of Experts(MoE)模型在标准神经网络之上增加了一组专门化的子网络,称为 专家(expert),再配一个轻量的 router(或称 gating network)来动态选择由哪些专家处理每个输入41011。在基于 Transformer 的语言模型12里,MoE 层通常替换掉 dense 的 Feed-Forward Network(FFN)块:不再是一个 FFN 处理所有 token,而是一个 MoE 层里含多个 FFN 专家,每个 token 根据学到的路由权重被送到其中一小部分(比如 top-$k$)专家。
形式化地讲,给定输入 token 表示 $\mathbf{x}$,router 在 $E$ 个专家上计算一个概率分布:
$$ \mathbf{p}(\mathbf{x}) = \text{Softmax}(\mathbf{W}_r \mathbf{x}) $$MoE 层的输出就是被选中专家输出的加权组合:
$$ \text{MoE}(\mathbf{x}) = \sum_{i \in \text{TopK}(\mathbf{p}(\mathbf{x}))} p_i(\mathbf{x}) \cdot E_i(\mathbf{x}) $$其中 $E_i$ 表示第 $i$ 个专家网络。这套架构提供三个关键优势:模型容量可以通过增加专家数量独立于计算成本增长(可扩展的容量);每个 token 只激活一小部分参数,相比同等规模的 dense 模型减少了 FLOPs(计算效率);不同专家可以针对不同输入类型做专门化(专门化)。附录 A 给出了本报告用到的全部符号与缩写对照。
这个概念本身可以追溯到 1990 年代初1011,但 MoE 与现代 Transformer 的结合才是重新点起兴趣的原因。GShard 首创了大规模分布式 MoE 训练,引入了 Expert Parallelism 和负载均衡辅助损失13。Switch Transformer 证明了 MoE 可以扩展到万亿参数并保持训练稳定5。GLaM 显示 MoE 能以 dense 模型一小部分的训练成本达到相当的质量14。Tutel15、DeepSpeed-MoE16 等框架进一步推进了 MoE 训练系统。
MoE 在研究界和工业界的采用都在加速。Mixtral-8x7B 表明开放权重的 MoE 模型可以匹配闭源 dense 模型,同时降低推理成本6。DeepSeek-V2 和 DeepSeek-V3 用细粒度专家架构把这条路推得更远,用几百个小专家来最大化容量与计算之比17187。NVIDIA 的 Nemotron-3 系列19采用了带 LatentMoE20 的 Mamba-Transformer 混合 MoE 架构,用 Megatron-Core 大规模训练。scaling law 的研究也确认细粒度 MoE 能取得更好的 compute-optimal 取舍2122,这加速了这一趋势——同时也放大了它带来的系统挑战。
1.2 训练大规模 MoE 模型的挑战
随着模型往几百个专家、每个专家容量更小的方向推进,MoE 训练的系统挑战也同比例增长。这些挑战都源于同一个根因:MoE 的稀疏性——它表现为一个 参数-计算失配(总参数量远超实际激活的计算量,本节)和一个 dense-sparse 失配(注意力层和 MoE 层需要互相冲突的并行配置,§3)。
核心的不对称来自稀疏性。在一个 dense transformer 里,每个参数都参与每个训练步。一个有 $N_{\text{total}}$ 参数的模型每 token 大约需要 $6N_{\text{total}}$ FLOPs(前向加反向),所以参数量和每 token 计算量是同步 scaling 的。把模型切到更多 GPU 上,通常也会按同样比例切分计算,让每张 GPU 都足够忙,通信开销的占比就能压得很小。
对 MoE 来说,稀疏性让这个耦合断掉了:每个 token 只激活 $E$ 个专家中的 $K$ 个,所以每 token 计算量大约是 $6N_{\text{active}}$ 而不是 $6N_{\text{total}}$,其中 $N_{\text{active}}$ 随 $K$ 增长而 $N_{\text{total}}$ 随 $E$ 增长,且 $K \ll E$。这造成了一个根本性的参数-计算失配:相比每 token 激活参数相同的 dense 模型,MoE 模型的总参数量要多得多,常常差一个数量级。DeepSeek-V3 就是个具体例子:总参数 685B,但每 token 只激活 37B,差 18 倍。
每 token 计算量这么少,模型切分就得比 dense 模型更小心——粗暴地把专家矩阵分片(Tensor Parallelism 就是这么干的)会把本来就很小的计算切得更碎,效率更低。因为 MoE 专家是彼此独立的网络,自然的策略是 Expert Parallelism(EP):把不同专家放到不同 GPU 上,保住完整尺寸的专家 GEMM。EP 引入了 all-to-all 通信来把 token 路由到它们被分配到的 GPU(§3 详述为什么优先选 EP,以及 Parallel Folding 如何应对随之而来的挑战)。这个参数-计算失配和 EP 的通信需求合在一起,制造出三个紧密耦合的挑战——三堵墙(Three Walls)——它们约束着每一个 MoE 训练步:
内存墙。训练期间,全部 $E$ 个专家的参数、梯度和优化器状态都必须驻留在内存里,尽管每个 token 只激活其中 $K$ 个。这带来的内存压力远超每 token 计算量相当的 dense 模型58。要缓解这个压力,就得在别处花钱:把参数分散到更多设备上要付 通信带宽;重算激活而不是存下来要付 额外计算;offload 到主机内存要付 PCIe 带宽。动态路由让事情更复杂:token 分布不均会在某些专家收到过多负载时造成不可预测的内存尖峰1323。
通信墙。EP 需要 all-to-all 集合通信来把 token 分发给它们被分配的专家、再把结果收回来13。每次 all-to-all 里每 GPU 的发送量大约是 $T \cdot K \cdot h \cdot \frac{EP-1}{EP}$,其中 $T$ 是本地 token 数,$K$ 是 top-$k$,$h$ 是 hidden 维度;一个完整的 dispatch-and-combine 循环会把这个量翻倍。随着 EP 增大,这个量会饱和,但通信会越来越多地从高带宽的节点内链路(比如 NVLink)转移到更窄的跨节点互联上,可用带宽在那里会掉一个数量级24。同时,稀疏激活模式提供的可与这段通信重叠的计算量也有限。在 DeepSeek-V3 这类专家跨多节点的架构里,未经优化的 all-to-all 可能吃掉总训练时间的 60%。
译注:这个发送量公式和 各种并行的通信量分析那篇是同一套算法,可以对着看。
计算效率墙。MoE 引入了 dense 模型里没有的计算低效:
- 小 GEMM。细粒度专家产生大量小矩阵乘,用不满 GPU 的计算单元25。在我们的测量里,GEMM 占 Llama-3 405B(dense)执行时间的约 70%,但在 DeepSeek-V3(MoE)里不到 50%。剩下的时间被那些随 tensor 数量而不是 FLOP 数量增长的操作吃掉了。
- 路由与 permutation 开销。token 路由和 permutation 是 dense 模型里没有的,即便优化过之后仍然给层执行时间增加约 9%。
- 负载不均。动态路由分给各专家的 token 数量不均,有些超载、有些闲着,浪费算力23。
- Host 开销。同样的 FLOPs 量,MoE 因为稀疏性和路由要发起更多 kernel,而每次发起都带固定的 host 侧成本,累积起来会让 GPU 在 kernel 之间空转。在 dropless MoE 里,动态 tensor 形状还要求昂贵的 host-device 同步。
这三堵墙是紧密耦合的:优化其中一堵往往把压力转移到另一堵。加大 batch size 能提升 GEMM 利用率,但会放大内存压力和通信量。CUDA Graphs 能消除 host 开销,但要求静态 tensor 形状,与 dropless 路由冲突。把跨专家的 token 分组能提升计算效率,但让负载均衡更难做。§3 给出 EP 与 Parallel Folding 背后详细的并行分析;§4 给出 Megatron-Core 应对全部三堵墙、同时管理它们之间相互作用的集成方案。
1.3 Megatron-Core MoE
这套 MoE 训练栈构建在 Megatron-Core 之内——一个基于 PyTorch 的大规模 transformer 训练库926——同时应对三堵墙:
多维并行。Expert Parallelism(EP)与 tensor、pipeline、sequence、data 并行集成。MoE Parallel Folding27 解耦了注意力层和 MoE 层的配置,打破了传统的 $\text{EP} \leq \text{DP}$ 约束,使得配置可以针对具体的模型架构和硬件拓扑定制。
内存优化。细粒度激活重计算、memory-efficient permutation、precision-aware 优化器和激活 offload 在不牺牲吞吐的前提下降低内存占用2829。在专家 GEMM、激活和通信上全面支持低精度训练(FP8 和 FP4),进一步降低激活存储,同时靠选择性精度策略维持收敛。
通信优化。高性能 token dispatcher(DeepEP、HybridEP)最大化带宽利用率。通信-计算 overlap 把 all-to-all 延迟藏到专家计算后面。
计算优化。Grouped GEMM kernel、kernel fusion、CUDA Graphs 和 sync-free 执行应对细粒度 MoE 架构固有的计算碎片化。
Production Features。负载均衡策略、带容量控制的 token dropping、distributed optimizer 与 FSDP 支持、支持灵活 resharding 的分布式 checkpoint、以及从 dense checkpoint 做 upcycling,让大规模部署成为可能。这套 MoE 栈的模块化设计支持快速实验,而它生产级的优化支撑从研究原型到万亿参数模型的训练30。这份报告解释这套栈提供了什么、为什么做出那些关键设计决策、以及它们如何应对 MoE 训练的挑战,并给出配置与调优的实践指导。
1.4 本文结构
报告余下部分按从架构到优化再到评测的顺序推进:
- §2:Megatron-Core MoE 架构。分两部分介绍 Megatron-Core MoE 的设计:MoE 层自身的内部设计(router、token dispatcher、experts)和四阶段前向传播(route、dispatch、compute、combine);然后是外部设计,覆盖与 transformer 模型的集成、并行 process group 的组织,以及针对专家参数的优化器处理。
- §3:用 Parallel Folding 和多维并行扩展 MoE。考察 MoE 的稀疏性如何打破 dense 训练的并行假设、为什么需要在传统策略之外引入 Expert Parallelism(EP)、以及注意力层与 MoE 层之间由此产生的 dense-sparse 失配如何被 MoE Parallel Folding 解决——后者解耦了两者的并行配置,从而在万亿参数规模上实现灵活高效的映射。
- §4:打破内存、通信与计算效率墙。给出 Megatron-Core MoE 对三个根本壁垒的解法:内存墙(激活管理、重计算策略、offload、分布式参数存储)、通信墙(包括 DeepEP 和 HybridEP 的优化 dispatcher、通信-计算 overlap)、计算效率墙(Grouped GEMM、kernel fusion、CUDA Graphs、dropless MoE 的 sync-free 执行)。
- §5:MoE 的 FP8/FP4 低精度训练。把低精度训练作为一项横切优化来讲,它同时影响三堵墙——降低激活内存、把通信量减半、加速 Tensor Core GEMM——同时给出维持训练稳定的选择性精度策略。
- §6:长上下文 MoE 训练。考察长上下文场景(16K 到 64K+ token)如何在注意力计算变成主导之后从根本上改变优化格局,并给出通过 Context Parallelism 和 Tensor Parallelism scaling 管理激活内存增长的技术。
- §7:Production Features。描述面向生产训练的运维特性:保证训练稳定的负载均衡与 token dropping、支持与并行策略无关 resharding 的分布式 checkpoint、从 dense checkpoint 做 upcycling,以及与 multi-token prediction 的集成。
- §8:性能评测。通过 DeepSeek-V3 和 Qwen3-235B 在 GB200 与 H100 平台上的实测 benchmark 验证框架的有效性,展示完整优化栈的效果。
- §9:性能最佳实践与 DeepSeek-V3 案例研究。给出一套系统性的工作流来确定最优并行配置,并用一份详细的 DeepSeek-V3 案例研究来验证——展示这些优化如何协同达到 state-of-the-art 的性能。
- §10:RL 场景下的 Megatron-Core MoE。应对新兴的 RL 后训练范式,覆盖 RL 负载独有的挑战(可变序列长度、内存 offload、在线权重导出)、Megatron-Bridge 与主流 RL 框架的集成,以及包括 packed sequence 支持、dynamic context parallelism 和 router replay 在内的 RL 专用优化。
2 Megatron-Core MoE 架构
本节分两部分介绍 Megatron-Core 的 MoE 实现架构。我们先描述 MoE 层自身的内部设计:它的模块化组件(router、token dispatcher、experts)和把输入 token 变成输出表示的四阶段前向传播。然后考察外部设计:并行 process group 如何组织以支持分布式执行,以及优化器如何用不同于 dense 层的方式处理专家参数。
2.1 MoE 层架构与前向传播
一个 MoE 层把 transformer block 里的 dense feed-forward network(FFN)替换成一组专家 FFN,每个 token 只由其中一部分处理。Megatron-Core 通过三个模块化组件(做 token-to-expert 分配的 router、做跨 GPU 通信的 token dispatcher、做计算的 experts)来实现这一点,它们由 Figure 1 所示的四阶段前向传播串起来。这种关注点分离带来了独立优化的可能:router 可以被 fuse 进 CUDA Graphs 而不影响 dispatcher 逻辑;dispatcher 可以在 all-to-all 和 DeepEP 之间切换而不改动专家计算;专家实现可以透明地换用不同的 GEMM 后端。
2.1.1 前向传播:Route、Dispatch、Compute、Combine
MoE 层按四个顺序阶段处理输入 token:
阶段 1:Route。TopKRouter 决定由哪些专家处理每个 token。一个学到的线性投影把每个 token 的 hidden state 映射到 $E$ 个 logits(每个专家一个),一个打分函数(softmax 或 sigmoid)把 logits 转成概率,top-$k$ 选择为每个 token 挑出得分最高的那些专家。router 输出两个 tensor:probs 装路由权重,routing_map 是一个布尔掩码,标记 token-expert 的分配关系。专家数很多时为了数值稳定,router 可以通过 --moe-router-dtype fp32 在 FP32 下运行。
阶段 2:Dispatch。token dispatcher 为跨 GPU 通信做准备。它先 permute token,让所有要去同一个专家的 token 在内存里连续;这个 permutation 对专家侧高效的 dense GEMM 是必需的。然后 dispatcher 用三种后端之一把 token 搬到承载其目标专家的 GPU 上:AllGather(简单但吃内存)、all-to-all(标准的基于 NCCL 的实现)、或 Flex(支持 DeepEP、HybridEP 这类优化后端)。
阶段 3:专家计算。每张 GPU 在收到的 token 上执行它本地的专家。所有本地专家通过 TEGroupedMLP 在一次 Grouped GEMM 调用里跑完,即便单个专家的工作量很小也能把 GPU 利用率拉满。
阶段 4:Combine。反向的通信把处理完的 token 送回原 GPU,随后 unpermute 恢复原始序列顺序。如果配置了 shared expert,它的输出(可选地与 routed expert 并行计算)在这一阶段被加进来。
2.1.2 Router:Token 到专家的分配
router 通过两步操作把一个全局 token batch 变成专家专属的工作负载4:
- Gating:一个线性投影 $\mathbf{W}_r \in \mathbb{R}^{h \times E}$ 把每个 token 的 hidden state $\mathbf{x} \in \mathbb{R}^{h}$ 映射到 logits $\mathbf{l} = \mathbf{W}_r^{\top}\mathbf{x} \in \mathbb{R}^{E}$。
- Top-$k$ 选择:一个打分函数把 logits 转成概率——softmax($p_i = e^{l_i}/\sum_j e^{l_j}$)或 sigmoid($p_i = \sigma(l_i)/\sum_j \sigma(l_j)$,DeepSeek-V3 用的是后者7)。每个 token 选出概率最高的 top-$k$ 个专家。
专家利用率不均会同时损害训练效率和模型质量。如 Figure 2 所示,Megatron-Core 支持多种负载均衡策略,这些在 §7 讨论。
2.1.3 Token Dispatcher:通信抽象
dispatcher 通过一条六阶段流水线管理 GPU 之间的 token 搬运:dispatch_preprocess → token_dispatch → dispatch_postprocess(前向),以及 combine_preprocess → token_combine → combine_postprocess(反向)。
通信后端。有三种 dispatcher 可选:
- AllGather(
allgather):每张 GPU 收齐所有 token,再筛出属于本地专家的那些。简单但吃内存;适合小 EP 规模。 - All-To-All(
all-to-all):标准的基于 NCCL 的点对点通信13。每张 GPU 只发送目标端需要的那些 token。扩展性好,但带同步开销。 - Flex(
flex):统一设计,支持 DeepEP(只集成了高吞吐 kernel,通过 overlap 隐藏 NVLink 通信延迟)24和 HybridEP(高带宽的 MoE 通信 kernel,在 NVL72 这类 NVLink 丰富的拓扑上表现更好)。
2.1.4 Experts:计算模块
每个专家是一个两层 MLP,可选带 gating(用于 SwiGLU/GeGLU 激活31)。Grouped GEMM 让多个专家计算能高效地 batch 起来2515。有两种实现:
- TEGroupedMLP:Transformer Engine32 的优化实现,支持 FP8 和 FP4 量化。
- SequentialMLP:在循环里逐个执行专家。适合 debug,但慢得多。
Shared Experts。一些架构(DeepSeek-V2/V3187、Qwen33)包含一个 shared expert,它处理所有 token 而不管路由结果。shared expert 的计算可以与 dispatch-compute-combine 流水线并行跑,把它的延迟藏起来。细节见 §7。
2.2 系统集成:并行与优化器
描述了 MoE 层内部如何工作之后,我们来看它如何集成进分布式训练系统:并行 process group 和优化器处理。
2.2.1 并行 Group 管理
MoE 层需要与 dense 层不同的 process group,因为不同组件有不同的通信模式。Megatron-Core 通过 ProcessGroupCollection 组织这些 group:
|
|
每个 MoE 组件根据自己的通信需求使用特定的 group:
Table 1:MoE 组件到 process group 的映射。
| 组件 | 用到的 Group | 原因 |
|---|---|---|
| Router | tp, cp, tp_cp |
权重在各 EP rank 上是复制的 |
| Token Dispatcher | ep, tp_ep |
在各专家 rank 间做 all-to-all |
| Experts | ep, expt_tp, expt_dp |
沿 EP 分片;梯度在 EDP 内 reduce |
| Shared Experts | tp |
与 dense MLP 相同 |
这种分离让 Parallel Folding(§3.3)成为可能:注意力层和 MoE 层可以用不同的 TP/DP 配置。比如注意力层可以用 TP=4,而 MoE 层用 ETP=1 配更高的 EP,让两类层各自独立优化。
2.2.2 优化器与梯度处理
专家参数在分布式优化里需要区别对待。Megatron-Core 用一个 ChainedOptimizer 包住 dense 参数和专家参数各自的优化器。三个关键设计决策支撑了正确的 MoE 优化:
- 参数识别。专家参数被标记
allreduce=False,从而与走标准 data-parallel 梯度 reduce 的 dense 参数区分开。 - 独立的 reduce group。dense 层在
dp_cp_group(完整 data parallelism)上 reduce 梯度,而专家在expt_dp_group(expert data parallelism)上 reduce。这保证梯度是在正确的副本数上求平均的。 - 梯度缩放。专家梯度按
edp_size / dp_size缩放,以反映专家(处理的是依赖路由的 token 子集)与 dense 层看到的有效 batch size 不同这一事实。
这套设计让 ZeRO 式的优化器状态分片34能与 MoE 无缝配合:专家参数的优化器状态在 EP group 上分片,而 dense 参数的状态走标准 DP 分片。
架构讲清之后,§3 处理如何把它分布到各设备上,§4 应对规模化之后出现的内存、通信和计算效率挑战。
3 扩展 MoE:Parallel Folding 与多维并行
§1.2 确立了 MoE 的稀疏性把模型规模与每 token 计算量解耦,制造出一个参数-计算失配,使 Expert Parallelism(EP)成为必需但同时引入了 all-to-all 通信。本节考察这个失配具体如何影响并行设计。我们先回顾 dense 模型的并行基线和它带来的取舍,再展示 MoE 的稀疏性如何打破这些策略背后的假设。最后给出 MoE Parallel Folding——Megatron-Core 对由此产生的 dense-sparse 失配 的解法:注意力层和 MoE 层的最优并行配置互相冲突,而 Parallel Folding 解耦它们的映射,让各自都能用自己的最优拓扑。
3.1 为什么需要并行,以及 MoE 为什么不一样
在考察 MoE 专属的挑战之前,我们先确立大模型训练为什么必须并行、以及并行引入了什么取舍。
3.1.1 大模型训练为什么需要并行
内存是根本约束。单张 GPU 的内存是有限的,但训练一个大模型需要同时存下模型参数、优化器状态、梯度和激活35。以用 BF16 精度加 Adam 优化器训练 Llama-405B363738 为例:
| 组成 | 内存(Llama-405B,BF16) |
|---|---|
| 模型参数 | 约 810 GB |
| 优化器状态(Adam) | 约 4860 GB |
| 梯度 | 约 1620 GB |
| 激活(8K 序列) | 约 5575 GB |
| 合计 | 约 12865 GB |
这远超任何单张 GPU 的容量。多卡不是可选项,而是 必须——只有这样才装得下模型。
计算吞吐是效率上的理由。除了内存,单张 GPU 的算力也有限。大模型训练需要天文数字的 FLOPs;聚合多张 GPU 能提升吞吐,把训练的 wall-clock 时间从年缩到周。
3.1.2 并行的取舍
并行不是免费的。每种并行策略都引入开销:
- 通信开销:GPU 之间交换数据要花时间和带宽。
- 同步:快的 GPU 必须在同步点等慢的 GPU。
- 流水线气泡:pipeline parallelism 在流水线边界引入空转时间。
- 计算强度下降:分片让每 GPU 的矩阵变小,降低 GEMM 效率。
结果就是:Model FLOP Utilization(MFU)受并行策略的影响很大,而好的并行设计能把理想 MFU 与实际 MFU 之间的差距压到最小。
dense 模型的关键洞察:并行的「成本」和「收益」是同比例 scaling 的。参数更多需要更多 GPU,但参数更多也意味着每次前反向的计算量更大。因为计算随模型规模增长,通信在每一步里的占比反而更小,模型变大时 MFU 能保持相对稳定。
MoE 的稀疏性破坏了这个平衡。因为每个 token 只激活 $E$ 个专家里的 $K$ 个,总参数随 $E$ 增长而每 token 计算量只随 $K$ 增长。为了内存需要更多 GPU,但每 token 计算量并没有相应增长,通信开销就暴露出来了。下一节考察这种不对称如何打破传统并行策略背后的假设、以及为什么需要一个新的并行维度。
3.2 MoE 并行的挑战
传统并行策略是为 dense Transformer 设计的。MoE 模型的计算模式根本不同,需要一个新的并行维度 Expert Parallelism(EP),并且在把 EP 与已有策略组合时会产生独有的挑战。
3.2.1 MoE 的并行悖论
Figure 3 通过把主流 LLM 的总参数量对每 token 前向 FLOPs 作图来展示这个悖论。dense 模型(圆点)贴着约 2N 参考线,在 scaling 时呈现一个 良性循环:参数增加需要更多 GPU 来装内存,但参数更多也意味着每 token 计算更多。因为计算随模型规模增长,通信在每一步里占的份额更小,MFU 保持稳定。
MoE 模型(三角)打破了这个循环。它们明显落在 2N 线以下,用远少的每 token FLOPs 达到等价能力。看这个根本性的不对称:
| 模型 | 总参数 | 激活参数 | 比例 |
|---|---|---|---|
| Llama-70B(Dense) | 70B | 70B | 1:1 |
| DeepSeek-V3(MoE) | 685B | 37B | 18:1 |
注:DeepSeek-V3 由 671B 的主模型权重加 14B 的 Multi-Token Prediction(MTP)模块权重组成(合计 685B)。表里给的是含 MTP 的总参数量。
DeepSeek-V3 的参数量比它的激活计算所暗示的多 18 倍,这在 Figure 3 里表现为 DeepSeek-V3 的位置与同等参数量 dense 模型本该所在位置之间的落差。这造成一个 复合效应:
- 内存增长快,迫使分散到很多 GPU 上:全部 $E$ 个专家的参数、梯度和优化器状态都必须驻留内存,而更大的 top-$k$ 还会通过 token 复制进一步增加激活内存。
- 通信更多:通过 EP 分散专家引入了 all-to-all 流量,其体量随 $K$ 增长(每个 token 要被分发到跨 EP rank 的 $K$ 个专家)。
- 但计算仍然低:每 token FLOPs 只随 $N_{\text{active}}$($\propto K$)增长,而不是 $N_{\text{total}}$($\propto E$),留下的计算量不足以与不断增长的通信重叠。
后果是:除非并行设计专门针对这种不对称,MoE 训练在根本上是通信受限的。这不是程度问题,而是与 dense 模型的 质的差别。
3.2.2 传统并行策略
dense Transformer 训练一般组合四种并行策略:
Tensor Parallelism(TP) 沿 hidden 维度把权重矩阵切到多张 GPU 上9。每张 GPU 算出一个部分结果,再用 AllGather 或 ReduceScatter 集合通信合并。TP 在矩阵足够大、能摊薄通信开销时表现良好。
Pipeline Parallelism(PP) 按层把模型切到不同 GPU 上3940941424344。micro-batch 流过流水线,各 stage 之间做点对点通信。PP 引入流水线气泡(边界处的空转时间)和 P2P 通信开销,但跨节点扩展性比 TP 好。
Data Parallelism(DP) 把模型复制到各 GPU 上,每张 GPU 处理不同的数据 batch4546。梯度通过 AllReduce 同步。DP 简单,但要求每张 GPU 装下完整模型。
Context Parallelism(CP)474849 沿序列维度把输入序列切到各 GPU 上。每张 GPU 处理序列的一个连续片段,只在注意力计算处需要通信——因为那里 token 必须跨片段边界做 attend。CP 对长上下文训练是必需的,那里激活内存随序列长度二次增长。
传统并行为什么在 MoE 上不奏效:
- 对 MoE 专家用 TP:专家的 hidden 维度很小;上 TP 会切出更小的分片,降低计算效率同时抬高通信在总时间里的占比。高 TP 能高效切分注意力,但会 伤害 MoE。
- 专家很多时用 PP:MoE 模型参数量巨大,只用 PP 的话需要非常多流水线 stage。这会造成过多流水线气泡、降低吞吐。
- 只用 DP:DP 在每张 GPU 上复制完整模型。对万亿参数模型来说这不可能,因为 DP 只能切分数据、不能切分参数。
3.2.3 Expert Parallelism:第五个维度
传统并行切分的是 层(PP)、权重矩阵(TP)、序列(CP)或 数据(DP)。但 MoE 有个独特结构:专家是彼此独立的子网络。这使得第五个并行维度成为可能——切分 专家本身。
Expert Parallelism(EP) 把专家分散到各 GPU 上13。当 EP 度等于 $E$ 时,每张 GPU 持有 $E/\text{EP}$ 个专家。前向传播按以下顺序进行:
- Route:router 为每个 token 选出 top-$k$ 个专家。
- Dispatch:all-to-all 通信把 token 发到持有其目标专家的 GPU。
- Compute:每张 GPU 只用它本地的专家处理 token。
- Combine:all-to-all 通信把结果送回原 GPU。
EP 独有的取舍:
- 通信:all-to-all 集合通信。体量随 token 数增长,与专家数无关。
- 计算:每张 GPU 跑更少的专家,但每个专家处理它 完整 的 hidden 维度。
- 内存:EP 越高 = 每 GPU 专家越少 = 内存压力越低。
EP 提供两个关键收益(Figure 4)。第一,把来自不同 GPU 的 token 汇集到同一个专家上提高了计算强度,改善 GEMM 效率。第二,专家数增加时 all-to-all 的通信量保持恒定,变的只是 GPU 数量。
但光有 EP 不够。EP 只适用于 MoE 层;注意力层没有专家,仍然需要 TP、CP 或其他策略。而且大规模 MoE 训练里两类层都需要 Pipeline Parallelism 来进一步切分模型参数。这种异构性制造出下面要讨论的 dense-sparse 失配。
3.2.4 把 EP 与传统并行组合起来的挑战
单个 Transformer block 里含有两种根本不同的计算模式50,总结在 Table 2 里:
Table 2:单个 Transformer block 内注意力层与 MoE 层并行需求的对比。
| 方面 | 注意力(Dense) | MoE(Sparse) |
|---|---|---|
| 计算 | 每个 token 都 attend 到所有其他 token | 每个 token 路由到 $E$ 个专家中的 $K$ 个 |
| TP | 大 QKV 矩阵受益于高 TP | 每专家维度小,高 TP 反而有害 |
| CP | 长序列受益于高 CP | 没有序列依赖;CP 不相关 |
| EP | 不适用(没有专家) | 分散大量专家的必备手段 |
Dense-Sparse 失配。传统框架强迫 两类层 用 同一套 并行配置,但它们的最优配置直接冲突:高 TP 对注意力有利但会把专家分片打碎;高 CP 对长上下文注意力有帮助但对 MoE 层毫无收益;高 EP 对 MoE 是必需的但对注意力不相关。和参数-计算失配(§1.2)一样,这也源于 MoE 的稀疏性,但表现在并行配置这个层面:它是 dense 层与 sparse 层需求之间的结构性失配,不是一个调参问题。
早前的 MoE 框架把 EP 当作 DP 的一个子维度135:
$$ \text{World Size} = \text{TP} \times \text{CP} \times \text{PP} \times \text{DP}, \quad \text{where } \text{EP} \subseteq \text{DP} $$这个约束之所以存在,是因为框架假定了统一的并行:注意力层用 (TP, CP, PP, DP),MoE 层就直接从 DP group 里切出 EP。这套设计造成三个关键挑战:
挑战 1:GPU 需求是乘法级的。传统框架至少需要 $\text{TP} \times \text{CP} \times \text{PP} \times \text{DP}$ 张 GPU。既然 EP $\subseteq$ DP,要 EP=8 就迫使 DP $\geq$ 8。再叠上长序列要的 CP=8,最少就变成 $1 \times 8 \times 1 \times 8 = 64$ 张 GPU——即便理论上注意力和 MoE 可以共用同样那 8 张卡。这抬高了 MoE 训练的入门门槛。
挑战 2:被迫接受次优并行。因为注意力和 MoE 共享同一个 TP 配置,实践者只能在两种次优配置里选:用高 TP(比如 TP=8)来高效切分大的注意力矩阵,代价是把小专家碎片化成低效的分片;或者用低 TP(比如 TP=1)来保住专家计算效率,代价是注意力层并行不足。两个选项都无法让两类层同时达到最优。
挑战 3:跨节点通信。EP 被约束在 DP 之内时,高 EP 往往迫使 all-to-all 通信跨越节点边界,那里的带宽比 NVLink 低 5–10 倍。同时注意力的 CP 通信也可能跨节点。没有能力把 EP 和 CP 独立映射到高带宽域,通信开销就会主导训练时间。
这些挑战不是彼此独立的调参问题。它们都源于「所有层必须共享一套并行配置」这个根本假设。问题于是变成:怎样在打破这些约束的同时,仍让每类层用上自己的最优并行?
3.3 Megatron-Core 的解法:Parallel Folding 与多维框架
Parallel Folding 是 Megatron-Core 对 dense-sparse 失配的回答。它不强迫注意力层和 MoE 层共享同一套并行配置,而是解耦两者的并行映射,让每类层用自己的最优拓扑27。本节讲关键概念和收益;并行转换处理、token dispatch 设计这类实现细节在配套论文27里描述。
3.3.1 解耦的并行映射
核心想法很简单:不要强迫注意力和 MoE 共享并行。让各自独立用最优配置。
Parallel Folding 为注意力层和 MoE 层引入各自独立的并行 group:
- 注意力层 在 $\text{TP} \times \text{CP} \times \text{DP} \times \text{PP}$ 上组 group,针对序列级的 dense 计算优化。
- MoE 层 在 $\text{ETP} \times \text{EP} \times \text{EDP} \times \text{PP}$ 上组 group,其中 ETP(Expert Tensor Parallelism)和 EDP(Expert Data Parallelism)是 MoE 专用的维度。
唯一的约束:Pipeline Parallelism(PP)在两套布局里必须保持一致,才能保证梯度在模型里正确流动。
3.3.2 完整的多维并行栈
以 Parallel Folding 为基础,Megatron-Core 编排五个并行维度:
| 维度 | 作用于 | 用途 |
|---|---|---|
| TP(Tensor) | 注意力 | 切分大的 QKV/投影矩阵 |
| CP(Context) | 注意力 | 分散长序列 |
| DP(Data) | 注意力 | 处理不同 batch |
| PP(Pipeline) | 两者 | 按层切分模型(必须一致) |
| EP(Expert) | MoE | 把专家分散到各 GPU |
| ETP(Expert Tensor) | MoE | 在专家内部分片(很少用) |
| EDP(Expert Data) | MoE | 复制专家以提升吞吐 |
关键配置原则:
- 注意力层:针对大矩阵(高 TP)和长序列(高 CP)优化。
- MoE 层:针对大量小专家优化(高 EP,一般 ETP=1)。
- PP:两边必须一致,以保证数据流正确。
除上述并行策略之外(Figure 6),Megatron-Core 的 Parallel Folding 框架还集成了带 EP 支持的 Distributed Optimizer 和 FSDP,以进一步降低内存占用:
Distributed Optimizer + EP。每个 rank 上只驻留本地专家的权重和梯度;优化器状态在同一个专家的各副本之间(通过 EDP)分片。这消除了非本地专家的冗余优化器内存,并把梯度同步限制在最小的 group 内。
FSDP + EP。为了更高的内存效率,Megatron-Core 自研的 FSDP(Megatron-FSDP)通过双 DeviceMesh 架构把参数、梯度和优化器状态在 data/expert group 上完全分片,在降低内存占用的同时把 AllGather 和 ReduceScatter 与计算重叠。它兼容 TP/EP/CP 和混合精度(BF16、FP8、FP4)。完整设计见 §4.1.7。
3.3.3 Parallel Folding 的收益
如 Figure 5 所示,MoE Parallel Folding 通过允许 EP 在注意力并行配置的任意子 group 上「折叠」,消除了 EP $\leq$ DP 的限制。这带来四项关键优势:
- 打破 EP $\leq$ DP 约束:EP 现在可以通过在 TP $\times$ CP group 上折叠而超过 DP。考虑注意力配置为 TP=4、CP=2、DP=8、PP=4(共 256 张 GPU):
- 传统方式:EP $\leq$ DP = 8,所以 EP 最大只能到 8。
- 有了 Parallel Folding:MoE 用 ETP=1、EP=64、EDP=1(PP 同为 4)。EP 在 TP $\times$ CP $\times$ DP group 上「折叠」,实现 8 倍高的专家并行度,同时注意力层保持它自己的最优配置 TP=4、CP=2。
- 降低最少 GPU 需求:传统配置下 CP=8、EP=8 至少要 64 张 GPU。有了 Folding,CP 和 EP 共享同一组 GPU,只需要 8 张。
- 支持独立优化:注意力可以为大矩阵用高 TP,而 MoE 用 ETP=1 保住完整专家宽度和更好的 GEMM 效率。
- 把高带宽通信留在 NVLink 域内:CP(给注意力)和 EP(给 MoE)的 all-to-all 通信都可以留在 NVLink 连通的 GPU group 内,避开更慢的跨节点传输。
3.3.4 小结
本节应对的是 MoE 训练的 分布挑战。核心问题是 dense-sparse 失配,它根植于 MoE 的稀疏性:注意力层与 MoE 层的最优并行配置互相冲突,而先前的框架却强迫它们共用一套配置。MoE Parallel Folding 通过解耦并行映射解决这个问题,让注意力用高 TP/CP 而 MoE 独立用高 EP。配合 Distributed Optimizer 和 FSDP,Megatron-Core 把所有并行维度编排进一个统一框架,把 MoE 训练扩展到数千卡。
然而,可扩展的分布本身并不保证训练效率。§1.2 提到的三个根本壁垒——内存墙、通信墙和计算效率墙——仍然要解决。下一节给出 Megatron-Core 突破这些墙的方案。
4 扩展 MoE:打破内存、通信与计算效率墙
§3 确立了 如何 把 MoE 训练分布到数千卡上。本节处理 如何让它高效。挑战不小:§1.2 指出了三个根本壁垒(内存墙、通信墙、计算效率墙),它们源于 MoE 由稀疏性驱动的参数-计算失配。这些墙不只是不便;不解决它们,大规模 MoE 训练要么不可行(内存),要么慢到不可接受(通信、计算)。
为什么是这三堵墙? 这些壁垒源于 GPU 加速训练的基本资源:内存存放参数和激活,互联在设备间搬运数据,计算单元执行运算。任何训练系统都必须拿到足够内存、把数据搬到需要它的地方、并高效地完成计算。其他潜在瓶颈——比如 checkpoint 的存储 I/O 或 host 侧的数据加载——确实存在,但它们工作在允许与训练迭代重叠的时间尺度上。三堵墙代表的是每一次前反向传播上的硬约束。
这些墙相互作用。这些挑战不是彼此独立的;它们构成一个紧密耦合的系统,修好一堵墙可能暴露或恶化另一堵。看一条具体的优化轨迹:某个团队训练 1000B MoE 模型,先撞上内存墙——光激活就超过 GPU 容量。他们启用激活重计算(§4.1),用反向时重新生成激活来换内存。内存约束解决了,profiling 却显示 all-to-all 通信现在成了主导;原来是计算墙挡住了通信墙。他们实现通信-计算 overlap(§4.2),把 all-to-all 传输与专家执行流水起来。但细粒度专家完成得比通信更快,限制了 overlap 的效果。他们启用低精度训练(§5)来降低内存占用,从而允许更大的 batch size,提供更多计算来隐藏通信。这管用,但引入了量化 kernel,它们把执行流打碎、抬高 kernel 发起开销。系统于是变成 host-bound:GPU 在 kernel 发起之间空转。他们部署 CUDA Graphs(§4.3)来减少发起开销,但 graph 要求静态 tensor 形状,与 dropless 路由的动态专家分配冲突。每个解法都会在系统里荡开涟漪,要求同时对三堵墙保持警觉。孤立地处理某一堵墙会得到次优方案;有效的优化要求把内存、通信、计算当作一个统一系统来对待。
Megatron-Core 通过集成的系统设计应对这些挑战。本节余下部分按各方案主要针对的墙来组织:
- §4.1:打破内存墙。怎样把训练塞进 GPU 内存约束内?本节覆盖激活管理、重计算策略、offload、分布式参数存储。
- §4.2:打破通信墙。怎样把设备间通信损失的时间压到最小?本节覆盖优化的 dispatcher、通信-计算 overlap、流水线集成。
- §4.3:打破计算效率墙。怎样让 GPU 一直有活干?本节覆盖 kernel fusion、CUDA Graphs、以及 dropless MoE 的 sync-free 执行。
FP8/FP4 低精度训练同时在三堵墙上带来收益,单独放在 §5 讲。长上下文 MoE 训练会在注意力主导计算时改变优化的平衡,在 §6 处理。
4.1 打破内存墙
内存是 MoE 训练的第一道硬约束:如果参数、优化器状态和激活的合计占用超过 GPU 容量,训练根本跑不起来。随着模型和专家数增长,内存需求增长得很快,让内存优化成为实际训练的必需。要有效优化,先得搞清内存 消耗在哪里。
4.1.1 内存解剖:内存消耗的来源
MoE 模型通常比同等 dense 模型消耗多得多的内存,原因是几处不同的额外开销。考虑用 BF16 精度、$\text{PP}4 \times \text{VPP}4 \times \text{EP}64$ 配置在 256 张 GPU 上训练 DeepSeek-V3。Table 3 给出每 GPU 的内存分解,总需求 199.5 GB,不优化的话远超 H100 GPU 的 80 GB 容量。
Table 3:DeepSeek-V3 在 BF16 训练下的每 GPU 内存分解($\text{PP}4 \times \text{VPP}4 \times \text{EP}64$,256 张 GPU)。
| 组成 | 每 GPU 内存 | 优化技术 |
|---|---|---|
| 权重与梯度 | 36.4 GB | PP、EP 或 TP 分片 |
| 主权重与优化器状态 | 32.1 GB | Distributed optimizer、BF16 moment |
| 激活 | 131.0 GB | 低精度、重计算、offload |
| 合计 | 199.5 GB |
三个组成贡献了这个占用:
权重与梯度(36.4 GB)。全部 $E$ 个专家的参数都必须留在内存里,尽管每个 token 只激活 $K$ 个。DeepSeek-V3 的 256 个专家配 top-8 路由就是例证:685B 参数但每 token 只激活 37B,差 18 倍。
优化器状态(32.1 GB)。Adam 为每个参数存 momentum 和 variance,在 FP32 下把参数内存翻三倍。用 BF16 moment 的混合精度训练能降低这部分开销,但消除不了。
激活(131.0 GB)。最大的内存消费者,超过权重和优化器状态之和。MoE 激活随层深、hidden 维度和 top-$k$ 增长,也随 batch size 和序列长度增长。这让激活成为首要优化目标。
还有一些 MoE 专属因素加剧内存压力:动态路由可能造成负载不均,某些专家临时收到过多 token;token 常常还要 padding 才能塞进高效的计算 kernel,消耗掉超出实际数据的内存。这些都是激活内存成本,可以通过下面讨论的重计算技术降低。负载不均则由 ECHO(§4.3.7)进一步应对,它动态克隆热门专家来平衡各 rank 上的 token 分布。
模型架构不能为了省内存而改,所以优化必须针对数据如何存储和管理。四条互补策略应对内存约束:
- 降低存储精度。更低精度格式(FP8/FP4 而非 BF16)能降低激活内存,对精度影响极小。Memory-Efficient Permutation 则彻底消除冗余的中间 tensor。
- 用计算换存储。激活重计算512852 在前向时丢弃中间结果、在反向时重新生成,用计算周期换内存容量。
- Offload 到主机内存。GPU 内存耗尽时,激活可以在前向时传到 CPU 内存、反向时取回,用 PCIe 带宽换 GPU 内存53。
- 分散到多设备。Fully Sharded Data Parallel(FSDP)54345556 把参数、梯度和优化器状态切到各 data-parallel rank 上,使超出单设备容量的模型也能训练。
下面几小节分两组给出这些技术。先是激活内存优化,按开销从小到大排:零开销(Memory-Efficient Permutation)、精度取舍(FP8/FP4 对 BF16)、计算取舍(重计算)、带宽取舍(offload)。然后是针对权重和优化器状态的优化:precision-aware 优化器和分布策略(FSDP)。
4.1.2 Memory-Efficient Permutation:零开销的激活削减
最理想的内存优化是那些没有计算开销的。Memory-Efficient Permutation 正是如此,它通过一次简单的代数重排消除冗余的中间 tensor。这项技术之所以零开销,是因为它只改变 router 概率 什么时候 被应用,而不改变它 是否 被应用;反向传播需要的 pre-activation buffer 在常规实现里本来也是要存的。
如 §2.1 所述,router 把每个 token 分配给它的 top-$k$ 个专家并附一个学到的路由权重,加权后的专家输出被合并成最终结果。
考虑一个被路由到其 top-$k$ 专家的 token $\mathbf{x}$。设 $\mathcal{T}(\mathbf{x}) \subset \{1,\dots,E\}$ 表示被选中的专家集合,路由权重为 $\{p_i\}_{i \in \mathcal{T}(\mathbf{x})}$。每个专家 $E_i$ 是一个两层 MLP,权重矩阵为 $\mathbf{W}_1^{(i)}, \mathbf{W}_2^{(i)}$,非线性激活为 $\phi$(比如 SwiGLU):
$$ E_i(\mathbf{x}) = \mathbf{W}_2^{(i)}\,\phi\!\bigl(\mathbf{W}_1^{(i)}\mathbf{x}\bigr). $$在标准写法里,路由权重是在专家计算 之后 应用的:
$$ \mathbf{y} = \sum_{i \in \mathcal{T}(\mathbf{x})} p_i \cdot \mathbf{W}_2^{(i)}\,\phi\!\bigl(\mathbf{W}_1^{(i)}\mathbf{x}\bigr). $$Memory-Efficient Permutation 把 $p_i$ 吸收进激活里,在第二个线性层 之前 就应用它:
$$ \mathbf{y} = \sum_{i \in \mathcal{T}(\mathbf{x})} \mathbf{W}_2^{(i)}\,\bigl(p_i \cdot \phi\!\bigl(\mathbf{W}_1^{(i)}\mathbf{x}\bigr)\bigr). $$当专家没有 bias 项时,$\mathbf{W}_2^{(i)}$ 是纯线性映射,标量乘法可交换:对任意向量 $\mathbf{h}$ 都有 $p_i \cdot \mathbf{W}_2^{(i)}\mathbf{h} = \mathbf{W}_2^{(i)}(p_i \cdot \mathbf{h})$,所以上面两式在数学上等价。
这次重排通过消除 router 反向传播所需保存的 tensor 来降低峰值内存。设 $\mathbf{z}_i = \mathbf{W}_1^{(i)}\mathbf{x}$ 表示 pre-activation 输入。在标准写法里,计算 $\partial\mathcal{L}/\partial p_i$ 要求在整个反向传播期间保留每个专家输出 $E_i(\mathbf{x})$。在 memory-efficient 写法里,$p_i$ 直接乘在 $\phi(\mathbf{z}_i)$ 上,所以 $\partial\mathcal{L}/\partial p_i$ 只依赖 $\phi(\mathbf{z}_i)$,而后者可以由一个 fused 反向 kernel 从 $\mathbf{z}_i$ 现场重算出来。既然不管用不用 Memory-Efficient Permutation,$\mathbf{z}_i$ 本来都必须为 SwiGLU 激活自己的反向传播存下来,那就没有引入任何额外 buffer,于是在零计算开销的前提下净降低了峰值内存。Figure 7 展示了这个变换。
对 DeepSeek-V3(Table 3)而言,Memory-Efficient Permutation 每 GPU 省下约 26.3 GB 激活内存,是一笔零计算成本的显著削减。
4.1.3 低精度训练:用 FP8/FP4 削减激活内存
把激活存成 FP8/FP4 而不是 BF16,能在对模型质量影响极小的前提下降低它们的内存占用。
前向传播时,每个线性层的输入必须保留下来,以便在反向时计算权重梯度。在 Transformer 模型里,这些线性层输入构成了激活内存的大部分:注意力投影(Q、K、V、output)和 MLP 层(包括 MoE 里的专家 MLP)各自都要存下自己的输入以供梯度计算。把这些输入 tensor 存成 FP8/FP4 而非 BF16,每个 tensor 的内存占用就降低 50%/75%。
对 Table 3 里的 DeepSeek-V3 配置,启用 FP8 训练能削减约 16 GB 激活内存,占 131 GB 激活预算的大约 12%。这相当于把大约 32 GB 有资格用 FP8 存储的线性层输入减半。其余激活(attention score、normalization 中间量、路由 tensor)要么为了数值稳定需要更高精度,要么本来就不跨前反向边界存储。这项削减与重计算、offload 等其他激活优化是正交的,可以叠加起来累积收益。
低精度训练的机制细节,包括 recipe、量化策略和 MoE 专属优化,见 §5,也可以对着 garden 里 fp16 到 fp8 的混合精度训练那篇看。
4.1.4 重计算:用计算换内存
激活重计算(或称 activation checkpointing)是一项成熟技术,在前向时丢弃中间激活、在反向时重算出来51。但是,粗暴的整层重计算会增加约 33% 的计算开销,而对 MoE 层来说代价更高——因为重算专家计算还会再触发一次 EP all-to-all 通信。Megatron-Core 引入细粒度重计算,只针对那些最吃内存但计算最便宜的操作,以极小开销换取可观的内存节省28。
两项技术组合成这套策略(Figure 8):
细粒度重计算。不是把 activation checkpointing 套在大块的、整体的区段上,而是让用户精确指定哪些计算在反向时重算。比如可以只重算专家 MLP 里的激活函数、LayerNorm 模块,或者 Multi-Latent Attention(MLA)里的 up-projection。通过重算单个算子或子模块,只用适度增加的计算量(额外计算开销不到 5%)就能取得可观的内存节省,因为只有被选中的那部分需要重算。
丢弃输出的重计算。常规的 activation checkpointing 流程会把被 checkpoint 的模块的输出传给下游层,并把这些输出存下来以备反向使用。但既然这些输出在反向时会被重算,这份存储就是冗余的。为了避免这一点,Megatron-Core MoE 在被 checkpoint 的模块的输出被后续层消费掉之后就立即释放它们。反向传播时,这些输出从重算结果里恢复。这个策略确保内存不会被那些可以廉价恢复的激活白占着,在不损害梯度正确性和训练动态的前提下降低内存占用。
Table 4 汇总了 Table 3 那个 DeepSeek-V3 配置下不同重计算目标带来的内存削减。
Table 4:DeepSeek-V3 用细粒度重计算的每 GPU 内存削减($\text{PP}4 \times \text{VPP}4 \times \text{EP}64$,256 张 GPU)。
| 重计算目标 | 每 GPU 省下的内存 |
|---|---|
| MLA Up-Projection | 30.4 GB |
| 激活函数(SwiGLU) | 3.8 GB |
| LayerNorm | 8.2 GB |
| 合计 | 42.4 GB |
4.1.5 细粒度激活 offload
当精度和重计算优化都做完、GPU 内存仍然不够时,把激活 offload 到 CPU 内存能提供额外容量。与用计算周期换内存的重计算不同,offload 换的是 PCIe 带宽。难点在于把传输延迟藏到计算后面,让 offload 看起来「免费」。
动机
细粒度 MoE 模型的参数膨胀极端:DeepSeek-V3 每 token 只激活 37B(18 倍,Table 3);Kimi-K2 总参数达到 1T 而激活 32B(31 倍)。这个参数-计算失配(§1.2)对 offload 尤其关键,因为激活内存不会随 Expert Parallelism(EP)或 Pipeline Parallelism(PP)减少——这些策略降的是参数内存,不是激活内存。
Transformer Engine 提供层级的 offload,但这个粒度太粗,效果受限。同一层内不同模块的内存-计算比差别很大:LayerNorm 的激活小且重算便宜,而 expert_fc1 的输入很大但计算代价高。层级 offload 分不清这两种情况,它只能全 offload(浪费带宽)或全不 offload(浪费内存)。
大致思路:overlap 与 prefetch
GPU 的 Copy Engine 和 Compute Engine 是独立运转的。当一个模块的计算时间超过它的激活传输时间时,D2H 拷贝就能与后续计算并发进行、零成本。Figure 9 展示了前向和反向两种情况下的 stream overlap 机制。
前向传播
前向传播时,输入激活在模块计算完成后立即被 offload 到 CPU,通过一条专用的 D2H stream 与下一个模块的计算并行进行。有一个例外:最后一层的激活不 offload,因为它们在反向时马上就要用,没有计算可以用来隐藏传输延迟。
反向传播
反向传播时,激活的回载遵循 Layer-Staggered Reload 模式:系统在为当前层算梯度的同时,回载 下一层 同一个模块(比如 expert_fc1)的激活。回载发生在每个模块反向完成之后,所以任何时刻 GPU 内存里每种模块类型只驻留一份激活,避免了把激活存储翻倍的需要。当单个模块的激活非常大、2 倍占用会造成意外内存峰值时,这一点是必需的。
PP 与 VPP 的处理
在 PP/VPP 场景下,一个 ChunkOffloadHandler 为每个 (microbatch, VPP stage) 组合管理 offload/reload 逻辑,与 PP=1 的情况一样。主要难点是管理各虚拟流水线 chunk 之间的数据依赖和执行顺序。handler 被入队到一个 deque 里,VPP stage 按逆序(FILO)、microbatch 按正序(FIFO)。反向时弹出的 handler 自动匹配 VPP chunk 的执行顺序,把 offload 出去的激活与对应的反向传播正确配对。
关键技术特性
- 模块级粒度:用户通过
--offload-modules指定 offload 哪些模块,支持混合策略。轻量模块(LayerNorm、激活)用重计算,而昂贵模块(注意力、专家)用 offload。 - 异步传输:专用的 D2H/H2D CUDA stream 与计算并行跑传输;CUDA event 只在必要时协调同步。
- 与重计算集成:可以和细粒度重计算(§4.1.4)组合。对
moe_act,两种策略同时适用:重算激活、同时 offload 它的输入,从而释放整条 fc1 → act 链。 - CUDA Graphs 兼容:用外部 event 而不是 stream 同步,让 offload 模块可以被 CUDA graph 之外的计算隐藏。
- 全场景兼容:支持 PP=1、PP $>$ 1、VPP $>$ 1、所有精度(BF16/FP8/MXFP8/NVFP4)、带 all-to-all overlap 的 1F1B、以及 MoE/MLA 架构。
相对整层重计算的峰值内存优势
整层重计算把每层的输入存在 GPU 上、释放中间激活;反向时从这些存下的输入重算中间量。对一个 $L$ 层模型,GPU 峰值内存是 $L \times \text{layer\_input} + 1 \times \text{layer\_intermediate}$。相比之下,offload 把层输入搬到 CPU;反向时每份输入在用之前才回载、用完立即释放。这把 GPU 峰值内存降到 $1 \times \text{layer\_input} + 1 \times \text{layer\_intermediate}$,与模型深度无关。对深模型(比如 60+ 层),这是整层重计算无法达到的根本性内存优势。
性能
细粒度 offload 和重计算作为互补策略协同工作(Figure 10):LayerNorm 这类轻量操作用重计算,而注意力、专家这类昂贵模块用 offload。异步传输与计算重叠以隐藏 PCIe 延迟。Table 5 给出多种配置下的结果。细粒度 offload 降低 10–18% 内存,只带来 1.6–2% 的吞吐开销。在训练 Qwen3-235B 的情形下,offload 使得降低 Tensor Parallelism 度成为可能,从而在内存成本几乎不变的情况下把吞吐提升 15.0%。
Table 5:细粒度激活 offload 对内存和吞吐的影响。
| 模型与配置 | Baseline | +Offload | 内存 $\Delta$ | 吞吐 $\Delta$ |
|---|---|---|---|---|
| DeepSeek-V3 full | 169 GB | 151 GB | $-$10.7% | $-$1.6% |
| TP1PP8EP32VPP4, MXFP8 | 945 TF/s | 930 TF/s | ||
| Qwen3-235B | 172 GB | 175 GB | $+$1.7% | $+$15.0% |
| TP2 $\rightarrow$ TP1 + EP16 $\rightarrow$ EP64 | 800 TF/s | 920 TF/s |
4.1.6 权重与优化器优化:低精度存储和 offload
前面这些优化针对的是激活内存,它在内存分解里占主导。不过 Table 3 显示主权重和优化器状态每 GPU 占 32.1 GB,是总占用的 16%。对几千亿参数的模型,这部分就成了值得优化的目标。Megatron-Core 提供两项技术:降低存储需求的 precision-aware 优化,以及把不活跃状态搬出 GPU 的 CPU offload。
Precision-Aware Optimizer
Adam 优化器57 为每个参数维护两个状态 tensor:一阶 moment(exp_avg)和二阶 moment(exp_avg_sq)。传统实现把它们存成 FP32,每参数吃 8 字节,对大规模训练是显著的内存瓶颈。关键洞察是:只要实际的更新计算保持在更高精度,优化器状态是能容忍更低存储精度而不影响收敛的。
precision-aware 优化器把存储精度与计算精度解耦58。一阶和二阶 moment 可以存成 BF16(各 2 字节)甚至 FP8(各 1 字节),把每参数存储从 8 字节降到 4 字节(BF16)或 2 字节(FP8)。在每个优化器步里,这些低精度状态在 TransformerEngine 的 FusedAdam kernel 内部被动态 cast 到 FP32;更新用全精度算术计算,以维持数值稳定。
实现提供四个可配置的精度档位:main gradient、main parameter、一阶 moment、二阶 moment。在典型配置里,main parameter 和 gradient 保持 FP32 以确保高质量的梯度更新,而 moment 估计存成 BF1659。这带来优化器状态内存约 50% 的削减(从 Table 3 的 32.1 GB 预算里省下大约 10–12 GB),对训练动态的影响可以忽略。
与 distributed optimizer 组合时——后者把优化器状态切到大小为 $d$ 的 data-parallel rank 上——理论上每 rank 的内存需求会进一步下降。在 DeepSeek-V3 用 BF16 moment 训练时,每 DP rank 每参数的内存消耗从 $6 + 12/d$ 字节降到 $6 + 8/d$ 字节。
状态 offload
前向和反向传播期间,优化器状态占着 GPU 内存但并不活跃;offload 能把这块内存回收给其他操作用。优化器状态 offload 把优化器步保留在 GPU 上,但在 optimizer.step() 之后把优化器状态(exp_avg、exp_avg_sq)和 master weight 传到 CPU,在下一步之前再回载。这个做法既用上了 GPU 算力,又在前向和反向期间回收了内存。
状态 offload 在互联带宽高的系统上尤其有效。在带 NVLink-C2C 的 GB200 上,异步传输能与计算重叠,pinned memory 让带宽利用率拉满。对 DeepSeek-V3,状态 offload 省下 15–20 GB GPU 内存(32.1 GB 优化器与权重预算的 47–62%),每次迭代只增加 0.1–0.2 秒开销。
这两项技术的取舍互补、配合良好。precision-aware 优化器没有性能开销,因为 FP32 cast 发生在 fused Adam kernel 内部;它把优化器状态内存降低最多 50%。CPU offload 省下的内存更多(全部优化器状态和 master weight),但引入适度的传输开销。重要的是这两项技术组合得很好:用 BF16 moment 的低精度存储减小了需要 offload 的状态体量,从而降低传输时间,让 offload 即便在互联带宽不是最高的系统上也变得可行。
4.1.7 面向 MoE 的 FSDP
Fully Sharded Data Parallelism(FSDP)5434 把模型参数、梯度和优化器状态切到各 data-parallel rank 上,让每张 GPU 只持有一份本地分片。MoE 模型里专家参数往往占主导,这让 FSDP 成为 Expert Parallelism 的天然补充。但两者必须正确组合。参见 garden 里的 FSDP。
为什么 MoE 需要 FSDP
Megatron-FSDP 能与多种并行策略无缝组合,包括 Expert Parallelism(EP)、Tensor Parallelism(TP)和 Context Parallelism(CP)。对大规模 MoE 模型,把 FSDP 与 EP 组合起来能把 FSDP 的通用优势转化成针对专家密集负载的具体收益。
- 更低的内存与通信:EP 给每张 GPU 分配一部分专家,而 FSDP 把这些本地专家切到 expert data-parallel(EDP)group 而不是完整的 DP group 上。每 GPU 内存和集合通信量都随 EDP 大小而不是总 DP 大小 scaling,所以在同样的硬件预算下 MoE 模型能支持更多专家或更大 batch。
- 免 PP 的灵活性:FSDP+EP 避开了 pipeline parallelism 的若干工程痛点,包括 PP/VPP stage 划分不均、DeepSeek 式模型里 output 层和 MTP 层的摆放、以及多模态场景下视觉编码器的切分。配置简化成选 EP 大小和 FSDP 分片度,不需要复杂的流水线 stage 设计。
- 广泛的模型兼容性:Megatron-FSDP 支持两条模型实现路径:用 Megatron-Core 自家模块搭的模型,以及用标准
torch.nn模块组成的 PyTorch 原生模型。两条路径都能得到同样的 FSDP+EP 分片、通信优化和 checkpoint 支持。为了与 HuggingFace 生态互操作,Megatron-Bridge 负责 HuggingFace 模型与 Megatron FSDP 之间的在线权重转换,用户可以加载一个预训练 HuggingFace checkpoint、用 FSDP+EP 训练、再把权重导出回去,不需要手动做格式转换。
FSDP+EP:双 DeviceMesh 设计
核心设计难点是 dense 层和专家层需要不同的分片范围。dense 模块(注意力、normalization)受益于在完整 DP group 上做 FSDP 分片。EP 先切分了专家模块,所以每张 GPU 只持有分给它的那些专家。因此 FSDP 应当在每个专家的 data-parallel(EDP)group 内分片,而不是全局分片。
Megatron-FSDP 用一套 双 DeviceMesh 架构解决这个问题。一个主 DeviceMesh 管 dense 模块的 DP-Shard、DP-Outer、TP 和 CP。一个辅助的 Expert DeviceMesh 管 EP 模块,FSDP 的范围限定在 EDP 维度上。每个 transformer 层自动把它的子模块路由到合适的 mesh。注意力和 normalization 用主 mesh,而 MoE 专家 FFN 层用 Expert DeviceMesh。结果是专家参数的 AllGather 和 ReduceScatter 留在小的 EDP group 内,而不是跨越所有 DP rank。
对大规模部署,这套设计扩展到 Hybrid Sharded Data Parallelism(HSDP),它增加一个外层 DP 复制维度。HSDP 在一部分 rank 内完全分片参数、在各子集之间做复制,把 AllGather 限制在组内规模。双 DeviceMesh 为专家参数和非专家参数保持各自独立的 outer-DP group,从而利用 NVLink(节点内)与 scale-out 互联(跨节点)之间的带宽差。
零拷贝通信
双 DeviceMesh 决定了哪些 rank 参与每次 FSDP 集合通信,但集合通信本身仍带有 buffer 管理和数据拷贝的开销。Megatron-FSDP 通过两项优化消除这部分开销。
1. 非均匀分片:标准的 FSDP2 沿主维度独立地切分每个参数,产生逐参数的均匀分片(Figure 11(a))。Megatron-FSDP 则把一个模块内的所有参数压平并拼接,然后在各设备间做非均匀分片(Figure 11(b))。这样分片边界就与通信 buffer 的布局对齐,集合通信可以直接从压平的存储上读,不需要冗余拷贝。在 Llama3 405B 训练里,这把通信开销降低了约 10%。
2. 持久双 buffer 加 NCCL User Buffer Registration:基线 FSDP 频繁地分配和释放通信 buffer,而 NCCL 在每次集合通信时都要在用户 buffer 和它内部的暂存区之间拷数据。Megatron-FSDP 在训练开始时预分配两个持久 buffer 并在它们之间循环(Figure 12),消除了分配抖动。然后 Megatron-FSDP 通过 User Buffer Registration(UBR)把这些 buffer 注册给 NCCL,于是 NCCL 直接读写这块预注册的内存,不经中间拷贝。合起来的效果是真正的零拷贝通信。在 NVLink 系统上,通信 kernel 的 SM 占用从 8–32 个 SM 降到 1–4 个。在支持 SHARP 的 InfiniBand 上,网络交换机接手 reduction,把 GPU SM 完全释放出来。
计算-通信 overlap
即便有了零拷贝集合通信,AllGather 和 ReduceScatter 仍然占着网络,而 GPU 在等参数或梯度同步。Megatron-FSDP 用专用 CUDA stream 把这些集合通信与计算流水起来:下一个 FSDP 单元的 AllGather 在当前单元的前向或反向还在执行时就发起,梯度的 ReduceScatter 与后续层的反向传播并发跑。加大 micro-batch size 能延长可用于隐藏通信的计算窗口,代价是更高的内存占用。用户可以通过 overlap_param_gather 和 overlap_grad_reduce 两个 flag 调这个取舍。
4.1.8 小结
Table 6 汇总本节描述的内存优化技术及它们的主要目标。
Table 6:内存优化技术汇总。
| 技术 | 内存目标 | 取舍 |
|---|---|---|
| 低精度训练 | 激活 | 数值精度与 CPU 开销 |
| Memory-Efficient Permutation | 激活 | - |
| 细粒度重计算 | 激活 | 计算开销 |
| 细粒度 offload | 激活 | CPU 开销与未被重叠的拷贝 |
| Precision-aware 优化器 | 优化器状态 | 数值精度 |
| FSDP(带 EP) | 参数 + 优化器 | 通信开销 |
这些内存优化互相补充:Memory-Efficient Permutation 以零开销消除冗余存储;FP8/FP4 激活以极小的质量影响降低精度;细粒度重计算为特定模块提供有利的计算-内存取舍;激活 offload 在其他技术都不够时提供额外余量;低精度存储和状态 offload 降低权重与优化器内存;FSDP 让规模突破单设备容量。合起来,它们把内存从一道阻塞性壁垒变成一个可管理的约束。
内存优化不是过一次就能忘掉的关卡。对大规模 MoE 模型,内存在整个优化生命周期里都是持续稀缺的资源。除了让训练能跑起来之外,内存余量还解锁其他优化:更大的 batch size 提供更多计算来隐藏通信延迟(§4.2.3),CUDA Graphs 需要额外的静态 buffer(§4.3.6),EP 通信 overlap 必须同时持有多个 microbatch 的激活。后面几节描述的许多优化都要消耗内存,而上面这些技术正是让那些消耗成为可能的东西。
4.2 打破通信墙
通信开销直接降低 GPU 利用率:花在集合通信上的每一微秒都是丢掉的算力。优化之前,EP all-to-all 通信通常吃掉 20–60% 的训练时间,具体取决于模型配置、EP 大小和硬件拓扑。当 EP 留在 NVLink 域内时(比如 GB200 NVL72 上 EP64 的 DeepSeek-V3),开销大约 20%;当 EP 跨越节点边界时(比如 H100 跨节点上 EP64 的 DeepSeek-V3),会升到 40–60%。本节的技术针对这个区间上的所有点。
Expert Parallelism(EP)把专家分散到各设备以把 MoE 模型扩展到超出单设备容量,但这种分散引入了一种独特的通信模式。与通信量与并行度无关的 AllReduce 不同,all-to-all 的量正比于 token 数和 hidden 维度,而更大的 EP 会把这段通信推到带宽受限的跨节点上。对 DeepSeek-V3、Kimi-K2 这类细粒度 MoE 架构,三个因素叠加加剧了这个挑战:
- 频率高:专家更多意味着每层更多 dispatch/combine 操作(DeepSeek-V3 有 58 个 MoE 层,每层需要两次 all-to-all)。
- 跨节点瓶颈:EP 规模大到跨多节点时,跨节点 all-to-all 延迟因带宽较低而占主导。
- 算术强度低:小专家计算完成得快,留给通信重叠的时间更少。
4.2.1 通信解剖:Expert Parallel 的模式
Figure 13 展示了 expert parallelism 的通信模式。在标准 EP 实现里,每个 MoE 层需要两次集合通信:dispatch 把 token 发给它们被分配的专家,combine 把处理完的 token 送回原 rank。对一个有 $L$ 个 MoE 层、每 batch $B$ 个 token、hidden 维度 $h$、EP 度为 EP 的模型,每次前向传播涉及 $2L$ 次 dispatch/combine 操作,每次在各 EP rank 间传输 $O(Bh)$ 的数据。
在 DeepSeek-V3 的规模上,这意味着 58 个 MoE 层 $\times$ 2 次操作/层 = 每次前向 116 次 dispatch/combine 操作。反向传播把这个数字再翻一倍。在 50 GB/s 的跨节点带宽下(比如 InfiniBand NDR),一次 200 MB 载荷的 dispatch 要花几毫秒,每次迭代累积到几百甚至上千毫秒。
通信量由 EP 配置决定,不改并行策略就降不下来。所以优化必须针对通信 如何执行和调度。两条策略协同应对通信开销:
- 把带宽利用率拉满。标准 NCCL all-to-all 实现并没有充分用上可用带宽,对细粒度 MoE 负载尤其如此。优化的 dispatcher(DeepEP、HybridEP)用专门的 kernel 把操作 fuse 起来、利用硬件原语来逼近峰值带宽。
- 把延迟藏到计算后面。即便带宽已经最优,all-to-all 操作还是要花时间。通过把通信与相邻 microbatch 的计算重叠,这段延迟可以被隐藏而不是暴露在关键路径上。
下面几小节按这两条策略组织:带宽优化(DeepEP、HybridEP),以及延迟隐藏(EP 通信 overlap、流水线集成)。
4.2.2 DeepEP 与 HybridEP:把 EP 带宽拉满
在常规 EP 实现里,token 交换依赖 all-to-all 通信。即便用上优化过的 NCCL 集合通信,这条路也有固有局限:dispatch 之前需要一个 permutation 阶段,它把每个 token 复制 top-$k$ 份,制造出冗余流量。在某些设定下,这个预处理还会表现为 host 开销。
为了解决这些问题,Megatron-Core 提供两个基于 token 的 dispatch 后端:HybridEP 和 DeepEP24。基于 token 的 dispatch 消除了 permutation 步骤、避免发送冗余 token,从而降低整体通信量、提升有效带宽。
HybridEP 由 NVIDIA 开发。它遵循与 DeepEP 相同的 token-based 原则,利用 TMA、IBGDA 这类硬件原语,目标是在更低 SM 占用下达到相当或更高的带宽,也包括 Multi-Node NVLink(MNNVL)部署。
对 dispatch,如 Figure 14 所示,HybridEP 按路由信息把数据从 global memory 读进 shared memory,然后通过一个 FIFO 队列把 token 写到目标位置。在跨节点的情形下,它不是直接把重复的载荷经网卡发出去,而是先用一个 RDMA warp group 在各节点上具有相同 local index 的 GPU 之间交换数据,然后在每个节点内转发。这降低了跨节点流量,也让跨节点与节点内的传输可以重叠。
对 combine,标准 all-to-all dispatch 只做通信,所以之后还需要一个独立的 unpermutation 阶段。HybridEP 则把 reduction fuse 进通信 kernel 里。如 Figure 15 所示,HybridEP 经 FIFO 队列读数据、执行 reduction、把结果直接写到目标位置。在跨节点设定下,HybridEP 先对必须跨节点通信的数据做跨节点 reduction,然后在每个节点内完成第二次 reduction。
我们在 GB200 和 H100 上评测 all-to-all dispatch 与 HybridEP,hidden size 7168、序列长度 4096、256 个专家,测多个 EP 规模。结果汇总在 Table 7 里。在所有测试设定下 HybridEP 都一致优于 all-to-all,跨节点场景下增益更大。表里只报了通信延迟;在端到端训练里,一旦把 permutation 和 host 开销算进去,差距通常更大。
Table 7:HybridEP 与 all-to-all 的 EP scaling 性能(单位 µs)。
| EP 大小 | GB200 HybridEP | GB200 all-to-all | H100 HybridEP | H100 all-to-all | |
|---|---|---|---|---|---|
| dispatch | 8 | 391 | 735 | 661 | 1265 |
| dispatch | 16 | 578 | 743 | 1485 | 5774 |
| dispatch | 32 | 612 | 769 | 3064 | 8059 |
| dispatch | 64 | 675 | 930 | 4626 | 9164 |
| combine | 8 | 353 | 741 | 624 | 1277 |
| combine | 16 | 527 | 765 | 1688 | 5628 |
| combine | 32 | 646 | 758 | 3088 | 7815 |
| combine | 64 | 744 | 827 | 4398 | 8727 |
4.2.3 EP 通信 overlap:隐藏 EP 通信延迟
优化的 dispatcher 提升了 all-to-all 的带宽利用率,但根本问题仍在:all-to-all 还是躺在端到端的关键路径上。对 EP64 的 DeepSeek-V3,EP all-to-all 通信仍可占迭代时间的 30–40%,直接限制吞吐。
关键观察是:只要 overlap 窗口里有足够多的独立工作,all-to-all 延迟就能藏到计算后面6061。对 DeepSeek-V3、Kimi-K2 这类细粒度 MoE 模型,这件事很难做,因为跨节点 EP 通信可能吃掉大约一半的层时间,而相邻算子(比如 GEMM)常常太短、藏不住它。
为了应对这种失衡,Megatron-Core 采用一套专门的 1F1B all-to-all overlap 方案62,它把相邻 micro-batch 的前向和反向合并起来,在多条 CUDA stream 上交错编排计算和 all-to-all kernel4039。具体来说,一个 micro-batch 的反向 all-to-all 与另一个 micro-batch 的前向 attention/MLP 重叠。概念上这是一套建在标准 1F1B 之上的、类似 DualPipe63 的双向调度,同时保持了 Megatron-Core 的兼容性。
为了在 1F1B 里启用 all-to-all overlap,我们合并相邻 micro-batch 的前向与反向,并评估了两种模式:
1. 合并 FWD-FWD / BWD-BWD:这个策略合并来自两个 micro-batch 的同类型传播(Figure 16)。在 FWD-FWD 里,micro-batch 0 和 1 的前向并行跑,用计算重叠 all-to-all。BWD-BWD 对反向套用同样的原理。这个做法取舍很清楚:
- 内存开销:峰值激活内存 2 倍。
- all-to-all overlap:能藏住的 all-to-all 更少,因为前向计算大约只有反向计算的一半。
2. 合并 FWD-BWD(等价于 DualPipe):这是首选策略(Figure 17)。它把一个 micro-batch 的前向与另一个的反向合并(比如 micro-batch 1 的 FWD 与 micro-batch 0 的 BWD)。相比 FWD-FWD,它的优势是:
- 内存开销:没有额外内存开销(前向的激活直接被反向复用)。
- 兼容性:与 DualPipe 的设计一致,但避开了复杂的调度。
- 局限:第一个 FWD 和最后一个 BWD 仍在端到端关键路径上,所以它们的 all-to-all 藏不住。
为了在合并 FWD-BWD 里最大化 all-to-all 的隐藏程度,我们用两项关键优化:
1. Stream 分离:我们把负载拆到两条 CUDA stream 上:
- Compute Stream:跑前向/反向计算(比如注意力、专家 MLP)。
- Comm Stream:跑 all-to-all 通信(比如 EP 的 token dispatch/combine)。
通过在两条 stream 之间交替派活,all-to-all 就与计算并行运行——把空转周期压到最小。
2. W/D 拆分(权重梯度 / 数据梯度拆分)7:有一处关键依赖阻碍 overlap:反向 dispatch(B/dispatch)需要反向 MLP(B/mlp)的输出。为了打破这个依赖,我们把反向 MLP 拆成:
- W/mlp:权重梯度计算(与 B/dispatch 无关)。
- D/mlp:数据梯度计算(喂给 B/dispatch)。
这个拆分降低了 compute stream 的空转时间:当 F/mlp 单独太短、藏不住 B/dispatch 时,W/mlp 可以与 F/mlp 重叠来补上。
Figure 18 汇总了这些优化。基线(上)顺序执行,all-to-all 阻塞计算。1F1B FWD-BWD 方案(下)交错相邻 micro-batch,让通信能藏到计算后面。W/D 拆分进一步增加 overlap 机会。合起来,这些方法把 DeepSeek-V3 在 H100 上训练时的 EP 通信开销从(用了 DeepEP 之后的)30–40% 降到迭代时间的 5% 以下。
为了把 overlap 扩展到大模型,合并 FWD-BWD 可以与 Interleaved Pipeline Parallelism(Interleaved PP)组合,后者把模型切成多个虚拟流水线 stage(VPP),在各流水线 rank 上扩大 overlap 机会。Figure 19 展示了一个交错 1F1B 调度的三个阶段:warmup、1F1B、flush。1F1B 阶段里相邻的 FWD-BWD 对可以套用上面的 EP overlap 模式。不过相邻的对如果属于同一个 micro-batch,仍可能有数据依赖。为了避免这一点,我们在 warmup 结束时多跑一个 micro-batch 再进入 1F1B,确保相邻的 FWD/BWD 对之间无依赖,从而在各虚拟 stage 上实现完全的 all-to-all overlap。
合并 FWD-BWD(1F1B)与 DualPipeV(一个改进的 DualPipe 变体)的延迟对比显示:
- VPP 较大时 1F1B 更快(虚拟 stage 更多,overlap 机会更多)。
- PP 规模和 micro-batch 数都很大时两者差距收窄(两种策略都接近最优 overlap)。
- 对混合模型(比如 dense 层与 MoE 层混排的 DeepSeek-V3),各 PP stage 之间的负载平衡至关重要。Megatron-Core 的 Flexible Asymmetric VPP 支持自定义每 stage 的层摆放来最大化 overlap。
有几个因素限制了 all-to-all overlap 能拿到的加速:
- 被重叠 batch 的比例:micro-batch 越多,被重叠的 batch 比例越高。
- all-to-all 的占比:all-to-all 占主导时 EP overlap 的收益更大,跨节点 EP 和计算-通信比低的细粒度 MoE 模型尤其如此。
- SM Carve-out 开销:为 all-to-all 预留 SM 会降低 GEMM 效率。在我们的 DeepSeek-V3 benchmark 里,DeepEP 每 GPU 用 20 个 SM,引入大约 20% 的 GEMM 效率损失。
在 DeepSeek-V3 这类细粒度 MoE 模型上,把 all-to-all overlap 与优化 dispatcher(§4.2.2)组合起来,能对专家通信延迟达到 93% 的 overlap 比例,把专家通信占迭代时间的份额从 30–40% 降到 5% 以下。
4.2.4 小结
通信墙的优化协同工作:DeepEP/HybridEP 把带宽利用率拉满;FWD-BWD overlap 把延迟藏到计算后面;Interleaved PP 把 overlap 机会扩展到各流水线 stage;Flexible VPP 实现最优负载平衡。合起来,它们把 all-to-all 对训练时间的贡献压到 10% 以下。
通信延迟被藏到计算后面之后,本以为 GPU 利用率就高了。然而 profiling 揭示出计算效率墙,它有两个不同的面向:kernel 效率(细粒度专家产生的小 GEMM 用不满 GPU 资源)和 host 开销(大量小操作在 kernel 执行之间制造出 GPU 等 host 派活的空隙)。
4.3 打破计算效率墙
计算效率是最后一道约束:即便内存够、通信也藏住了,如果 kernel 本身低效、或者 host 派活不够快,GPU 资源仍会被浪费。随着模型和专家数增长,计算低效会复合,要达到峰值硬件利用率就必须优化。
4.3.1 计算解剖:低效的来源
MoE 训练里的计算低效来自两个不同的源头:
Kernel 效率。细粒度 MoE 架构产生的负载用不满 GPU 资源。DeepSeek-V3 的 256 个小专家产出的 GEMM,M 维度只有每专家约 128 个 token,远达不到峰值 Tensor Core 利用率所需的数千。除了专家计算之外,MoE 还引入了由许多小 kernel 组成的复杂路由和 dispatch 操作,它们单独都填不满 GPU 算力。
Host 开销。大量小操作制造出 kernel 发起开销,CPU 派活的速度跟不上 GPU 的消耗。三个因素让情况更糟:
- 细粒度专家:许多单独的 GEMM 而不是少数大的,每个都要单独发起一次 kernel。
- 低精度训练:额外的量化 kernel 加重发起开销。
- Dropless 路由:动态 token 数要求 host-device 同步来确定 tensor 形状,把 CPU 摆到关键路径上。
结果就是 host-bound:GPU kernel 之间被几微秒的 host 侧延迟隔开,在 profiling trace 上表现为算力空转的空隙。
kernel 形状由模型架构决定,不改模型就改不了。所以优化必须针对计算 如何组织和调度。五条策略协同应对计算低效:
- 把相关操作 fuse 起来。kernel fusion 把路由、permutation 和辅助损失计算合并成更少、更大的 kernel。
- 用低精度加速。FP8/FP4 低精度训练用上更快的 Tensor Core 操作,在同样的 kernel 形状下提升吞吐。
- 消除每迭代的 CPU 逻辑。CUDA Graphs 捕获 kernel 序列并以最小的 host 介入回放。
- 移除 host-device 同步。sync-free 执行让 GPU kernel 不必等 host 给形状信息就能推进。
- 平衡专家负载。ECHO 动态地把热门专家克隆到利用率低的 rank 上,减少让部分 rank 等待、部分 rank 计算的负载不均。
下面几小节按这些策略组织:面向 kernel 效率的 batching 与 fusion(§4.3.2)、低精度加速(§4.3.5)、消除 host 开销的 CUDA Graphs(§4.3.6)、dropless MoE 的 sync-free 执行以及通过 ECHO 做负载均衡(§4.3.7)。
4.3.2 Grouped GEMM 与 kernel fusion:改善 kernel 效率
提升 kernel 效率最直接的办法就是把小操作合并成大的。Grouped GEMM 把专家计算 batch 起来以获得更好的硬件利用率25,而 kernel fusion 把路由和 permutation 操作合并成更少的 kernel64。Megatron-Core 在三个层面实现这些优化:专家计算(Grouped GEMM)、token 路由(Permutation Fusion)、router 逻辑(Router 与 Aux-Loss Fusion)。
Grouped GEMM
专家的计算本质上就是一串独立的 GEMM。Grouped GEMM 通过重叠各 kernel 的 wave tail 效应,比顺序执行独立 GEMM 性能更好。
Megatron-Core 提供两种 Grouped GEMM 实现:
- i. 多 stream 发起 cuBLASLt GEMM。把单个 cuBLASLt GEMM 发到多条 CUDA stream 上,让它们彼此重叠。这个方法支持多种精度和 scaling 模式,包括 BF16、per-tensor FP8、blockwise FP8、MXFP8 和 NVFP4,性能与高度优化的 Grouped GEMM 实现相当。
- ii. CUTLASS Grouped GEMM。把单个 GEMM fuse 成一个 kernel,在 GEMM 数量很大时性能更好。但不同精度、scaling 模式和硬件平台都需要单独的开发投入,不同 problem size 还需要额外调优。TE 里当前的实现只支持 Hopper 平台上的 BF16。
还有两种实现正在开发中:
- iii. 通过
CUBLASLT_BATCH_MODE_GROUPED的 cuBLASLt Grouped GEMM。这个实现假定形状信息在 device 上,在单个 kernel 里完成计算,从而解锁了对专家部分应用 CUDA Graphs 的可能。它以内置启发式覆盖所有精度和 scaling 模式,目标是在支持的场景下取代 (i) 和 (ii)。 - iv. 带 fusion 的 cuteDSL Grouped GEMM。通过 cuteDSL 把 Grouped GEMM 与激活函数、量化(为下一层准备)和 scaling factor swizzling fuse 在一起。这些 kernel 针对 Blackwell 平台上的 MXFP8 和 NVFP465 优化,目标是 FC1 的 fprop GEMM 和 FC2 的 dgrad GEMM。通过把专家计算、激活和量化合进单个 kernel,这个方案显著减少 kernel 数量,并且与 CUDA Graphs 兼容。
4.3.3 Permutation Fusion
当前 Grouped GEMM 要求分配给同一个专家的 token 连续存放(也就是 permute 过的 token)。如果启用了 all-to-all 集合通信,还需要进一步 permute 以保证发往每个 rank 的 token 聚在一起。用原生 PyTorch 实现高效的 permutation 很困难,因为它往往会发起很多小 kernel、带来额外 CPU 开销。
训练负载里的 permute fusion 流水线(Figure 20)由三个阶段组成:
- Preprocessing:permutation 本质上是一个数据搬运过程,要求 token 在每个专家对应的 buffer 里连续存放。preprocessing 步骤的目的是生成一张 offset map(Figure 20 里的 Row ID map),指出每个 token 在输入和输出 buffer 里的偏移。这保证了 permute 和 unpermute 操作能高效执行。preprocessing kernel 在前向传播里只在 permute 之前调用一次,其他组件复用它的结果。
- Permute:permute kernel 的功能很直接:按 offset map 把 token 从输入 buffer 搬到输出 buffer,然后交给后续的专家 MLP。在 memory-efficient permute 里,概率也需要被 permute,结果直接进到专家的激活部分。
- Unpermute:unpermute 是 permute 的逆操作。因为一个 token 在 permutation 期间会被复制多份、发往不同目的地,这些 token 必须在 combine 阶段被加和到一起。在 memory-efficient permute 里,combine 步骤就是直接相加;否则概率被用作权重。所有累加都在 FP32 精度下进行。
4.3.4 Router 与 Aux-Loss Fusion
router 和辅助损失5661516 的计算是另一处产生大量小 kernel、引入可观 CPU 开销的地方。随着设备算力增长,fusion 变得越来越重要。
难点在于 router 里含有难以 fuse 的组件,比如 GEMM 和通信。我们把余下的操作分解成三段,每段 fuse 成一个 kernel(Figure 21)。
- 打分计算,含 top-$k$ 与 softmax/sigmoid:这是最复杂的函数,针对不同模型有各种分支。它包含不同的 top-$k$ 算法(naive top-$k$、group top-$k$187)、打分函数(sigmoid、softmax)、它们的组合,以及可能的 scaling 操作。
- 辅助损失的打分计算:这是第一段的一个子集。在打分函数和 naive top-$k$ 计算之后,生成辅助损失计算的输入。
- MoE 辅助损失的计算:在第 2 步之上,把辅助损失计算 fuse 成单个 kernel。
启用 router fusion 之后,router 模块内的工作流如 Figure 21 所示。未来的工作可能把通信也纳入 fusion 范围。
4.3.5 低精度训练:低精度加速
带小专家 GEMM 的细粒度 MoE 负载常常是内存带宽受限的,限制了 Tensor Core 利用率。低精度训练通过减少数据搬运、在 Hopper 和 Blackwell GPU 上使用更快的 FP8/FP4 Tensor Core 操作来解决这个问题。
从计算效率角度看,低精度训练带来一个关键收益:FP8/FP4 GEMM 通过用 FP8/FP4 做线性层计算把 Tensor Core 利用率拉满。配合 §4.1.3 的内存收益,FP8 训练在不同硬件平台上为大规模 MoE 训练带来约 10% 到 25% 的端到端性能提升,而 FP4 的加速更多。
低精度训练 recipe、MoE 专属挑战和量化策略的完整内容见 §5。
4.3.6 CUDA Graphs:消除 host 开销
kernel fusion 减少的是 kernel 发起的 数量;CUDA Graphs 消除的是这些发起的 每迭代成本。本小节介绍 CUDA Graphs 如何应对 CPU 开销、drop-and-pad 与 dropless MoE 各自的策略、以及让 CUDA Graphs 变得实用所必需的内存优化。
在 MoE 训练里,CPU 开销常常成为主导的性能瓶颈。这个开销来自三处:
- Python 执行:解释器开销、C foreign function interface(CFFI)和垃圾回收引入延迟。
- 框架开销:即便像
torch.empty这样简单的操作也会增加几微秒的 host 侧成本;库层(TransformerEngine、cuBLAS、cuDNN)再加一层。 - Kernel 发起开销:每次 kernel 发起带几微秒的 CPU 侧成本。
现代训练里有三个趋势放大了这个问题:
- GPU 更快:kernel 执行速度提升,留给重叠 CPU 工作的时间更少,CPU 开销就越来越显眼。
- MoE 的复杂度:MoE 给 FFN 层加了大量复杂度——router、dispatch、Grouped GEMM、combine 操作——除 GEMM 本身之外还产生许多小 kernel,因而带来可观的 CPU 开销。
- 低精度训练:额外的量化 kernel,在算子数量随专家数增长的细粒度 MoE 里尤其明显。
Figure 22 展示了传统执行流水线里巨大的 CPU 开销和由此产生的 GPU 气泡,可以看到 GPU kernel 之间大片的空隙。这清楚说明 CPU 发起 kernel 的速度不够快、喂不满 GPU,这个现象被称为 CPU 开销或 host-boundedness。CUDA Graphs 的应对办法是在初始迭代里把一串 GPU kernel 捕获成一个可回放的 graph6768。后续迭代只发出一次 Graph Launch 调用,基本绕开了逐算子的 Python/框架开销和逐 kernel 的发起开销,从而降低 CPU 侧延迟。
然而 CUDA Graphs 要求所有被捕获的操作有静态、预先确定的形状。动态形状或控制流会让 graph 捕获失败。这与 MoE 形成根本张力,因为每个专家收到的 token 数是动态变化的。下一段描述 Megatron-Core 的两种 CUDA Graph 模式——full CUDA Graphs 和 layer-wise CUDA Graphs——以及它们如何在 MoE 里应对这个静态形状约束。
Full CUDA Graphs 与 Layer-wise CUDA Graphs
如 Figure 23 所示,Megatron-Core 支持两类 CUDA Graphs:full CUDA Graphs 和 layer-wise CUDA Graphs。Full CUDA Graphs(也记作 full-iteration CUDA Graphs)把整个前反向传播捕获成单个 CUDA Graph。这只对 dense 模型、或者用带 padding 的 token dropping 的 MoE 模型支持——也就是每个专家有固定的接收容量,超出容量的 token 被丢弃,输入被 padding 以保证 tensor 形状恒定。所有形状都静态之后,整个迭代都能被 graph 化,把 CPU 开销的削减最大化。至于 DeepSeek-V3 这类 dropless MoE 模型,被路由的 token 一个都不能丢,full CUDA Graphs 就不可行了,因为分配给每个专家的 token 数会随路由决策动态变化。专家 GEMM 的形状每次迭代都在变,违反了静态形状要求。此外,为了 dispatch 还需要 device-host 同步来取每专家的 token 数。§4.3.7 描述我们为克服动态 dispatch 和同步问题所做的工作,让 full CUDA Graphs 也能用在 dropless MoE 上。作为替代,这里我们先介绍一个基于 layer-wise CUDA Graphs 的更简单方法,称为 partial CUDA Graphs。
Layer-wise CUDA Graphs 逐个 transformer 层分别捕获。和 full CUDA Graphs 一样,捕获整个 transformer 层仍会在 MoE 部分撞上动态形状。Partial CUDA Graphs 通过只捕获每层的静态组件、把动态部分留在 graph 之外来规避这一点,即便不捕获完整模型,也能为 dropless MoE 带来很大的性能收益。在一个 dropless MoE 层里:
- 静态组件(可以 graph 化):
- 注意力层(处理固定数量的 token)
- Router 计算(输入/输出形状固定)
- Expert Parallelism(EP)preprocessing(静态的 permutation 元数据)
- Shared expert(如果有,处理所有 token、形状固定)
- Dense transformer block 里的 MLP 层
- 动态组件(不能 graph 化):
- Token dispatch(从固定形状到动态形状)
- 专家 GEMM 操作(M 维度随 token 分配动态变化)
- Token combine(从动态形状回到固定形状)
Partial CUDA Graphs 独立地捕获每层的静态部分,让 routed expert 的计算照常执行。如 Figure 24 所示,每个 transformer 层的前向路径被切成几个「scope」,用户可以选择要 graph 化哪些 scope。启用 "attn+moe_router+moe_preprocess" 这几个 scope 时,一个 graph 就捕获了该层 dispatch 操作之前的所有静态组件:注意力、router(投影、top-$k$ 选择和辅助损失计算)、以及 EP preprocessing(token 元数据计算和 permutation)。如果有 shared expert,它也被捕获在这个 graph 里。这个方案在消除 CPU 开销的同时保住了正确性,因为它允许动态的专家计算以变化的形状执行。尽管没有捕获整个模型,收益依然可观。
Megatron-Core 已经做了大量适配,让 CUDA Graphs 能与 Multi-Token Prediction(MTP)、EP 通信 overlap、细粒度 offload 与 checkpoint、以及灵活 PP 布局这些关键特性正确配合。启用 partial CUDA Graphs 后的 Nsight Systems profiling 显示 CPU 开销大幅下降(Figure 25)。剩余的 CPU 开销主要局限在 token dispatch 与 combine 之间那段区域,它因为动态性而无法被 graph 化。如何解决这个残留瓶颈、让整层都可 graph 化,细节见 §4.3.7。
CUDA Graphs 的内存优化
CUDA Graphs 从三处带来内存开销:
- Graph 结构:存 graph 拓扑的内存(通常可忽略)。
- 独立的内存池:被 graph 化和未被 graph 化的操作需要各自独立的 PyTorch 内存池,不能共享内存。用 partial CUDA Graphs 时,单个 transformer 层需要两个池:一个给 graph 化的工作、一个给未 graph 化的。
- 静态 buffer:每个 graph 需要专用的输入/输出静态 buffer,一旦分配就从 PyTorch 的内存分配器里被拿走了。
如果不仔细优化,partial CUDA Graphs 会从上述三处招来高得多的内存开销。Megatron-Core 和 Transformer Engine 实现了几项优化把这个开销压到最小:
减少 graph 数量。用 pipeline parallelism(PP)时,每个 microbatch 需要一个专用的 CUDA Graph。原因是:如果各 microbatch 共享一个 graph,microbatch $i+1$ 的前向会在 microbatch $i$ 的反向执行之前覆盖掉保存的上下文,造成内存损坏(Figure 26)。这导致总共需要 $L \times M \times 2$ 个 graph(其中 $L$ = 每 GPU 层数,$M$ = microbatch 数,$\times 2$ 是前向/反向)。
反过来,如果不用 PP,各 microbatch 可以共享同一个 graph,数量降到 $L \times 2$。为了实现 graph 共享,引入了一个 is_first_microbatch GPU flag 来控制共享 graph 内 microbatch 专属的行为,比如只在第一个 microbatch 上跑的量化。回放之前把这个 flag 设成 0 或 1,就能让被捕获的量化 kernel 按需被跳过或执行。
内存池共享。虽然 graph 化和非 graph 化的操作需要各自独立的内存池,但只要按执行顺序捕获,所有 graph 可以共用一个池。Transformer Engine 的 make_graphed_callables() API 接受一个 _order 参数来指定 microbatch 间的调度,从而允许顺序捕获并正确共享内存池。
Buffer 复用。静态输入/输出 buffer 可以按 PP 执行顺序在各 graph 之间复用。一旦某个 microbatch 完成了某层的反向传播,它的前向 buffer 和反向输入 buffer 就可以给接下来的前向和反向 graph 复用。只有反向输出 buffer 必须保留到数据被发往下一个 PP stage 为止。
有了这些内存优化,我们在 DeepSeek-V3 的 GB200 训练上以约 7 GB 额外内存换来 10% 的端到端加速。
残留问题:动态的专家计算
Full CUDA Graphs 不支持 dropless MoE。Partial CUDA Graphs 消除了静态组件的 CPU 开销,但动态的专家计算区域仍在 graph 之外。有没有办法让 CUDA Graphs 覆盖动态部分?要做到这一点至少得解决两个挑战:
- 动态形状需要 host–device 同步:host 必须从 GPU 查询每专家的 token 数,才能发起 token dispatch 和 Grouped GEMM。
- 内存分配需要已知的 buffer 大小:按最坏情况分配 buffer 浪费内存;按实际大小分配又需要同步。
下一小节给出 Megatron-Core 如何通过三项互补技术为 dropless MoE 实现完整的 CUDA Graphs 覆盖:sync-free kernel、ECHO 和 Paged Stashing。
4.3.7 通过 sync-free kernel、ECHO 与 Paged Stashing 为 dropless MoE 实现完整 CUDA Graphs 覆盖
kernel fusion 和 CUDA Graphs 有助于降低 host 开销,但 dropless MoE 引入了一个新挑战:动态 tensor 形状需要 host-device 同步,阻止 CUDA Graphs 覆盖专家计算。
在 dropless MoE 里,每个 EP rank 收到的 token 数随路由决策动态变化。router 在 GPU 上生成这个形状信息,但 host 需要它来发起后续操作和分配内存。要拿到这个信息就得做一次 device-to-host 拷贝和同步,把 CPU 摆到关键路径上、也阻止了对动态部分的 CUDA Graphs 捕获。要为带动态路由的 dropless 训练实现 sync-free MoE,有两个根本挑战:
挑战 1:不知道实际 problem size 就要发起 kernel。按惯例,很多 GPU 算子假定(动态的)形状信息在 host 上可用。host 侧代码根据运行时拿到的形状信息决定发起配置(比如 grid size 和 tile size)以及 kernel 必须完成的工作量。在 dropless MoE 里这就制造了一个 host-device 同步点:host 必须在发起 Grouped GEMM 或通信 kernel 之前从 device 查询每专家的 token 数。要解决这个问题,处理动态形状的 GPU kernel 需要重新设计。我们开发了 device-initiated GPU kernel,包括 device-initiated Grouped GEMM 和用 HybridEP 的 sync-free dispatch。
挑战 2:不知道实际大小就要分配内存。sync-free 执行意味着 buffer 的大小需要在 host 上预先确定。为了避免 token 溢出,buffer 常常需要开得过大,比如按最坏情况——所有 token 都被路由到同一个专家——来定尺寸。但这种超额供给会导致严重的内存碎片:如果实际的 token 分布是均衡的,绝大部分预分配内存就一直闲着,却又不能被 graph 内的其他操作回收。最坏情况的 buffer 可能需要比实际工作集多 $O(\text{EP\_size})$ 倍的内存。
这种内存碎片可以通过两条互补策略缓解:降低负载不均,让最坏情况的供给更接近实际使用;以及在 CUDA Graph 内做 动态内存管理 来回收未用的 buffer 空间。ECHO 走第一条路,动态地把热门专家克隆到利用率低的 rank 上,降低各 rank token 数的方差、从而缩小最坏情况与平均内存需求之间的差距。Paged Stashing 走第二条,在 CUDA Graph 内实现细粒度内存管理,让预分配 buffer 里未用的部分可以被其他操作挪用。ECHO 和 Paged Stashing 详述如下。
Device-Initiated GPU Kernel
要消除常规 host-initiated GPU 算子所需的强制 host-device 同步,kernel 必须是 device-initiated 的。这要求三件事:
- i. GPU kernel 需要从 GPU 内存读形状信息并用它来指导计算。
- ii. GPU kernel 需要把它实际执行的工作量(只在运行时才知道)与它的静态发起配置解耦。
- iii. GPU kernel 需要跳过不必要的计算,比如对 padding 数据的操作。
要满足这些要求,现有 kernel 必须重新设计。Grouped GEMM kernel 和 HybridEP kernel 是下面两个具体例子。
Device-Initiated Grouped GEMM。如 §4.3.2 所讨论,TE 提供两种 Grouped GEMM 实现:多 stream 的 cuBLASLt GEMM 和 CUTLASS Grouped GEMM。两者都是 host-initiated 的。它们需要一份 CPU 侧的每专家 token 数列表来确定每个 GEMM 的形状和发起配置,从而形成一道 host-device 同步屏障。为了消除它,TE 引入了 device-initiated Grouped GEMM,直接从 device 内存读形状信息。提供两种实现:
- i. cuBLASLt Grouped GEMM。从 CUDA 13.1 起,cuBLASLt Grouped GEMM 支持以 device 数组形式传入矩阵形状。这个实现内置了选择最优 kernel 配置的启发式,随着近几个 CUDA 版本已覆盖多种精度和 scaling 模式。
- ii. 带 fused 激活与量化的 cuteDSL Grouped GEMM。在基于 cuteDSL 的实现里,SwiGLU 激活可以被 fuse 进 epilogue 阶段,而对 FP8 训练量化也可以被 fuse 进来。当前支持仅限于特定精度和 fusion 模式,但持续的开发正在扩大对新配置的覆盖。
用 HybridEP 的 sync-free dispatch。§4.2.2 介绍的 HybridEP 提供了 MoE 通信的高效实现,它在 sync-free MoE 里也扮演重要角色。每次 dispatch 和 permutation 之后,给定 rank 收到的 token 数是动态的,正常情况下需要同步才能拿到 buffer 大小。通过估计一个上界并把它传给 HybridEP,dispatcher 可以按这个界预分配输出 buffer,以额外 GPU 内存为代价消除所有同步。
即便有了 device-initiated kernel,最坏情况的 buffer 分配仍然浪费内存。两条互补策略应对这一点:ECHO 降低负载不均、让最坏情况供给逼近实际使用;Paged Stashing 实现细粒度内存管理以回收未用的 buffer 空间。
ECHO:热门专家的弹性克隆
实现细节见 Megatron-LM PR#236869。
负载不均是 MoE 固有的:热门专家收到的 token 远多于其他专家,这造成两个问题。第一,承载热门专家的 EP rank 成为计算瓶颈,导致其他 rank 在同步点等待。第二,高的负载方差意味着按最坏情况供给 buffer 相比实际使用浪费大量内存。ECHO 通过把热门专家动态克隆到利用率低的 rank 上的空闲 slot,同时解决这两个问题。
Figure 27 展示了 ECHO 的工作流。前向传播时,ECHO planner 识别出热门专家并生成两个输出:一张 hot expert map,指明把哪些专家克隆到哪些空闲 slot;以及一张更新后的 routing map,把溢出的 token 重定向到这些克隆体。Expert Dispatch 通过基于 HybridEP 的 sync-free 通信把热门专家的权重拷到空闲 slot;Token Dispatch 把 token 路由到 home 专家和克隆专家;专家计算在所有专家上进行。反向传播时,Expert Gradient Dispatch 从克隆体收集梯度并 reduce 回 home 专家,保证一致性。克隆专家在计算完成后被丢弃以节省内存。
为训练克隆专家比为推理克隆贵得多:每个克隆专家在前向时需要一次权重通信、反向时需要一次梯度 reduction 以保证一致性。把所有热门专家克隆到所有空闲 slot 会让负载均衡最优,但引入过多通信开销。ECHO planner 的目标是在达到足够负载均衡的前提下最小化克隆数量。它为每个专家计算 spillover(超出每 EP rank 平均负载的那部分 token),以及每个 rank 上的 spare capacity(低于平均的可用容量)。planner 用一个 bin-packing 算法高效地把溢出 token 匹配到空闲容量上,确定具体克隆哪些专家、克隆到哪些 rank,以最小的克隆开销产出 hot expert map 和 routing map。
ECHO 提供两个关键收益:(1) 降低 CUDA Graphs 的内存碎片,因为降低各 EP rank 之间的负载方差让最坏情况的 buffer 定尺寸更接近实际使用,从而允许更小的静态 buffer;(2) 提升计算效率,因为均衡的负载减少了 straggler 问题,改善整体 GPU 利用率。
Paged Stashing
实现细节见 Megatron-LM PR#269070。
ECHO 降低负载不均以缩小最坏情况与典型内存需求之间的差距,而 Paged Stashing 解决的是互补的问题:在 CUDA Graph 内实现细粒度内存管理以减少内部内存碎片。
碎片问题的根源是 CUDA Graphs 涉及的 buffer 大小必须预先确定。为了避免不期望的溢出,这些 buffer 通常需要开得过大。在直接的做法里,所有层的 buffer 都按最坏情况尺寸分配(Figure 28 中间)。对 dropless MoE 来说,这意味着每层都按每 rank 可能收到的最大 token 数预留一块 buffer,即便实际 token 数低得多,也要付 $O(\text{layers} \times \text{worst\_case})$ 的内存开销。
Paged Stashing 基于一个关键观察:存放反向计算所需激活的实际内存,与为避免溢出而在每层分配的最坏情况内存之间,存在巨大差距(常常超过一个数量级)。Paged Stashing 通过解耦这两块内存来应对,如 Figure 28(右)所示。它不为每层都分配最坏情况尺寸的 buffer,而是由 paged-stashing manager 维护:(1) 一块按最坏情况 token 数定尺寸的 tmp buffer,它在所有层之间共享用于计算;以及 (2) 一块以 page 形式组织的 stashing buffer,只存每层实际用到的 token。当一层完成前向计算后,它的激活被从 tmp buffer 拷到 stashing buffer,只存实际 token 数而非最坏情况的分配量。tmp buffer 于是可以立即被后续层的计算复用。这把内存消耗从 $O(\text{layers} \times \text{worst\_case})$ 显著降到 $O(\text{worst\_case} + \text{actual\_total})$。
stashing buffer 用了经典的分页思想来做灵活而高效的内存管理。PagedStashBuffer 按固定大小的 page 组织(默认每 page 64 个 token)。page stashing manager 通过一个用循环 buffer 实现的 free list 追踪可用 page。分页操作以 device-initiated 的 stash kernel 实现。stash kernel 把激活拷到从 free list 头部分配的空闲 page 上,并记录目标 page。reload kernel 从记录的 page 取回激活,并把用完的 page 归还到 free list 尾部。
为了把开销压到最小,Paged Stashing 用专用 CUDA stream 把内存传输与计算重叠(Figure 29)。系统在当前计算执行的同时预取下一个反向层的激活,隐藏 reload 延迟。这个 overlap 因为双缓冲会带来轻微的内存开销。
合起来看:完整的 CUDA Graphs 覆盖
这三项技术合起来,为 dropless MoE 实现了完整的 CUDA Graphs 覆盖,在消除 CPU 开销的同时保住了动态路由的灵活性:
- Device-initiated kernel 通过直接从 GPU 内存读动态形状,彻底移除 host-device 同步,让所有 MoE 操作都能被 CUDA Graphs 捕获。
- ECHO 在各 EP rank 之间平衡专家负载,缩小最坏情况与实际 buffer 需求之间的差距。这允许更小的静态分配,也缓解了本来会限制 GPU 利用率的 straggler 问题。
- Paged Stashing 通过跨层共享一块最坏情况 tmp buffer、只在 paged stashing buffer 里存实际激活,把内存从 $O(\text{layers} \times \text{worst\_case})$ 降到 $O(\text{worst\_case} + \text{actual\_total})$。
开销。这些收益伴随着实测的取舍。ECHO 在前向时为克隆热门专家、反向时为 reduce 它们的梯度引入了额外通信;实践中只有一小部分专家被克隆(由超出平均负载的 spillover 决定),所以相对标准 all-to-all dispatch,额外通信保持在适度水平。Paged Stashing 增加了内存拷贝操作(前向 stash、反向 reload)以及一块持久的最坏情况尺寸 tmp buffer 加双缓冲开销;这些拷贝跑在专用 CUDA stream 上、与层计算重叠,所以延迟基本被隐藏,主要成本是额外的内存带宽。相比消除逐迭代 host-device 同步、以及让 CUDA Graphs 下的静态内存规划成为可能这两项收益,两种开销都相对小。
4.3.8 小结
计算效率优化同时应对两个低效来源。对 kernel 效率:Grouped GEMM 把专家计算 batch 起来,kernel fusion 合并路由和 permutation 操作,低精度训练加速 Tensor Core 吞吐。对 host 开销:CUDA Graphs 消除每迭代的 CPU 逻辑,Sync-Free MoE 移除 host-device 同步屏障。ECHO 应对同时加剧这两个问题的负载不均。合起来,这些优化把计算低效从主导瓶颈变成一个可管理的约束。
4.4 小结:打破三堵墙
本节应对了高效 MoE 训练的三个根本壁垒,每一个都是不解决就无法在规模上实际训练的硬约束。
- 内存墙:内存是第一道约束。如果参数、优化器状态和激活超过 GPU 容量,训练就跑不起来。Memory-Efficient Permutation、FP8/FP4 激活、细粒度重计算、offload、precision-aware 优化器和 FSDP 合起来,把 DeepSeek-V3 每 GPU 199.5 GB 的需求变成一个可行的训练配置。
- 通信墙:通信开销直接降低 GPU 利用率。优化的 dispatcher(DeepEP、HybridEP)把带宽拉满,FWD-BWD overlap 把延迟藏到计算后面。合起来,它们把 DeepSeek-V3 训练里 all-to-all 的开销从最多 60% 的训练时间降到 10% 以下。
- 计算效率墙:计算低效来自两处:用不满 GPU 资源的小 kernel,以及让 GPU 空转的 host 开销。Grouped GEMM 和 kernel fusion 改善 kernel 效率;CUDA Graphs 和 Sync-Free MoE 消除 host 开销。这些把计算低效从主导瓶颈变成可管理的约束。
这些优化不是彼此独立的。它们构成一个连贯的系统,对一堵墙的解法会使能或增强对其他墙的解法。低精度训练把这一点体现得很清楚:它既降低内存(激活)又提升计算效率(更快的 Tensor Core 操作),同时需要仔细集成以避免引入新瓶颈(量化 kernel 开销)。
接下来两节处理两个横切议题:§5 覆盖同时影响三堵墙的 FP8/FP4 低精度训练,§6 处理注意力主导计算时改变优化平衡的长上下文 MoE 训练。
5 MoE 的 FP8/FP4 低精度训练
前面几节用针对性的优化各自应对了三堵墙。低精度训练很特殊:它同时影响内存、通信和计算效率,是一个横切的优化,值得单独一节来讲。
混合精度训练一直是高效深度学习的基石,BF16 是标准做法717273。FP8/FP4 低精度训练是更激进的一步,把精度从 16 位降到 8 位甚至 4 位,收益和风险也相应更大7475。对 DeepSeek-V3,FP8 训练把激活内存每 GPU 降低约 16 GB,通过更快的 Tensor Core 操作改善专家 GEMM,并加速参数 AllGather 通信。这些收益在三堵墙上复合,让 FP8 成为高效大规模 MoE 训练的必需品。
不过低精度训练引入的收敛风险必须被系统性地处理。本节介绍 Megatron-Core 的低精度训练方案:理解精度在哪里重要、选择合适的量化 recipe、以及实现 MoE 专属的优化来在维持训练稳定性的同时拿到低精度的收益。
5.1 为什么低精度训练对 MoE 尤其重要
与 dense 模型相比,MoE 架构把低精度训练的收益和风险都放大了:
收益被放大。有几百个专家,激活内存按比例增长,FP8/FP4 激活带来的绝对内存节省比 dense 模型更大。参数 AllGather 的通信量减半76,而主导 MoE 计算的专家 GEMM 也受益于更快的 Tensor Core 吞吐。
风险被放大。router 决策依赖精确的打分来把 token 分配给专家。量化噪声可能让专家选择不稳定,导致训练不稳定、模型质量下降、或者专家坍塌77。数值敏感的组件必须被保护起来,不做激进量化。
策略:选择性精度。解法是「该精确的地方精确,其他地方追求效率」。三条原则指导 MoE 里的低精度训练部署:
- 保护路由决策。router 保持在 FP32,以确保专家选择稳定。
- 为关键组件保留精度。embedding、输出层、main gradient、master weight 和优化器状态保持原精度,以维持模型质量。
- 量化大头计算。构成计算主体的专家 GEMM 用低精度训练,配合精心设计的量化方案(称为 recipe)。
5.2 低精度训练对三堵墙的影响
低精度训练在三堵墙上都带来收益,这让它成为一个统一性的优化。Table 8 汇总了这些影响。
Table 8:低精度训练对三堵墙影响的汇总。
| 墙 | 低精度训练的收益 | 详见 |
|---|---|---|
| 内存 | 激活减少 50%(FP8)/ 75%(FP4) | §4.1.3 |
| 内存 | 消除 BF16 权重副本 | §5.3.5 |
| 内存 | BF16 优化器状态 | §4.1.6 |
| 通信 | 参数 AllGather 减少 50%76 | §5.3.5 |
| 计算 | 更快的 Tensor Core GEMM | §4.3.5 |
| 计算 | 量化 kernel 开销(负面) | §4.3.2 |
5.2.1 用低精度训练打破内存墙
低精度训练从三处降低内存消耗:
激活内存(详见 §4.1.3)。激活内存节省的主要来源是为反向计算保存的线性层输入 tensor。把这些 tensor 存成 FP8/FP4 而非 BF16,内存占用降低 50%/75%。比如 FP8 为 DeepSeek-V3 每 GPU 节省约 16 GB 激活内存。
参数内存。常规的低精度训练维护三份参数副本:FP32 master weight、BF16 模型权重、FP8/FP4 计算权重。我们的 native FP8/FP4 方案(详见下面 §5.3.5)通过把权重直接从 FP32 cast 到 FP8/FP4、并用低精度做参数 AllGather,消除了 BF16 那一份。
优化器状态内存(详见 §4.1.6)。Adam 优化器状态(一阶和二阶动量)可以存成 BF16 而非 FP32,把优化器内存减半。这一点与低精度训练正交,BF16 训练也适用。
5.2.2 用低精度训练打破通信墙
FP8/FP4 的参数 AllGather。当 FP8/FP4 primary weight 与 distributed optimizer 配合使用时,参数 AllGather 通信减少 50%(每参数 1 字节而非 2 字节)。注意:
- 对 NVFP4,我们需要 gather 行方向和列方向两份 FP4 权重,所以参数 AllGather 减少的流量也是 50%。
- 对 MXFP8,即便用了 FP8 primary weight,我们仍然先把 FP32 master weight 拷进一块临时 BF16 buffer、用 BF16 精度通信。原因是 MXFP8 在前向和反向需要不同的量化方向,通信 MXFP8 权重就得同时传行方向和列方向两个量化版本,实际上每参数还是两字节,相比 BF16 通信没有任何优势。所以对 MXFP8,我们用 BF16 精度通信参数。
5.2.3 用低精度训练打破计算效率墙
更快的 GEMM。FP8(Ada/Hopper 及之后)和 FP4(Blackwell 及之后)Tensor Core 的吞吐高于 BF16,前向和反向都能加速。
取舍:量化开销。低精度训练引入了额外的量化 kernel、增加 CPU 开销,对已经在给 CPU 施压的细粒度 MoE 尤其棘手。这个开销通过以下手段管理:
- kernel fusion,把量化与其他操作(比如 normalization 和激活函数)fuse 起来。
- 把 padding/unpadding 与 permutation kernel 或 routing map fuse 起来,避免显式的 padding/unpadding。
- grouped quantization,把多个输入 tensor 的量化 fuse 进单个 kernel。
- CUDA Graphs 捕获量化 kernel 以消除发起开销,详见 §4.3.6。
5.3 低精度训练 recipe:Per-Tensor FP8、Blockwise FP8、MXFP8 与 NVFP4
确立了低精度训练对三堵墙的影响之后,我们来看定义一个低精度训练配置的技术组件。
一个低精度训练 recipe 由这几部分构成:
-
数据格式。FP8 有两种格式:E4M3 和 E5M27478。训练里通常用两种组合:
- E4M3:输入、权重和梯度全部量化成 E4M3 格式。
- Hybrid:输入和权重量化成 E4M3,梯度量化成 E5M2。
而 FP4 只支持 E2M1 一种格式。
-
Scaling 粒度。可选 per-tensor 和 per-block(也叫 group 或 tile)。对 per-block scaling,block 大小在不同 recipe 之间也不一样。
-
哪些部分被量化。通常所有线性层都被量化,而 embedding、输出层(LM head)和优化器保持原来的高精度。部分通信可以用 FP8/FP4(比如 TP AllGather、DP AllGather、EP all-to-all)。当前的低精度训练 recipe 通常把 SDPA 保持在 BF16。
-
额外机制,比如 stochastic rounding、Random Hadamard Transform(RHT)等。这些机制是 recipe 专属的,需要仔细验证。
Megatron-Core 提供三种 FP8 recipe,每种都经过大规模验证(Figure 30)。从 per-tensor 到 blockwise 再到 MXFP8 的演进,既反映了硬件进步,也反映了大规模部署里学到的经验。Per-Tensor Scaling 每个 tensor 用一个 scale factor,简单但精度有限,适合做实验。Blockwise FP8(128$\times$128 block)提供更细的粒度、在规模上有验证过的稳定性,Hopper 上推荐用它。MXFP8(1$\times$32 粒度)有 Blackwell Tensor Core 的原生支持和硬件加速的 scaling,Blackwell 上推荐用它。此外 Megatron-Core 还在 Blackwell 上提供一个 NVFP4 recipe79。
5.3.1 Per-Tensor FP8 recipe
per-tensor FP8 recipe 在 Hopper 和 Blackwell 上都支持,通常采用 hybrid 格式(输入/权重用 E4M3,梯度用 E5M2)。它计算输入 tensor 里所有值的绝对最大值(amax),用它把 tensor 量化成 FP8。
per-tensor scaling 有两个变体:
- Delayed scaling:用一个历史窗口里的 amax。这打破了 amax 计算与 scaling 之间的数据依赖,性能最好,但因为 amax 是估计出来的,有丢精度的风险。
- Current(live)scaling:即时计算 amax,精度更好。
因为精度问题,我们不推荐 delayed scaling;current scaling 提供更好的精度和模型收敛性。Figure 31(a) 和 31(b) 分别展示了 Hopper 和 Blackwell 平台上带 per-tensor current scaling 的线性层计算。由于 Hopper 上只支持 TN layout 的 FP8 GEMM,我们必须为反向计算存下转置过的 FP8 激活。在 Transformer Engine 里,量化 kernel 把 cast 和 transpose fuse 成单个 kernel 以减少 global memory 访问。在 Blackwell 平台上,所有 layout 的 FP8 GEMM 都支持,所以我们不需要转置版本,进一步降低内存占用。
per-tensor scaling 因为简单、精度和性能都还不错,是迁移到 FP8 训练的一个好起点。
5.3.2 Hopper 上的 Blockwise FP8
blockwise FP8 recipe 对所有 tensor 采用 E4M3 格式,把激活和梯度按 $1\times 128$ tile 量化、权重按 $128\times 128$ block 量化。blockwise scaling 用更细的 scaling 粒度换取更好的精度,并且已在超大规模 MoE 模型的生产中被证明成功,包括 DeepSeek-V37、 MiniMax-M280、Ant Ling-2.081 等。因此 blockwise FP8 是 Hopper 平台上推荐的 FP8 recipe。
Transformer Engine 为 blockwise recipe 提供高度优化的量化 kernel 和 GEMM kernel(通过 cuBLAS),以及层级 API 和一个 autocast context,让 FP8 训练只需几行代码就能启用。
Figure 31(c) 展示了 Hopper 平台上带 blockwise FP8 recipe 的线性层计算,除了 tile-based 量化之外,它与 Hopper 上的 per-tensor current scaling 非常相似。
5.3.3 Blackwell 上的 MXFP8
在 Blackwell 平台上,得益于第五代 Tensor Core 对 MXFP8 格式的原生支持82,我们采用 MXFP8——一种更细粒度的训练量化方案。激活和权重都按 $1\times 32$ 粒度量化,scaling factor 用 E8M0。理论上 MXFP8 因为 scaling 粒度更细而应该更精确,又因为 Tensor Core 原生支持 MXFP8 而性能更好。所以 MXFP8 是 Blackwell 平台上默认的 FP8 recipe。
Figure 31(d) 展示了 Blackwell 平台上带 MXFP8 recipe 的线性层计算。虽然 Blackwell 上所有 layout 的 FP8 GEMM 都支持,但因为要在另一个方向上量化,我们仍需为激活和权重额外存一份列方向量化的 FP8 数据。
MoE 训练里针对 MXFP8 recipe 正在进行的优化包括:
- grouped quantization,降低量化开销。
- 把激活和量化(为下一个 GEMM 准备)与 GEMM kernel fuse 起来。
- 确保 MXFP8 量化、scaling factor swizzling、Grouped GEMM 的整条流水线都是可 CUDA-graph 化的。
5.3.4 Blackwell 上的 NVFP4
在 79 里我们提出 NVFP4,一种用于 LLM 训练的 4 位 microscaling 格式,它在保留 FP4 效率收益的同时改善数值保真度。
NVFP4 用 E2M1 格式的 FP4 元素,因为裸 FP4 的可表示范围非常有限,所以值需要带 scaling 地量化。一个核心设计选择是 两级 microscaling。NVFP4 施加一个 per-tensor 的 FP32 scale 和一个 per-block 的 8 位 E4M3 scale。tensor 级的 FP32 scale 先把 tensor 分布重映射到与 block scaling 兼容的范围,然后 block 级的 E4M3 scale 把每个 block 映射进 FP4 范围。NVFP4 用 16 个连续元素为一个 block,每个 block 的 amax 被 scale 到 FP4 的最大值。
在这个格式设计之上,我们还发现稳定的 NVFP4 训练依赖于围绕量化的若干算法选择。具体来说我们加了三项实用技术:
- Random Hadamard Transform(RHT):应用于权重梯度计算,降低 outlier 的影响。
- 2D scaling:具体是 16$\times$16 的权重 block scaling(同时保留 FP32 tensor scale),让权重量化在前向和反向之间更一致,减少前反向的量化不匹配。
- Stochastic rounding:用在梯度上,降低 FP4 转换时的 rounding bias,因为确定性 rounding 引入的偏差会损害收敛。
这些补充是该论文所用训练 recipe 的关键部分,对更大规模上的收敛很重要。
5.3.5 FP8/FP4 primary weight:消除冗余存储
常规的低精度训练维护一个三层参数体系:给优化器更新用的 FP32 master weight、作为中间表示的 BF16 模型权重、以及给前反向 GEMM 计算用的 FP8/FP4 权重。同时维护 BF16 和 FP8/FP4 两份参数副本引入了冗余的内存开销。
Native FP8/FP4 通过建立一条从 FP32 master weight 到 FP8/FP4 计算权重的直接 cast 路径来消除这份冗余,完全绕过 BF16 中间层,既降低内存占用又加速参数 AllGather。
实现 native FP8/FP4 的核心难点在于管理 FP8/FP4 tensor 所需的量化元数据。我们提供一个统一接口支持不同的 FP8/FP4 recipe(delayed scaling、current scaling、blockwise scaling、MXFP8 和 NVFP4),它处理了 TransformerEngine 各实现之间的版本差异,保证跨版本兼容。
在 distributed optimizer 的每个 optimizer step 里,FP32 参数分片根据 reduce 后的梯度被更新。随后这些更新过的 FP32 分片用合适的 scaling factor 直接量化成 FP8/FP4 格式。量化过程从 FP32 值计算 amax,跨各 data-parallel rank 更新 scaling factor 以保持一致性,然后把量化后的 FP8/FP4 值写进模型参数 buffer。这条直接路径消除了维护 BF16 参数的内存开销,对参数内存可能超过几百 GB 的大语言模型影响尤其显著。
更具体地,对 delayed scaling 和 current scaling,FP32 分片的量化分三步进行(Figure 33):
- 第 1 步:从 master weight 拿到 local abs-max。如果某个权重在当前 rank 上没有分片,就把它对应的 abs-max 设为 0。
- 第 2 步:通过 AllReduce 拿到 global abs-max。
- 第 3 步:用 global abs-max 和 master weight 做 partial cast。
在 Blockwise recipe 里,一个权重的 abs-max 是在 2D block 上计算的,所以我们不能直接用一个 reduction kernel 来拿 local abs-max。我们实现了专门的 kernel,它们知道权重的 2D layout、也知道 master weight 与原始权重之间的对应关系,据此计算 abs-max(Figure 32)。
FP8/FP4 primary weight 的另一个好处是参数 AllGather 的通信量大幅减少(见 §5.2.2)。
5.4 低精度训练里 MoE 专属的挑战与解法
前面几小节覆盖的是适用于任何深度学习模型的低精度训练基础。MoE 架构给低精度训练引入了独特挑战,需要超出标准实现的专门解法。
5.4.1 Padding 与 Unpadding 的 fusion:动态形状对齐
FP8 和 FP4 GEMM 要求 tensor 维度对齐到特定倍数:per-tensor 和 blockwise recipe 是 16,MXFP8 和 NVFP4 是 32。这些 padding 要求既来自每种量化 recipe 的 block scaling 粒度,也来自 Tensor Memory Accelerator(TMA)访问必须在 GEMM 点积维度上保持 16 字节对齐这一约束。前向传播里,hidden 维度 K 是点积维度,通常在模型设计上就已经对齐了。但在权重梯度 GEMM 里,点积维度对应 token 维度 M,它动态变化、经常不满足对齐要求。因此沿 token 维度做零填充是必要的。为了进一步启用 CPU 开销更低、且与 CUDA Graph 兼容的 grouped quantization kernel,我们把每专家 token 数的 padding 提高到 128。§5.4.3 解释为什么 NVFP4 grouped quantization kernel 需要这个更大的 padding;同样的原理也适用于 MXFP8 grouped quantization。总之,最终施加在 token 维度上的 padding 由 routed expert 层里所有相关 kernel 的要求共同决定。
我们的基线方案是为不同专家的多个输入 tensor 做 fused padding 和 unpadding kernel。但显式的 padding 和 unpadding kernel 引入了不可忽略的开销。我们提出两个解法:
- Routing map padding,padding 的是 routing map 而不是收到的 token。这保证了每专家的 token 数对齐到要求,代价只是多发送少量额外 token,避免了昂贵的 per-tensor padding 操作。
- 把 padding fuse 进 permutation,避免一次 global memory 读写,性能应该是最好的。可用时它是默认选择。
5.4.2 Grouped Quantization 与 Grouped GEMM
Grouped Quantization。由于量化数据与 scaling factor 之间的 padding 要求不同,加上行方向和列方向量化都引入了复杂的数据 layout,量化不同专家的输入 tensor 最朴素的做法就是逐个对它们套用量化 kernel。然而这会产生大量极小的 kernel,既加重 CPU 开销、从 GPU 角度看也不高效。
为了解决这个问题,我们实现了 grouped quantization kernel,把多个 tensor 的量化 fuse 进单个 kernel,大幅降低 CPU 开销、提升 GPU 利用率,并且可以被 CUDA Graph 捕获。
Grouped GEMM。优化过的、可 CUDA-graph 化的 grouped GEMM 见 §4.3.7。
5.4.3 NVFP4 量化 fusion
NVFP4 量化格外复杂,因为我们用了一些数值技巧来保持数值稳定性,所以必须做激进的 kernel fusion 才能把量化开销控制住。实践中我们的 NVFP4 量化 kernel 不只是一个简单的「scale + cast」kernel:它需要仔细地把训练 recipe 的逻辑吸收进来,包括 Random Hadamard Transform、2D scaling 和 stochastic rounding。
其中 RHT fusion 对延迟最敏感。如果实现成一个独立 kernel,Hadamard 变换会在 global memory 里引入一次额外的全精度(BF16)tensor 读写,显著增加带宽成本。通过把 RHT 与 NVFP4 量化 fuse 起来,我们在单个 kernel 里完成 Hadamard 变换和 FP4 量化,避开了那份额外的 BF16 流量。
第二个实现挑战是 Blackwell NVFP4 Tensor Core GEMM 是 TN 取向的,而 Wgrad 用的是转置过的激活和梯度。所以当 RHT 被 fuse 进 Wgrad 路径的 NVFP4 量化时,kernel 还必须把 transpose 也吸收进来,而不是依赖一个独立的 BF16 transpose kernel(那又会招来一次经 global memory 的全 BF16 tensor 读写)。
结果就是这条 fused 量化流水线必须支持从同一个高精度输入产出多个输出:(1) 给前向 GEMM 用的标准 FP4 量化;(2) 给反向/Wgrad 路径用的 transpose + RHT + FP4 量化。训练时我们在前向传播里发起这个 fused kernel,生成两份 FP4 副本:一份立刻被前向 GEMM 消费,一份存下来给反向。原始的高精度输入随后被丢弃,所以我们既避免了存 BF16 激活,又高效地为反向路径做好了准备。
还有一处复杂来自 NVFP4 里的 per-tensor FP32 scale。当前的 NVFP4 预训练 recipe 里,tensor 级 abs-max 是在 Hadamard 变换 之后 测的,这意味着我们需要一个专门的 Hadamard+amax kernel,它只计算 amax(不 materialize 一个变换后的 BF16 输出 buffer)。结果 Hadamard 变换实际上被算了两次——先在 Hadamard-amax kernel 里,再在 fused 量化 kernel 里。虽然这重复了变换的计算,但端到端仍然比把一整份变换后的 BF16 tensor 写到 global memory 更快,因为它避开了额外的高带宽 BF16 读写流量。
相比之下,2D 量化在端到端吞吐上不太显眼,因为它只作用于权重:我们可以只量化一次权重(比如在第一个 microbatch 上),把量化后的权重和 scale 缓存起来给后续 microbatch 复用,所以它的成本基本被摊掉了。
不过 stochastic rounding 在关键路径上加了一个 kernel 级要求:量化 kernel 还必须支持一条可选的 stochastic rounding 路径,随机数在 kernel 内生成(对 dY 这类梯度 tensor 启用,其他情况关闭),这样 FP4 量化才能保持在单趟里被 fuse 住。为了保证效率,我们用 NVIDIA cuRANDDx83 在 fused CUDA 量化 kernel 内部做 device 侧随机数生成。
MoE 的 Grouped NVFP4 量化。为 MoE 层支持 NVFP4 量化格外困难,因为 NVFP4 recipe 本身的算法复杂度就高。在 full-iteration CUDA Graph 的要求下,MoE 量化路径也必须是 CUDA Graph 安全的:host 不能依赖 CPU 上的动态专家 token 数,只能依赖一个 device 侧的 tokens-per-expert tensor。这意味着输入激活不能再被单独切分和量化,我们必须为整条 NVFP4 流水线实现 grouped quantization kernel。NVFP4 输出的内存必须在不知道中间形状的情况下作为一整块 flat buffer 分配。
一个关键观察是:dense NVFP4 激活量化 kernel 的大部分都可以复用到 grouped quantization 上,因为经过 routing/permutation 之后,不同专家的 token 在内存里是连续 packed 的。主要差别在于为保住性能和 graph 安全性所需的实现约束:
- 一个量化 thread block 不应跨越来自多个专家的 token,因为那会引入显著的控制流和索引开销;所以我们把每个专家的 token 数在 token 维度上零填充到 thread block 形状的整数倍(通常是 128)。
- NVFP4 的 transpose 必须按专家进行(独立地转置每个专家 packed 的激活,然后拼接),这与把整个 grouped buffer 当成一个 tensor 来转置并不等价。
- NVFP4 GEMM 需要 scale-factor swizzling:scaling factor 必须被 pad 到 128$\times$4 对齐的形状,并在 GEMM 之前 swizzle 成 32$\times$16 layout,如 cuBLAS 文档84 所述。这个 scaling factor 形状约束是施加在 per-expert GEMM 上的,也就意味着确定所需的 padding 大小隐式地要求 CPU 能看到 tokens-per-expert,那不是 CUDA Graph 安全的。为了避开这一点,我们从构造上就强制每专家 128 token 对齐。
第 1 点和第 3 点都要求每专家 token 数是 128 的倍数,这需要一个专门的零填充 kernel 来满足。为了消除这种零填充开销,我们把 per-expert 零填充能力 fuse 进 token-permute kernel,让 routed expert 的激活一产出就已经是对齐且零填充好的 layout,下游的 grouped NVFP4 量化/GEMM kernel 于是可以直接假定所需对齐、无需额外预处理。
对 NVFP4 recipe 里的 per-tensor FP32 二级 scale,MoE 场景下的「per-tensor 二级 scale」指的是 per-expert 二级 scale,而不是让所有专家共享同一个 amax,这样才能在训练中保持数值稳定性。这些 amax 值无法提前知道:它们必须从被路由的 token 在线计算,每次迭代都生成各专家不同的 amax 值。为了高效做到这一点,我们利用每专家 128 对齐这个保证,把 dense 输入的实现改造成一个 CUDA-Graph 安全的 grouped Hadamard-amax kernel。
6 长上下文 MoE 训练
前面几节应对三堵墙时,都假定 MoE 层主导计算成本。这个假定对序列长度 4K–8K 的典型训练任务成立,那时专家计算和 all-to-all 通信是主要瓶颈。然而推理模型(OpenAI o1、DeepSeek-R185 等)和基于 RL 的训练的兴起创造了新的边界:16K、64K,甚至更长的序列。在这些长度下,MoE 训练的性能特征发生了根本转变:注意力主导计算,三堵墙的相对重要性也随之改变。本节考察长上下文训练如何为 MoE 模型重塑三堵墙,以及 Megatron-Core 如何应对这些新挑战。
6.1 当注意力开始主导:计算重心的转移
在 §4 里我们确立了 MoE 训练面对内存、通信、计算效率三堵基本的墙。那里给出的优化(激活管理、dispatcher 优化、Grouped GEMM、CUDA Graphs)针对的是 MoE 层主导计算画像的场景。长上下文训练从根本上改变了这个假定。关键洞察是 MoE 层与注意力层之间计算复杂度的差异,如 Figure 34 所示:
- MLP 组件:计算成本随序列长度线性增长,复杂度 $O(s)$。
- 注意力组件:scaled dot-product attention(SDPA)主导计算成本,随序列长度二次增长,复杂度 $O(s^2)$。这种二次增长催生了大量关于高效注意力机制的研究8687。
比如在 64K token 下,SDPA 吃掉 69% 的 FLOPs,而短序列场景下只有 10–15%。注意力在计算上占据主导。所幸的是 SDPA 在大多数 GPU 加速库里都被高度优化过,比如 FlashAttention888990 和 cuDNN(Table 9),所以 SDPA 不会成为性能瓶颈。优化重心于是转向内存和通信:只要我们在不引入过多开销的前提下解决这两个挑战,训练性能就能保住。
Table 9:DeepSeek-V3 在 cuDNN 里的 SDPA 性能。
| 平台 | 序列长度 | 前向 TFLOPS | 反向 TFLOPS |
|---|---|---|---|
| Hopper | 4,096 | 553 | 422 |
| Hopper | 16,384 | 638 | 523 |
| Blackwell | 4,096 | 1,324 | 1,083 |
| Blackwell | 16,384 | 1,698 | 1,298 |
6.2 管理激活内存的增长
激活内存随序列长度增长是长上下文训练的首要挑战。为了应对这道被加剧的内存墙,我们套用一组协同工作的技术,它们来自 DeepSeek-V37 和 Qwen391 等大规模 MoE 负载上的经验。这些技术是 §4.1 内存优化原则的延伸。
Context Parallelism 与 Tensor Parallelism。把 CP 和 TP 组合起来能把激活内存分散到各设备上,从而支持本来会超出 GPU 容量的序列长度。关键原则是保持 sub-sequence 长度(每个 CP/TP 分片的序列长度)大致恒定,通常是 4096 或 8192。让 $\text{CP}\times\text{TP}$ 随序列长度一起放大,就能让每设备内存维持在基线水平附近。这让负载看起来很像短上下文训练,除了 SDPA 和 TP/CP 专属的通信之外,大部分短上下文优化仍然适用。
优化器 CPU offload。大模型里优化器状态常常每 GPU 吃掉几十 GB。CPU offload 几乎能把这块内存全部回收,代价是传输开销和 host 侧优化器开销。所以这是内存余量与吞吐之间的取舍。对 256 张 H100 上的 DeepSeek-V3,在约 50% MFU、序列长度至少 16K 的条件下,最坏情况开销约 2%。虽然确切取舍取决于模型规模、序列长度和集群规模,但对长上下文训练来说,优化器 CPU offload 通常是值得的。
选择性重计算。重计算用算力换内存:前向丢掉中间激活、反向再生成。关键是基于内存-算力取舍做模块级的选择。在长上下文设定下 SDPA 主导总计算,所以重计算 SDPA(Megatron-Core 里的 core attention recomputation)通常太贵。相比之下,重计算 MLP 相关模块这类成本更低的组件往往净收益更好。对 64K 序列长度的 DeepSeek-V3,SDPA 贡献最多 72% 的总计算;重算它增加约 18% 计算开销、造成约 16% 性能损失,却只省下 9 GB 内存。改为重计算非 SDPA 组件,全局能省 89.8 GB,性能影响相当甚至更小。所以我们建议关掉 core attention recomputation,优先重计算其他模块。
实践中 CP 和 TP 是长上下文内存效率的主要手段。优化器 CPU offload 和选择性重计算再提供额外余量,在内存受限的平台上尤其有用。这些技术彼此配合良好,可以按需组合。
6.3 Context Parallelism 与 Tensor Parallelism 的取舍
因为 CP 和 TP 是长上下文训练里最有效的两个策略,实际问题就变成:当注意力主导计算时,该怎么把它们组合起来。
CP 和 TP 都能降低长上下文训练的激活内存,但它们的通信模式和内存效果不同(Figure 35)。CP 沿序列维度切分激活,把激活内存降低 CP 倍。对 SDPA,CP 有两种通信模式:point-to-point(P2P)和 all-to-all。P2P CP 在 CP 组内以 ring 方式交换 KV cache,这个交换与 SDPA 计算重叠489249。all-to-all CP 在 SDPA 之前把 tensor 从 sequence-sharded 变成 head-sharded,之后再恢复成 sequence-sharded。TP 额外切分线性层权重,降低参数内存但在线性层里引入额外集合通信。在 all-to-all CP 和 TP 两种情况下,SDPA 都跑在 head-sharded 的 tensor 上,这增加了每分片的 sub-sequence 长度、可以改善 SDPA kernel 效率。
实践中 P2P CP 是更常见的 CP 模式。相对 TP,它的主要优势是通信天然与计算重叠,而 TP 通信往往更难藏住。不过 TP 能降低 dense 参数内存、还能改善 SDPA 效率。所以在节点内通常首选 TP,那里通信快、内存收益也大。跨节点则常常首选 P2P CP,因为 TP 通信开销会增长而 P2P 的 overlap 依然有效。all-to-all CP 处在两者之间:它避开了 P2P CP 的多步 ring 交换,也避开了 TP 引入的线性层权重切分开销。实践中,当 ring 式交换不合适而 TP 通信成本又高时,all-to-all CP 常与 TP 组合用于节点内执行。
Megatron-Core 支持把两种 CP 后端组合成 hierarchical CP,实现与 TP 的拓扑感知配对。一个实用的起点是:节点内用 all-to-all CP 配 TP,以改善 SDPA kernel 性能并降低参数内存;跨节点用 P2P CP,以保住通信-计算的重叠。这些是指导原则而非硬性规则,但作为 TP/CP 调优的初始配置很好用。
6.4 变长训练的 Packed Sequence
前面几小节考察了序列长度增加时三堵墙如何转移。这些分析假定每个 batch 内序列是定长的。但实际训练场景——尤其是强化学习和监督微调——常常涉及变长序列,这带来额外的效率挑战。
传统 batching 要求把所有序列 pad 到该 batch 内的最大长度,序列长度差异很大时会造成大量浪费93。比如一个 batch 里有 4K、8K、32K token 三条序列,所有序列都得 pad 到 32K,约 60% 的计算和内存浪费在 padding token 上。这种低效与上面讨论的内存和计算挑战复合在一起。
为此 Megatron-LM 引入 packed sequence,在单个 batch 内处理变长序列而不浪费在 padding 上。此外,为了应对变长序列造成的 DP 不均衡和 CP 低效,我们引入 Dynamic Context Parallelism(Dynamic-CP),它与 sequence packing 方案联合地、按 microbatch 自适应选择有效 CP 大小。
6.4.1 Packed Sequence 支持
Megatron-LM 支持 packed sequence,让多条变长序列可以被拼接、在单个 batch 内处理,序列之间无需 padding。这是通过 THD(Total tokens $\times$ Heads $\times$ Dimension)tensor 格式实现的,与常规的 SBHD(Sequence $\times$ Batch $\times$ Heads $\times$ Dimension)格式相对。THD 把注意力 tensor 表示成 [total_tokens, num_heads, head_dim],不同样本的序列沿 token 维度拼接。
核心机制依赖累积序列长度追踪,它标出每条独立序列在 packed tensor 里的起止位置。这些参数被沿注意力机制传递给 Transformer Engine 的 fused kernel,从而实现尊重序列边界的高效 SDPA 和 RoPE 操作。注意力计算用累积序列长度来保证来自不同序列的 query 和 key 不会互相 attend,在消除 padding 开销的同时维持正确性。
Megatron-LM 的实现支持多种训练场景下的 packed sequence:
- 强化学习:bin-packing 算法把变长轨迹分组进固定大小的 bin,最大化 GPU 利用率。
- 多模态训练:图像 token 数可变的视觉-语言模型用 packed sequence 来处理来自不同数量视觉 patch 与文本 token 组合的异构序列长度。
- 与 Context Parallelism 集成:当 Context Parallelism 与 padding 要求同时启用时,系统自动切到 THD 格式,同时提供 padded 和 unpadded 两套累积长度,以在保住计算效率的同时处理通信对齐。
packed sequence 支持的收益在长上下文场景下格外明显。对 thinking 模型的 RL 训练,序列长度可以从几百变化到几万,相比传统的基于 padding 的做法,packed sequence 可以降低 40–60% 内存占用、把训练吞吐提升 1.5–2$\times$。随着上下文长度超过 32K,padding 浪费本会主导内存消耗、显著降低有效 batch size,这个效率收益因而愈发关键。
6.4.2 Packed Sequence 的 Dynamic Context Parallelism
THD 格式的 packed sequence 消除了 padding 浪费,但它并没有消除各 data-parallel(DP)rank 之间的计算不均衡。即便各 packed 样本的总 token 数完全相同,因为点积注意力对 sub-sequence 长度是二次复杂度,它们的注意力负载仍可能差别很大。比如 Figure 36 里的序列被 pack 到相同总长度,但 Figure 37 揭示它们的计算负载差异显著。
这种 DP 不均衡导致 GPU 在梯度同步点空转,还会进一步放大流水线气泡。此外,启用 Context Parallelism(CP)时,通常要选一个静态的 CP 大小来满足 batch 里最长那条 packed 样本的内存需求。结果那些本可以放进更少设备的短样本也被迫用同样的 CP 大小,招来不必要的 CP 通信。实践中 CP 通信本应被注意力计算藏住;然而在由许多短子序列组成的 packed sequence 下,每个 CP 分片的计算量可能不足以藏住 CP 集合通信,CP 跨节点时尤其如此。我们把这个效应称为 CP 低效。
为了同时应对 DP 不均衡和 CP 低效,我们引入 Dynamic Context Parallelism(Dynamic-CP)94,它与 packing 方案联合地、按 microbatch 自适应选择有效 CP 大小。关键观察是:调整 CP 大小相对轻量,它只改变 token 切片如何划分、以及注意力算子用哪个 CP 通信组,不需要任何参数重分布或优化器状态迁移。因此 Dynamic-CP 为变长训练提供了一种实用的动态并行形式,框架开销极小。相关工作包括 ByteScale95 和 WLB-LLM96,它们应对类似的负载均衡挑战。
举个例子,考虑 Figure 38(a) 里三条序列的情形,假设只有两张 GPU 可用。如果采用标准的 CP2 配置(也就是 DP=1、CP=2),负载会像 Figure 38(b) 那样分两个 micro-batch 执行。microbatch-0 里橙色序列足够长,必须切分到两张 GPU 上。microbatch-1 里蓝色和绿色的 packed 序列也被切分到两张 GPU 上。但这给蓝色和绿色序列引入了不必要的切分,因为它们各自本来都能放进单张 GPU 而不需切分。
如 Figure 38(c) 所示,有了 Dynamic-CP,microbatch-0 里的橙色序列仍然需要 CP=2。而 microbatch-1 里我们不再需要切分蓝色和绿色序列,而是让它们各自以 CP=1 独立运行。
为了让按 microbatch 调整 CP 在 Megatron-Core 里变得实用,Dynamic-CP 避免任何昂贵的模型状态重分布。它把 CP 调整当作一个轻量的运行时选择:只改变 token 切片如何划分、以及注意力算子用哪个 CP 通信组。具体来说,框架不是把每个 rank 绑到单个静态构造的 cp_group 上,而是在初始化时为每个 rank 预先构造多个 CP 组,候选 cp_size 从 1 到 dp$\times$cp(限制为 2 的幂)。运行时调度器为每个 microbatch 选定有效 cp_size 和对应的预建 cp_group,从而在不承担动态创建通信组开销的前提下实现 Dynamic-CP。
Dynamic-CP 是为 THD layout 设计的,那里变长序列在一个长度约束下被 pack,原本的 batch/sequence 维度坍缩成一个 token 维度。一个直接后果是 microbatch 数不再是由 global_batch_size 和 micro_batch_size 推出的固定量,它会在各迭代之间变化,因为 pack 进每个 microbatch 的原始序列数并不恒定。为了把对现有 Megatron-Core 流水线调度器的侵入式修改降到最低,Dynamic-CP 在原本的 data_iterator 外面引入一个轻量的 data_iterator_wrapper。对每个 global batch,这个 wrapper 会 (i) 重新调度并 pack 序列,以在各 DP rank 之间创造均衡的负载;(ii) 为每个 microbatch 选择有效 cp_size,避免对短的 packed 样本过度切分;(iii) 返回当前迭代的有效 num_microbatches。因为在 PP/VPP 下通常只有一部分 rank 驱动调度决策,框架会广播动态调度元数据(包括 num_microbatches、max_seqlen 和 cu_seqlens),以保证各流水线 stage 执行一致。此外 PackedSeqParams 被扩展成同时携带选定的 cp_size 和 cp_group;所有依赖 CP 的组件(比如 position embedding 和注意力)都从 PackedSeqParams 取 CP 配置,而不再依赖全局静态的 CP 变量。
在 THD 里给定变长序列,不同样本贡献的有效 token 数不同。因此损失按 token 计算,以避免 padding 带来的偏差:
$$ \mathcal{L}=\frac{\sum_{t\in\mathcal{V}}\ell_{t}}{|\mathcal{V}|} $$其中 $\mathcal{V}$ 表示 packed 表示里有效(非 padding)token 的集合。在规划一侧,Dynamic-CP solver 在 GPU 内存约束下联合决定 packing 和每 microbatch 的 CP 大小。由于注意力计算相对 sub-sequence 长度大致按 $\mathcal{O}(S^2)$ 增长,而激活内存更接近 $\mathcal{O}(S)$,同时平衡计算和内存很困难。所以 Dynamic-CP 在面向负载和面向内存的决策之间交替:估计负载超出目标配额的 microbatch 被分配更大的 cp_size 以降低每 rank 计算量,之后内存变成主导约束,剩余容量通过挑选计算量较轻的样本来填满,同时保持可行性。
最后,Dynamic-CP 在工程上让运行时开销可以忽略。构造一个 plan 需要对 global batch 多扫一遍,以探测轻量的形状和序列长度元数据;这份 I/O 压力通过把探测分散到整个集群、只 gather 紧凑的元数据来缓解。solver 本身异步运行(比如在 data_sampler 里)以与训练迭代重叠。为了让搜索空间可控,穷举搜索被替换成对 microbatch 数的一维 grid search,并约束所有 DP rank 用同一个 num_microbatches。这个数从 PP$\times$1 扫到 PP 的一个小倍数,捕捉每 microbatch 负载与流水线气泡之间的取舍;实践中搜索还能进一步收窄——在负载-microbatch 曲线上找「膝点」,只探索它的邻域。
总的来说,Dynamic-CP 通过 (i) 减少 DP 不均衡造成的同步停顿、(ii) 为那些不受益于大 CP 的 microbatch 避免不必要的 CP 通信,来补足 packed sequence 训练,从而提升长上下文变长训练的吞吐。根据 benchmark,在序列长度分布高度不均衡的真实场景(比如多模态训练)里,Dynamic-CP 带来 35–60% 的端到端性能提升97。
6.5 小结
长上下文 MoE 训练代表一个独特的优化 regime,§4 的基本假定必须被重新审视。三堵墙(内存、通信、计算效率)仍是基本障碍,但随序列长度增加,它们的相对重要性发生了显著转移:
- 计算重心转移:注意力的 $O(s^2)$ 复杂度把 MoE 层从主导计算瓶颈的位置上挤了下来,在超过 64K token 时吃掉超过 69% 的 FLOPs。§4.3 那些聚焦 MoE 的优化仍然有价值,但已不再针对主要瓶颈。
- 内存墙被加剧:激活内存随序列长度增长,需要 §3 的并行策略(具体是 CP 和 TP)配合 §4.1 的内存技术(重计算、offload)才能维持可行的内存占用。
- 通信取舍:CP 与 TP 之间的选择涉及通信开销、内存节省、SDPA 计算效率之间微妙的取舍,这个平衡与 §4.2 里以 MoE 为中心的 all-to-all 优化重心明显不同。
通过组合 context parallelism、tensor parallelism、选择性重计算和优化器 offload,Megatron-Core 能在各种序列长度上高效训练。我们在 256 张 Hopper GPU、256K 序列长度上的实验证明了这一点:DeepSeek-V3 用 TP、优化器 CPU offload 和选择性重计算达到其短上下文 MFU 的 88%,而 Qwen3-235B-A22B 用 TP、CP 和选择性重计算达到短上下文 MFU 的 129%(后者超过 100% 是因为在更长序列上 SDPA 主导计算时,SDPA kernel 效率很高)。packed sequence 和 Dynamic-CP 这两项特性进一步把这些能力扩展到 RL 和 SFT 训练里的变长负载。
7 Production Features
前面几节聚焦于应对三堵墙的性能优化。然而生产级的 MoE 训练还需要健壮性、灵活性和易用性。本节描述应对这些运维需求的特性:保证训练稳定的负载均衡、shared expert 与 latent expert 架构、支持灵活部署的分布式 checkpoint、用现有 dense 模型的 upcycling、以及与 multi-token prediction 等高级训练范式的集成。
7.1 负载均衡与 Token Dropping
MoE 模型里的动态路由会导致显著的负载不均衡,某些专家收到的 token 远多于其他专家。这造成计算低效、内存瓶颈和硬件利用率下降。我们用两个协同的机制应对:负载均衡和 token dropping。
负载均衡。Figure 39 展示了 Megatron-Core MoE 里可用的负载均衡策略。我们提供多种策略来鼓励 token 均匀分布。主要做法采用辅助损失135,这是一个可微的惩罚项,抑制把所有 token 都路由到一小部分专家。我们也支持 expert choice routing98,它把均衡路由建模成一个最优传输问题;以及通过可学习的 expert bias 项实现的 auxiliary-loss-free 均衡23,它根据历史负载动态调整路由决策。
Token Dropping。我们支持两种容量管理策略不同的 dispatch 策略。在 dropless 模式(默认)下,所有被路由的 token 都被处理、没有容量约束,以每专家负载可变为代价最大化模型表达力。在 droppable 模式下,显式的专家容量上限被强制执行135:当分配给某个专家的 token 超过它的容量时,多余的 token 被丢弃、经残差连接绕过。这提供了可预测的内存上界,在 router 初始化不佳的训练早期很有用。
Pad-to-Max 实现静态形状。droppable 模式还启用了 pad-to-max 功能,把所有专家输入 pad 到同一个容量上限。这把 MoE 本质上动态的每专家 token 数转成静态形状,从而使能那些要求 tensor 维度在各迭代间固定的优化,比如 CUDA Graphs(§4.3.6)。
7.2 Shared Expert
有些 MoE 架构(DeepSeek-V2/V3187、Qwen33)包含一个 shared expert,它不管路由如何都处理所有 token,为所有 token 提供一致的基线容量(Figure 40)。启用 overlap(--moe-shared-expert-overlap)时,shared expert 的计算与 all-to-all 通信及 routed expert 计算并行进行,把它的延迟藏到 dispatch-compute-combine 流水线后面。
7.3 Latent MoE
在标准 MoE 里,all-to-all 通信按完整 hidden 维度 $d$ dispatch token,每个专家由 $\mathbb{R}^{m\times d}$ 和 $\mathbb{R}^{d\times m}$ 的权重矩阵参数化。LatentMoE20 通过在 dispatch 之前插入一个 shared down-projection $W_{\downarrow}\in\mathbb{R}^{\ell\times d}$、在 combine 之后插入一个 up-projection $W_{\uparrow}\in\mathbb{R}^{d\times\ell}$(其中 $\ell 压缩比 $\alpha=d/\ell$ 把 all-to-all 通信量降低 $\alpha$ 倍(token 按维度 $\ell$ 而非 $d$ dispatch),并把每专家权重大小降低同样的倍数(专家矩阵从 $\mathbb{R}^{m\times d}$ 缩到 $\mathbb{R}^{m\times\ell}$)。Elango 等人20 定义了两种利用这份节省的方式。在 $\ell\text{-MoE}_{\text{eff}}$ 里,总专家数 $E$ 被放大 $\alpha$ 倍而 top-$K$ 不变,在降低推理成本的同时保住基线精度。在 $\ell\text{-MoE}_{\text{acc}}$(推荐)里,$E$ 和 top-$K$ 都被放大 $\alpha$ 倍,这把推理成本恢复到标准 MoE 水平,但指数级扩展了专家选择的组合空间($\binom{\alpha E}{\alpha K}\geq\binom{E}{K}^{\alpha}$),在等成本下带来更高精度。在最多 95B 参数的规模上,$\ell\text{-MoE}_{\text{acc}}$ 在每 FLOP 精度和每参数精度上都一致优于标准 MoE。这个架构已被 NVIDIA 的 Nemotron-3 Super 和 Ultra 模型采用。 实现。LatentMoE 通过 传统 checkpointing 把保存的模型状态与特定并行配置紧耦合,改 TP、EP、PP 或其他并行设置时需要复杂的离线转换。我们的分布式 checkpoint 库通过与并行配置无关的 checkpointing 加自动 resharding 解决了这个问题。 核心抽象是 这实现了任意到任意的并行重配置:用 $\text{TP}=2$、$\text{EP}=4$ 保存的 checkpoint 可以用 $\text{TP}=4$、$\text{EP}=8$ 加载,无需离线转换。这个库支持 Zarr(默认)和 PyTorch Distributed54 两种存储后端。 Interleaved pipeline parallelism99 把每个物理流水线 stage 切成多个虚拟 stage 并交错调度,以缓解流水线气泡。传统 VPP 要求层均匀分布(比如 24 层模型在 $\text{PP}=4$、$\text{VPP}=2$ 下按 $[6,6,6,6]$ 分层)。但 MoE 模型呈现出显著的负载异质性:MoE 层、dense 层、embedding、loss,以及 MTP 这类特殊层的计算成本差别很大。 我们引入 Flexible Asymmetric VPP,允许每个虚拟 stage 有不同数量和类型的层。考虑 61 层 decoder 加 1 层 MTP 的 DeepSeek-V3 在 $\text{PP}=16$、$\text{VPP}=2$ 下的情形(Table 10)。首个 rank 把轻量的 embedding 与 3 个 dense decoder 放在一起(成本与 2 个 MoE 层相当)。大多数 rank 每个 stage 放 2 个 MoE decoder。最后几个 rank 策略性地摆放重的 MTP 层和轻的 loss 层来平衡负载。 这个做法让任意层数的模型都能用上 VPP,通过考虑每层计算成本实现真正的负载均衡,还为特殊层提供灵活的摆放位置。 Table 10:DeepSeek-V3 在 flexible asymmetric VPP 下的层分布($\text{PP}=16$、$\text{VPP}=2$)。 这种灵活非对称的做法(Figure 41)带来几个关键优势:(1) 让任意层数和层组成的模型都能用 VPP,移除了对模型架构设计的人为约束;(2) 通过允许对层分布做细粒度控制、把不同层类型之间巨大的计算成本差异纳入考虑,实现真正的计算负载均衡;(3) 为特殊层(MTP、encoder-decoder 结构等)提供灵活的摆放策略,让多样的模型架构都能充分利用流水线并行的延迟隐藏收益。实践者可以设计出把计算昂贵的层分散到各 rank 的布局,以避免内存压力、最小化流水线气泡。 Upcycling 把预训练好的 dense 模型转换成稀疏 MoE 架构,在不从头重训的前提下扩展模型容量100101102。我们支持 virtual group 初始化和专家权重 scaling,让 dense checkpoint 能无缝适配成细粒度 MoE 模型。这个做法采用 softmax-then-topK 路由,把 token 经一个选择性的专家子集路由,在推理成本不变或更低的情况下提升表达力,如 Figure 42 所示。 Multi-token prediction(MTP)优化模型在每个位置预测多个连续的未来 token,从而稠密化监督信号1037。与并行的独立预测不同,MTP 通过 hidden state 转移在各预测之间维持因果依赖,加速收敛、改善生成质量。推理时模型退回标准的单 token 预测以保证部署兼容性。 我们把 MTP 与灵活的流水线并行集成,允许 MTP 层被策略性地摆放在 VPP 布局里以实现均衡的负载分布。 与 AdamW 这类做逐元素更新的传统优化器不同,Muon 优化器引入了一种矩阵感知的做法,对整个权重矩阵做正交化104。这改善了优化轨迹的 conditioning,相比 AdamW 显著减少训练步数。 我们的集成为大规模分布式训练提供生产可用的支持,有几个关键优势:(1) 完整支持切分的 query-key-value(QKV)权重 layout,即便注意力投影矩阵被存成独立 tensor 也能高效正交化;(2) 与 distributed optimizer 无缝集成,允许优化器状态在各 data-parallel rank 上被切分,同时维持正确的正交化语义;(3) 当 GPU 内存受限时,为 Muon 的正交化 buffer 提供 CPU offload。 MuonClip。训练万亿参数模型带来稳定性挑战,query-key 点积可能无界增长、造成注意力爆炸105。MuonClip 通过 cuDNN、cudnn-frontend 和 Transformer Engine 里的硬件加速实现来应对这个问题。 前面几节描述了 Megatron-Core MoE 的架构、并行策略和优化。本节通过实证评测验证这些贡献,展示 Megatron-Core MoE 栈在多种模型配置和硬件平台上的性能。 Benchmark 模型。我们在两个最先进的细粒度 MoE 架构上评测 Megatron-Core MoE,它们对框架能力构成压力测试:DeepSeek-V3-685B7 和 Qwen3-235B3391。两者都有细粒度专家设计,把三堵墙的挑战放大:高专家数消耗内存容量、top-k 路由增加 all-to-all 通信量、更小的每专家计算降低 GEMM 效率。这些特征让它们成为验证我们优化策略的严苛 benchmark。 硬件与软件栈。benchmark 在 NVIDIA GB200 和 H100 GPU 上进行,以评估跨硬件代际的性能。我们用 Megatron-LM 的 dev 分支106 配最新的 TransformerEngine107。§4 描述的完整优化栈全部启用,包括但不限于: 指标。我们报告每 GPU 的 TFLOPS 吞吐和每 GPU 每秒处理的 token 数。这两个指标分别给出计算效率和实际训练吞吐两个视角。 概览。Table 11 汇总了在 NVIDIA GB300、GB200 和 H100 系统上、§4 描述的优化栈完全启用时,代表性细粒度 MoE 训练负载的每 GPU 吞吐。所有配置都用 force-balanced routing 以保证 token 在各专家间均匀分布。 DeepSeek-V3 和 Qwen3-235B(大规模训练)。Table 11 报告了两个当代细粒度 MoE 训练负载在标准序列长度 4,096 下的结果,包括一次 1,024 卡 H100 上的 DeepSeek-V3 run。在这些设定下,测量结果显示三个平台上都有很强的每 GPU 效率和高的持续吞吐。 Table 11:两个 mixture-of-experts 模型在 NVIDIA GB300、GB200 和 H100 上的统一吞吐 benchmark(每 GPU 数字)。所有配置都用 force-balanced routing。Dtype 列指明 FP8 recipe:FP8-BLK 表示 Hopper 上的 blockwise FP8,MXFP8 表示 Blackwell 上的 microscaling FP8。 FP8 训练的有效性。如 Table 11 所示,DeepSeek-V3 的 FP8 训练在 H100 上达到每 GPU 368 TFLOPS,说明 blockwise FP8 recipe 让大规模 MoE 模型在 Hopper 架构上高效训练成为可能。§4.3.5 讨论的优化成功应对了把 FP8 套用到有动态专家路由和多变计算模式的 MoE 模型上的挑战,在维持数值稳定性的同时提供低精度的性能收益。 GB200 与 GB300 的性能。对两个 MoE 模型,GB200 和 GB300 平台在相当或更少的 GPU 数下,交付了约比 H100 高 3$\times$ 的 token 吞吐。这份效率提升来自更高的内存带宽、更高的算力和原生 MXFP8 Tensor Core 支持。配合我们的软件优化(§4.1 和 §4.3),GB200 为 DeepSeek-V3 训练达到每 GPU 超过 1,048 TFLOPS,GB300 进一步推到每 GPU 1,233 TFLOPS。 可扩展性。这些结果证明 Megatron-Core MoE 能高效扩展到数千 GPU、几千亿参数的生产负载。DeepSeek-V3 在 1,024 卡规模上的成功部署验证了 parallel folding(§3.3)、优化 dispatcher(§4.2)和 kernel 优化(§4.3)的有效性。 长上下文压力测试。为了捕捉标准 4,096 token regime 之外的行为,Table 11 包含一次 GB300 上序列长度 131,072 的长上下文 Qwen3-235B run。即便在这个内存和通信双重密集的设定下,系统仍维持每 GPU 1,150 TFLOPS,说明在完整优化栈下长上下文效率很强。 配置细节。Table 11 里每个 benchmark 条目所用的并行配置(TP、PP、CP、EP、DP、VPP 等)、batch 设置和其他优化细节见附录 B。注意这些配置是我们通过实证调优找到的最佳设定,可能不是全局最优——对这个规模的模型,穷举搜索完整的并行与优化配置空间是不可行的。 §8 证明了 Megatron-Core MoE 在最新 GPU 平台上、在多种 MoE 架构上都达到了 state-of-the-art 性能。下面我们考察这些结果背后的系统化方法,并通过详细案例研究展示 §4.1–§4.3 的优化如何转化为部署决策。 §3–§7 给出了一整套 MoE 训练的优化武器库:parallel folding、优化 dispatcher、Grouped GEMM、低精度训练、CUDA Graphs 等等。然而拥有很多优化技术并不自动转化为高性能。每种技术针对特定瓶颈,但也引入自己的取舍。低精度训练降低内存但要求仔细的精度管理。通信 overlap 提升吞吐但增加内存开销。CUDA Graphs 消除 host 延迟但与动态 tensor 形状冲突。挑战不在于有没有优化可用,而在于知道该用 哪些 优化、什么时候 用、以及它们彼此 如何 相互作用。 这份复杂度意味着达到峰值性能需要系统化的方法,而不是试错式的调参。通过在 GB200 和 H100 上优化多种 MoE 模型(从 Mixtral 到 DeepSeek-V3),我们发展出一套可重复的工作流来定位瓶颈、套用针对性解法。本节把那些经验提炼成可操作的指导。§9.1 在三堵墙框架下给出方法论,每个阶段配具体例子。§9.2 展示这些原则如何在 GB200 和 H100 上为 DeepSeek-V3 组合起来,包括为什么同一个模型在不同硬件上需要不同策略。 我们给出一个三阶段的 MoE 性能优化工作流。这套方法论来自在 GB200 和 H100 平台上调优 Mixtral、DeepSeek-V3 和 Qwen3 的经验。这个过程本质上是迭代的:解决一个瓶颈往往暴露下一个,需要持续 profiling 和精化。 内存可行性是第一道约束。在为吞吐优化之前,配置必须先装进 GPU 内存。Table 12 汇总了各并行策略如何影响每 GPU 内存和通信开销。 Table 12:并行策略对内存和通信的影响。$d$ = 并行度。†需要 distributed optimizer( 快速可行性测试。用 交互式内存估算器。我们开发了一个带 web GUI 的交互式内存模拟器,用于针对内存占用快速调优并行配置。细节见博客111。 例子。对一个 685B 参数的细粒度 MoE 模型,仅 BF16 激活就可能超过每 GPU 130 GB(见 Table 3),即便用上并行也立刻排除了 80 GB 设备上的基线配置,这也就引出了阶段 3 的内存优化。 一旦识别出内存可行的配置,就选那个在维持吞吐的前提下最小化通信开销的策略。最优选择取决于模型架构、序列长度和硬件拓扑。 注:对 DeepSeek-V3 这类超大 MoE 模型,EP 通信量可能超出 NVLink 能藏住的范围。这种情况下要启用通信-计算 overlap(见 §4.2),把 EP 延迟藏到计算后面。 对 MoE 层,EP 相比 TP 有几个优势: 例子:对 Mixtral-8$\times$7B,EP8$\times$TP1 优于 EP4$\times$TP2。 例子。考虑一个 NVL72 系统上的 256 专家 MoE 模型。套用准则 4,Parallel Folding 把 expert TP 设为 1,让每个专家跑在单张 GPU 上,最大化 GEMM 效率。有了准则 2,EP64 完整装进 NVLink 域。准则 1 随后驱动余下的选择:可用内存预算决定注意力层的 TP 和 PP,DP 填掉剩下的。 有了能跑的并行配置之后,profile 训练来判断哪堵墙在主导。按瓶颈类型套用针对性优化。 症状:为避免 OOM 而被迫使用完全重计算或过大的并行度。 解法:套用省内存的技术(Table 13)降低内存消耗。 Table 13:内存瓶颈的解法。 症状:profiling 显示大量时间花在集合通信操作上。 解法:识别哪种通信是瓶颈,套用对应的优化(Table 14)。 Table 14:通信瓶颈的解法。 计算效率墙有两种不同的表现形式:CPU 开销和计算低效。CPU 开销发生在 host 发起 GPU kernel 不够快、在 GPU 执行里制造出空隙时。 症状:Nsight Systems 时间线显示 GPU kernel 之间有空隙,CPU 发起 kernel 跟不上。 解法:通过减少 kernel 发起数量和启用 CUDA Graphs 来降低 host 侧开销(Table 15)。 Table 15:CPU 开销瓶颈的解法。 计算低效发生在 GPU kernel 自身用不满硬件资源时,通常是因为细粒度 MoE 架构里 GEMM 尺寸太小。 症状:没有通信或 CPU 瓶颈,但 GPU SM 利用率低。 解法:通过 batching、fusion 和更低精度改善 kernel 效率(Table 16)。 Table 16:计算瓶颈的解法。 例子。在 EP 跨节点的 NVL8 系统上,profiling 可能显示 all-to-all 通信吃掉 30–50% 的 step 时间,直指通信墙。而在同样 EP 度数留在 NVLink 域内的 NVL72 系统上,启用 FP8 之后主导瓶颈常常转成 CPU 开销——更快的 GPU 计算把 host 侧发起延迟暴露了出来。同一个模型在不同硬件上可能需要完全不同的优化策略。 三个阶段的顺序很重要:内存约束是必须先解决的硬障碍(阶段 1),并行决策决定通信拓扑(阶段 2),只有这之后 profiling 才能揭示该攻哪堵墙(阶段 3)。关键是这个过程是迭代的。内存优化可能让更小的并行度成为可能,于是回到阶段 1。有些阶段 3 的优化本身有内存成本(EP 通信 overlap 需要额外 buffer,CUDA Graphs 消耗额外内存),这可能要求重新审视之前的决策。每轮之后持续 profiling 来指导下一次迭代。 DeepSeek-V3 是 MoE 训练的一个极端压力测试:685B 总参数,带 Multi-Token Prediction(MTP)和 Multi-Latent Attention(MLA)。三堵墙全部适用:256 个专家带来的内存压力、top-8 路由带来的高 all-to-all 量、以及小的每专家 GEMM。前面的工作流(§9.1)展示了如何单独处理每个阶段;这个案例研究考察那些决策如何相互作用、形成一个连贯的优化栈,以及为什么同一个模型在 GB200 和 H100 上需要根本不同的策略。 Table 17 汇总了 GB200 和 H100 平台上最终的优化配置和性能。 Table 17:DeepSeek-V3 在 GB200 和 H100 上的最终优化配置。†用了 Parallel Folding;TP 只作用于非 MoE 模块,expert TP 恒为 1。 GB200 配置用 HybridEP(针对 NVL72 优化)和 CUDA Graphs 来最小化 CPU 开销,后者在 Blackwell 上的 FP8 训练里成为主导瓶颈。H100 配置用 FP8-blockwise 精度配 DeepEP 和 EP all-to-all overlap 来隐藏通信延迟。 Table 18 汇总了各平台上套用的关键优化。 Table 18:DeepSeek-V3 按平台的优化汇总。 两个平台都用 Parallel Folding(准则 4)来解耦注意力和专家的并行:TP 只作用于注意力层以降低内存,而专家用 EP 配 TP $=1$ 以获得更好的 GEMM 效率。选 EP64 是为了让每张 GPU 正好持有 256 个专家中的 4 个,消除局部 token permutation 开销。 关键的平台差异是内存容量。GB200 每 GPU 192 GB(对比 H100 的 80 GB)允许用 TP1/PP4 而不是 TP2/PP8,把流水线深度减半,从而减少流水线气泡、简化负载平衡(准则 1)。Flexible VPP 随后微调流水线:H100 上,布局 FP8 训练是第一根杠杆:它在两个平台上都把激活内存减半,释放出本来需要激进重计算才能腾出的 GPU 内存。在 H100 上,这份释放出来的内存预算至关重要,因为它使 EP 通信 overlap 成为可能——后者需要额外 buffer 来流水化 dispatch 和 combine 操作(见下面的通信墙)。细粒度重计算(mlp、mla_up_proj、layernorm、moe_act)削掉剩余的内存压力。Memory-efficient permutation 消除冗余的激活存储,低精度优化器状态(BF16)再提供额外余量。 在 GB200 上,这条内存链的解法不同。NVL72 让 EP 保持在本地,所以通信 overlap 不再关键、它的内存预算被完全释放。GB200 更高的 C2C 带宽(NVLink-C2C)让优化器状态 offload 变得有效:传输开销低到 offload 能为激活释放出大量 GPU 内存。这反过来把细粒度重计算的需要缩小到只有 mlp,远比 H100 上温和。 在 H100/B200(NVL8)上,EP64 跨越 8 个节点,用标准 all-to-all dispatcher 的话跨节点 all-to-all 延迟会吃掉近 50% 的 step 时间。DeepEP 通过 fused permutation 和优化的集合通信操作降低这份开销。即便如此,EP 通信 overlap 仍是必需的,用来把剩余的跨节点延迟藏到专家计算后面;而这个 overlap 之所以可行,正是因为 FP8 为那些额外 buffer 释放了足够内存(见上面的内存墙)。 在 GB200(NVL72)上,EP64 完整留在 NVLink 域内。HybridEP 无需通信 overlap 就能充分利用 1.8 TB/s 的双向带宽。通信墙实际上仅靠硬件拓扑就被解决了,主导瓶颈随之转向计算效率。 FP8 在两个平台上都加速 GEMM(H100 上是 blockwise FP8,GB200 上是有原生 Blackwell Tensor Core 支持的 MXFP8),但这份加速有个副作用:更快的 GPU 计算把 CPU 开销暴露出来。在 GB200 上,NVL72 已经消除了通信瓶颈,CPU 开销于是成为主导约束——host 发起 kernel 的速度不足以让 GPU 饱和。 Partial CUDA Graphs 通过把注意力、router 和 MoE preprocessing 捕获成静态 graph 来应对,同时把动态的专家计算留在 graph 之外。Kernel fusion(router fusion、permute fusion、MLA RoPE fusion)减少 kernel 发起数量。CPU/NUMA 绑定112 降低 host 侧内存访问延迟。 这些优化的横切特性贯穿全程:FP8 同时降低内存(激活减半)和提升计算吞吐,但引入的量化 kernel 又放大了 CPU 开销。同一个模型在不同硬件上落到了根本不同的优化栈:H100 的栈以隐藏通信延迟为中心(DeepEP + EP overlap),而 GB200 的栈以消除 CPU 开销为中心(CUDA Graphs + kernel fusion)。这个对照验证了 §9.1 那套 profile 驱动的迭代方法。 四点洞察可以推广到 DeepSeek-V3 之外: 继 OpenAI o1 和 DeepSeek-R185 之后,强化学习(RL)后训练已成为一个关键范式。许多领先的 RL 训练模型都是 MoE 架构(DeepSeek-R1、Kimi-K2 等),而 RL 负载引入的挑战把本报告通篇讨论的 MoE 专属问题进一步放大:高度可变的序列长度给内存墙施压,交错的推理与训练阶段要求快速的内存 offload,推理引擎与训练引擎之间的路由差异威胁稳定性。为了在规模上达到高训练吞吐,RL 框架把 Megatron-Core 跑在 Ray worker 里作为分布式训练后端,并用 Megatron-Bridge 做 Hugging Face 与 Megatron checkpoint 之间的双向转换。 RL 后训练与预训练的差异,给训练引擎带来了新要求: 一个典型的 RL 工作流从预训练的 Hugging Face 模型出发,在 RL 框架里微调它。在规模上,这些框架常常把 Megatron-Core 跑在 Ray worker 里做分布式训练,而 rollout 通常用一个期望 Hugging Face 格式 checkpoint 的推理栈。这产生两项集成需求:把模型定义映射到 Megatron-Core 模块,以及在 Hugging Face 与 Megatron 格式之间高效转换 checkpoint。Megatron-Bridge 解决 checkpoint 互操作这一层,为初始化、训练和导出提供快速的 HF-to-Megatron 转换。这个模式被多个 RL 框架采用,包括 veRL113、Slime114 和 NeMo RL。 Packed Sequence 支持。如 §6.4.1 所讨论,我们可以 pack 序列来去掉变长 batch 里的 padding。在 packed sequence 支持之上,RL 框架一侧我们推荐用 packing-aware 的动态 batch size,保证每个 batch 有相近的有效 token 总数。此外我们引入一个额外的负载均衡策略,它显式地考虑 transformer block 的异质计算成本。因为注意力模块随序列长度二次增长($O(L^2)$)而前馈网络(FFN)线性增长($O(L)$),一个在 token 数上很均衡的 batch 在 wall-clock 时间上仍可能高度不均衡。为了缓解这个问题,我们先为 batch 里每个样本计算其序列长度的平方,然后在 mini-batch 内求和。接着按这个「注意力成本」指标排序 micro-batch,并按小-大-小的模式调度(也就是先升序再降序)。这种蛇形排序带来两个好处:(i) 对数据并行、流水线并行和专家并行,它减少同步气泡,因为相邻 micro-batch 的注意力负载相当;(ii) 在流水线并行内部,传统上各 stage 利用率不足的 warm-up 和 cool-down 阶段空闲时间更短,因为更轻的 micro-batch 在调度里出现在更早和更晚的位置。 Dynamic Context Parallelism。如 §6.4.2 所讨论,常规训练设置必须挑一个固定的 context-parallel(CP)度,来保证负载里最长的序列不会触发 OOM。这个保守的选择迫使内存占用小得多的大多数短序列也用不必要的大 CP 度运行,从而浪费跨设备带宽、降低整体吞吐。Dynamic CP 通过按 micro-batch 自适应选择 CP 度来移除这种一刀切的限制。运行时我们按长度把进来的序列分桶,为每个桶选择能让它留在设备内存预算内的最小 CP 度,然后据此派发 micro-batch。长序列拿到更高的 CP 度以保证内存安全,而短序列以更低的 CP 度执行,从而最大化算术强度、降低通信量。 CPU 优化器 offload。为了进一步缓解 GPU 内存压力,我们在前向和反向传播期间把优化器状态 offload 到 host DRAM,只在参数更新那一步把它们载回 GPU。因为优化器 tensor 可以占到模型参数的数倍大小,临时把它们腾出去能释放大量高带宽 GPU 内存,转而用于缓存激活或容纳更长的序列。 FP16 训练支持。虽然 bfloat16(BF16)在大语言模型预训练里被广泛采用,若干研究观察到在强化学习训练期间、在某些超参选择下,半精度浮点(FP16)能带来更好的数值稳定性。因此 Megatron-Core MoE 实现提供了一条功能完整的 FP16 路径,包括 loss scaling 和混合精度优化器 kernel,让实践者能在不牺牲吞吐的前提下选择最符合其稳定性要求的精度模式。 Router Replay。近期工作115 显示:捕获 MoE router 在推理期间产出的路由决策、并在后续训练阶段回放它们,可以改善收敛一致性。得益于一个社区贡献的 patch,Megatron-Core MoE 现在支持这个特性:推理引擎记录每个 token 的专家分配,训练栈可以摄入并强制执行同样的路由模式。这个机制把路由的可变性与权重更新解耦,带来更稳定的优化轨迹,在 on-policy 数据分布频繁漂移的 RL 设定下尤其有效。 本报告介绍了 Megatron-Core MoE,一个训练大规模 Mixture-of-Experts 模型的开源栈。MoE 的稀疏性引入两个根本挑战:表现为三堵墙(内存、通信、计算效率)的参数-计算失配,以及要求注意力层与 MoE 层解耦并行的 dense-sparse 失配。通过在这些挑战上组合集成式的解法,Megatron-Core MoE 使高效的万亿参数级训练成为可能。 关键技术贡献包括: 这些优化让 Megatron-Core MoE 达到很强的训练吞吐:DeepSeek-V3(685B 参数、256 专家)在 256 张 GB300/GB200 上以每 GPU 1,233/1,048 TFLOPS 训练,在 1,024 张 H100 上每 GPU 368 TFLOPS;Qwen3-235B 在 GB300/GB200 上达到 974/919 TFLOPS,在 H100 上 320 TFLOPS。GB300 和 GB200 平台交付了约比 H100 高 3$\times$ 的 token 吞吐,展示了这个框架利用下一代硬件的能力。 通过开源这些能力,Megatron-Core MoE 为研究者和实践者提供了生产级的工具,用于规模化的 MoE 实验和部署。模块化的架构支持快速原型验证,而完整的优化栈支撑从研究原型到万亿参数生产模型的训练。 Table 19:本报告通篇使用的符号与缩写。 本节列出复现 Table 11 中性能数字所需的并行配置和详细设置。每个配置由其系统、GPU 数和精度格式标识,对应的超参字符串汇总了并行布局和 batch 配置。 Table 20 列出与 Table 11 吞吐结果对应的 benchmark 配置。这些设置是撰稿时通过实证调优找到的最佳配置,可能不是全局最优。 Table 20:Table 11 中各 benchmark 条目的并行与训练配置详情。 译注:Table 20 里 H100 上 DeepSeek-V3 的并行配置是 TP2 PP4 EP64,而 Table 17 和 §9.2.2 的正文都写作 TP2/PP8;下面 §B.2 说 DeepSeek-V3 在 H100 上重计算 本节汇总我们 benchmark run 里最影响性能的那些优化特性的配置。虽然完整的优化空间更大,但这些特性对吞吐一致地有一阶影响,并且在不同负载和平台上配置不同。 DeepSeek-V3 这些设置是为了在每个平台上平衡内存余量、通信开销和 kernel 效率,以达到最大持续吞吐。 Qwen3-235B 由此得到的配置模式反映了在模型与硬件各自约束下、吞吐优先的调优策略。 Table 11 里的 benchmark 数字可以通过两条受支持的路径复现。第一,用户可以用 Megatron-Bridge116 跑端到端训练,它为模型、并行和优化配置提供更高层的接口。第二,用户可以用 Megatron-MoE-ModelZoo117 里针对具体模型的脚本直接用 Megatron-Core(MCore)启动训练,那里暴露了完整的底层启动 flag 用于性能调优。两条路径都从 Table 20 出发,只需调整少数与集群相关的运行时设置(比如 hostfile、environment module 和作业调度器选项)。 Scalable Training of Mixture-of-Experts Models with Megatron Core, https://arxiv.org/abs/2603.07685 ↩︎ Scaling Laws for Neural Language Models, https://arxiv.org/abs/2001.08361 ↩︎ Training Compute-Optimal Large Language Models, https://arxiv.org/abs/2203.15556 ↩︎ Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, https://arxiv.org/abs/1701.06538 ↩︎ ↩︎ ↩︎ Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, https://arxiv.org/abs/2101.03961 ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ Mixtral of Experts, https://arxiv.org/abs/2401.04088 ↩︎ ↩︎ DeepSeek-V3 Technical Report, https://arxiv.org/abs/2412.19437 ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ A Survey on Mixture of Experts in Large Language Models, IEEE TKDE 2025, https://arxiv.org/abs/2407.06204 ↩︎ ↩︎ Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, https://arxiv.org/abs/1909.08053 ↩︎ ↩︎ ↩︎ ↩︎ Adaptive Mixtures of Local Experts, Neural Computation 3(1):79–87, 1991 ↩︎ ↩︎ Learning Factored Representations in a Deep Mixture of Experts, https://arxiv.org/abs/1312.4314 ↩︎ ↩︎ Attention Is All You Need, NeurIPS 2017, https://arxiv.org/abs/1706.03762 ↩︎ GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, https://arxiv.org/abs/2006.16668 ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ ↩︎ GLaM: Efficient Scaling of Language Models with Mixture-of-Experts, ICML 2022, https://arxiv.org/abs/2112.06905 ↩︎ Tutel: Adaptive Mixture-of-Experts at Scale, MLSys 2023, https://arxiv.org/abs/2206.03382 ↩︎ ↩︎ ↩︎ DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale, ICML 2022, https://arxiv.org/abs/2201.05596 ↩︎ ↩︎ DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models, https://arxiv.org/abs/2401.06066 ↩︎ DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, https://arxiv.org/abs/2405.04434 ↩︎ ↩︎ ↩︎ ↩︎ NVIDIA Nemotron 3: Efficient and Open Intelligence, 2026 ↩︎ LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts, 2026 ↩︎ ↩︎ ↩︎ Scaling Laws for Fine-Grained Mixture of Experts, ICML 2024, https://arxiv.org/abs/2402.07871 ↩︎ Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient, https://arxiv.org/abs/2502.05172 ↩︎ Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts, https://arxiv.org/abs/2408.15664 ↩︎ ↩︎ ↩︎ DeepEP: an efficient expert-parallel communication library, https://github.com/deepseek-ai/DeepEP ↩︎ ↩︎ ↩︎ MegaBlocks: Efficient Sparse Training with Mixture-of-Experts, MLSys 2023, https://arxiv.org/abs/2211.15841 ↩︎ ↩︎ ↩︎ PyTorch: An Imperative Style, High-Performance Deep Learning Library, NeurIPS 2019, https://arxiv.org/abs/1912.01703 ↩︎ MoE Parallel Folding: Heterogeneous Parallelism Mappings for Efficient Large-Scale MoE Model Training with Megatron Core, https://arxiv.org/abs/2504.14960 ↩︎ ↩︎ ↩︎ Reducing Activation Recomputation in Large Transformer Models, MLSys 2023, https://arxiv.org/abs/2205.05198 ↩︎ ↩︎ ↩︎ Learning (with) Distributed Optimization, 2023 ↩︎ NVIDIA Megatron Core MoE user guide, 2024 ↩︎ GLU Variants Improve Transformer, https://arxiv.org/abs/2002.05202 ↩︎ NVIDIA Transformer Engine, https://github.com/NVIDIA/TransformerEngine ↩︎ Qwen2.5 Technical Report, https://arxiv.org/abs/2412.15115 ↩︎ ↩︎ ↩︎ ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020, https://arxiv.org/abs/1910.02054 ↩︎ ↩︎ ↩︎ DeepSpeed: System Optimizations Enable Training Deep Learning Models with Over 100 Billion Parameters, KDD 2020 ↩︎ The Llama 3 Herd of Models, https://arxiv.org/abs/2407.21783 ↩︎ LLaMA: Open and Efficient Foundation Language Models, https://arxiv.org/abs/2302.13971 ↩︎ Llama 2: Open Foundation and Fine-Tuned Chat Models, https://arxiv.org/abs/2307.09288 ↩︎ GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019, https://arxiv.org/abs/1811.06965 ↩︎ ↩︎ PipeDream: Generalized Pipeline Parallelism for DNN Training, SOSP 2019 ↩︎ ↩︎ Zero Bubble Pipeline Parallelism, ICLR 2024, https://arxiv.org/abs/2401.10241 ↩︎ Pipeline Parallelism with Controllable Memory, NeurIPS 2024, https://arxiv.org/abs/2405.15362 ↩︎ Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines, SC 2021, https://arxiv.org/abs/2107.06925 ↩︎ PipeOptim: Ensuring Effective 1F1B Schedule with Optimizer-Dependent Weight Prediction, https://arxiv.org/abs/2312.00839 ↩︎ Measuring the Effects of Data Parallelism on Neural Network Training, JMLR 2019, https://arxiv.org/abs/1811.03600 ↩︎ Modern Distributed Data-Parallel Large-Scale Pre-training Strategies for NLP Models, 2022 ↩︎ Sequence Parallelism: Long Sequence Training from System Perspective, ACL 2023, https://arxiv.org/abs/2105.13120 ↩︎ Ring Attention with Blockwise Transformers for Near-Infinite Context, https://arxiv.org/abs/2310.01889 ↩︎ ↩︎ DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models, https://arxiv.org/abs/2309.14509 ↩︎ ↩︎ A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training, ICS 2023, https://arxiv.org/abs/2303.06318 ↩︎ Training Deep Nets with Sublinear Memory Cost, https://arxiv.org/abs/1604.06174 ↩︎ ↩︎ Optimal Checkpointing for Heterogeneous Chains: How to Train Deep Neural Networks with Limited Memory, https://arxiv.org/abs/1911.13214 ↩︎ ZeRO-Offload: Democratizing Billion-Scale Model Training, USENIX ATC 2021, https://arxiv.org/abs/2101.06840 ↩︎ PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, https://arxiv.org/abs/2304.11277 ↩︎ ↩︎ ↩︎ ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning, SC 2021, https://arxiv.org/abs/2104.07857 ↩︎ ZeRO++: Extremely Efficient Collective Communication for Giant Model Training, https://arxiv.org/abs/2306.10209 ↩︎ Adam: A Method for Stochastic Optimization, ICLR 2015, https://arxiv.org/abs/1412.6980 ↩︎ 8-bit Optimizers via Block-wise Quantization, ICLR 2022, https://arxiv.org/abs/2110.02861 ↩︎ DeepSeek-V3 Technical Report(原文此处与 7 为重复条目), https://arxiv.org/abs/2412.19437 ↩︎ Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models, ASPLOS 2023 ↩︎ FLUX: Fast Software-based Communication Overlap on GPUs through Kernel Fusion, https://arxiv.org/abs/2406.06858 ↩︎ 1F1B-based EP A2A overlapping for MoE models, https://zhuanlan.zhihu.com/p/28463368206 ↩︎ DualPipe: A bidirectional pipeline parallelism with zero bubble on NVIDIA GPUs, https://github.com/deepseek-ai/DualPipe-with-HybridFlow ↩︎ TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018, https://arxiv.org/abs/1802.04799 ↩︎ ↩︎ Pretraining Large Language Models with NVFP4, https://arxiv.org/abs/2509.25149 ↩︎ ST-MoE: Designing Stable and Transferable Sparse Expert Models, https://arxiv.org/abs/2202.08906 ↩︎ CUDA Graphs, CUDA C++ Programming Guide §3.2.8, https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-graphs ↩︎ PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024 ↩︎ ECHO 实现, Megatron-LM PR#2368, https://github.com/NVIDIA/Megatron-LM/pull/2368 ↩︎ Paged Stashing 实现, Megatron-LM PR#2690, https://github.com/NVIDIA/Megatron-LM/pull/2690 ↩︎ Mixed Precision Training, ICLR 2018, https://arxiv.org/abs/1710.03740 ↩︎ FP8-LM: Training FP8 Large Language Models, https://arxiv.org/abs/2310.18313 ↩︎ COAT: Compressing Optimizer States and Activation for Memory-Efficient FP8 Training, https://arxiv.org/abs/2410.19313 ↩︎ FP8 Formats for Deep Learning, https://arxiv.org/abs/2209.05433 ↩︎ ↩︎ Scaling FP8 Training to Trillion-Token LLMs, https://arxiv.org/abs/2409.12517 ↩︎ 对 Blackwell 上的 NVFP4 recipe,因为需要 gather 行方向和列方向两份权重 tensor,节省下来的通信量同样是 50%。 ↩︎ ↩︎ A White Paper on Neural Network Quantization, https://arxiv.org/abs/2106.08295 ↩︎ FP8 Quantization: The Power of the Exponent, NeurIPS 2022, https://arxiv.org/abs/2208.09225 ↩︎ Pretraining Large Language Models with NVFP4(原文此处与 64 为重复条目), https://arxiv.org/abs/2509.25149 ↩︎ ↩︎ MiniMax-M2, https://github.com/MiniMax-AI/MiniMax-M2 ↩︎ Every Activation Boosted: Scaling General Reasoner to 1 Trillion Open Language Foundation(Ling 2.0), https://arxiv.org/abs/2510.22115 ↩︎ Microscaling Data Formats for Deep Learning, https://arxiv.org/abs/2310.10537 ↩︎ Random number generation using cuRANDDx, https://docs.nvidia.com/cuda/curanddx/get_started/introduction.html ↩︎ 1D block scaling factors layout, cuBLAS 文档, https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout ↩︎ DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning, https://arxiv.org/abs/2501.12948 ↩︎ ↩︎ Longformer: The Long-Document Transformer, https://arxiv.org/abs/2004.05150 ↩︎ Big Bird: Transformers for Longer Sequences, NeurIPS 2020, https://arxiv.org/abs/2007.14062 ↩︎ FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022, https://arxiv.org/abs/2205.14135 ↩︎ FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, https://arxiv.org/abs/2307.08691 ↩︎ FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision, https://arxiv.org/abs/2407.08608 ↩︎ Qwen3 Technical Report, https://arxiv.org/abs/2505.09388 ↩︎ ↩︎ Striped Attention: Faster Ring Attention for Causal Transformers, https://arxiv.org/abs/2311.09431 ↩︎ Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance, https://arxiv.org/abs/2107.02027 ↩︎ Speeding up variable-length training with dynamic context parallelism and NVIDIA Megatron-Core, https://developer.nvidia.com/blog/speeding-up-variable-length-training-with-dynamic-context-parallelism-and-nvidia-megatron-core/ ↩︎ ByteScale: Efficient Scaling of LLM Training with a 2048K Context Length on More Than 12,000 GPUs, https://arxiv.org/abs/2502.21231 ↩︎ WLB-LLM: Workload-Balanced 4D Parallelism for Large Language Model Training, https://arxiv.org/abs/2503.17924 ↩︎ 变长序列训练的入口见 Megatron-LM PR#2000 与 PR#2959, https://github.com/NVIDIA/Megatron-LM/pull/2000 、 https://github.com/NVIDIA/Megatron-LM/pull/2959 ↩︎ Mixture-of-Experts with Expert Choice Routing, NeurIPS 2022, https://arxiv.org/abs/2202.09368 ↩︎ Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021, https://arxiv.org/abs/2104.04473 ↩︎ Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints, ICLR 2023, https://arxiv.org/abs/2212.05055 ↩︎ Upcycling Large Language Models into Mixture of Experts, https://arxiv.org/abs/2410.07524 ↩︎ Llama 3 Meets MoE: Efficient Upcycling, https://arxiv.org/abs/2412.09952 ↩︎ Better & Faster Large Language Models via Multi-token Prediction, ICML 2024, https://arxiv.org/abs/2404.19737 ↩︎ Muon is Scalable for LLM Training, https://arxiv.org/abs/2502.16982 ↩︎ Kimi K2: Open Agentic Intelligence, https://arxiv.org/abs/2507.20534 ↩︎ Megatron-LM dev 分支, https://github.com/NVIDIA/Megatron-LM/tree/dev ↩︎ TransformerEngine main 分支, https://github.com/NVIDIA/TransformerEngine/tree/main ↩︎ HybridEP, https://github.com/deepseek-ai/DeepEP/tree/hybrid-ep ↩︎ 探索用 Megatron-Core 训练框架提升大模型训练的显存效率, https://developer.nvidia.cn/blog/explore-using-the-megatron-core-training-framework-to-improve-gpu-memory-efficiency-in-large-model-training/ ↩︎ bindpcie 脚本按 local rank 自动探测 GPU/NUMA 拓扑,用 numactl 把每个进程的 CPU 和内存绑到离它的 GPU 最近的 NUMA 节点上, https://github.com/NVIDIA/mlperf-common/blob/main/client/bindpcie ↩︎ HybridFlow: A Flexible and Efficient RLHF Framework(veRL), https://arxiv.org/abs/2409.19256 ↩︎ slime: An LLM post-training framework for RL scaling, https://github.com/THUDM/slime ↩︎ Router Replay: Improving Sample Efficiency in Mixture-of-Experts Reinforcement Learning, 2026 ↩︎ Megatron-Bridge 性能脚本, https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/scripts/performance ↩︎ Megatron-MoE-ModelZoo 最佳实践, https://github.com/yanring/Megatron-MoE-ModelZoo/tree/main/best_practice/readme.md ↩︎--moe-latent-size 启用,它设定 $\ell$。MoELayer 实例化两个 TELinear 投影:fc1_latent_proj(preprocess 步骤里的 down-projection,在路由之后、dispatch 之前)和 fc2_latent_proj(postprocess 步骤里的 up-projection,在 combine 之后)。专家后端(TEGroupedMLP、SequentialMLP)自动把它们的输入输出维度适配到 $\ell$。7.4 分布式 Checkpoint
ShardedTensor descriptor,它编码每个局部 tensor 的全局形状、偏移和 sharding 模式。保存时每个 rank 独立写自己的局部分片(Fully Parallel Saving),消除了协调者瓶颈。加载时每个 rank 根据新的 sharding 规格判断自己需要全局 tensor 的哪些部分,只读那些切片。7.5 Flexible Asymmetric Virtual Pipeline Parallelism
PP rank
VPP rank 0
VPP rank 1
0
embedding + 3$\times$decoder
2$\times$decoder
1–13
2$\times$decoder
2$\times$decoder
14
2$\times$decoder
MTP
15
2$\times$decoder
loss
7.6 Upcycling
7.7 Multi-Token Prediction
7.8 Muon 优化器
8 性能评测
8.1 实验设置
8.2 关键性能结果
模型
系统
GPU 数
序列长度
Dtype
每 GPU TFLOPS
Tokens/s/GPU
DeepSeek-V3
GB300
256
4,096
MXFP8
1,233
4,730
DeepSeek-V3
GB200
256
4,096
MXFP8
1,048
4,020
DeepSeek-V3
GB200
256
4,096
BF16
857
3,298
DeepSeek-V3
H100
1,024
4,096
FP8-BLK
368
1,412
Qwen3-235B
GB300
256
4,096
MXFP8
974
6,583
Qwen3-235B
GB200
256
4,096
MXFP8
919
6,212
Qwen3-235B
GB200
256
4,096
BF16
750
5,100
Qwen3-235B
H100
256
4,096
BF16
320
2,132
Qwen3-235B
GB300
128
131,072
MXFP8
1,150
1,556
9 性能调优最佳实践
9.1 一套系统化的优化工作流
9.1.1 阶段 1:确立内存可行的并行方案
--use-distributed-optimizer)。
策略
峰值激活
权重内存
优化器状态
通信(每层)
TP
$1/d$(配 SP 时 $1/d$)
$1/d$
$1/d$
高
EP
约 1(取决于负载)
$1/d$(仅 MoE)
$1/d$
中
PP
1(配 VPP 时 $>1$)
$1/d$
$1/d$
中
CP
$1/d$
1
$1/d$†
中
DP
1
1
$1/d$†
低
--fake-init-process-group 在单张 GPU 上模拟分布式训练,无需申请整个集群就能快速迭代并行配置。用法细节见 Megatron-LM 文档110。9.1.2 阶段 2:选择最优并行策略
准则 1:最小化模型并行,最大化数据并行
--use-distributed-optimizer)把优化器状态切分到各 DP rank,为更大的 DP 腾出内存。准则 2:把 EP 和 TP 通信留在 NVLink 域内
准则 3:用 Pipeline Parallelism(PP)做多节点扩展
--pipeline-model-parallel-layout 控制 Pipeline Parallelism 设置(见 §7.5)。更大的 VPP 通常减少流水线气泡但增加 P2P 通信;取中间值往往能达到最佳平衡。各 VPP rank 的负载应当均衡以最大化吞吐。准则 4:专家层优先用 EP 而不是 TP
准则 5:长序列启用 Context Parallelism(CP)
--context-parallel-size 把序列切分到各 GPU。9.1.3 阶段 3:Profile 并优化瓶颈
内存瓶颈(内存墙)
优化
开销
配置
参见
FP8 训练
低
--fp8-format --fp8-recipe§4.1.3
选择性重计算
低
--recompute-granularity --recompute-modules§4.1.4
Precision-Aware 优化器
低
--use-precision-aware-optimizer§4.1.6
激活 offload
中
--fine-grained-activation-offloading --offload-modules§4.1.5
优化器 offload
中
--offload-optimizer-states§4.1.6
通信瓶颈(通信墙)
通信类型
配置
参见
DP 梯度 reduce 与参数 gather
--overlap-grad-reduce --overlap-param-gather—
TP 通信
--tp-comm-overlap—
EP dispatcher
--moe-token-dispatcher-type§4.2.2
EP all-to-all 隐藏
--overlap-moe-expert-parallel-comm§4.2.3
PP send/recv
--pipeline-model-parallel-layout§7.5
CPU 开销瓶颈(计算效率墙)
优化
配置
参见
关掉 Python GC
--manual-gc --manual-gc-interval 10—
减少 kernel 发起
降低 TP 或增大 MBS
—
启用 CUDA Graphs
--cuda-graph-impl transformer_engine§4.3.6
计算瓶颈(计算效率墙)
优化
配置
参见
Grouped GEMM
--moe-grouped-gemm§4.3.2
Kernel fusion
--moe-router-fusion --moe-permute-fusion§4.3.2
FP8 精度
--fp8-format --fp8-recipe§5
9.1.4 小结
9.2 案例研究:在 GB200 和 H100 上调优 DeepSeek-V3
9.2.1 最终优化配置与性能
配置项
GB200
H100
硬件
256$\times$GB200
1024$\times$H100
并行(TP/PP/EP)†
1/4/64
2/8/64
VPP
4
4
GBS / MBS / SeqLen
8192 / 1 / 4096
8192 / 1 / 4096
精度
MXFP8
FP8-Blockwise
Dispatcher
HybridEP
DeepEP
重计算
mlp
mlp, mla_up_proj, moe_act, layernorm
CUDA Graphs
启用
—
EP all-to-all Overlap
—
启用
性能(TFLOPS/GPU)
1048
368
类别
GB200
H100
并行
Parallel Folding;Flexible VPP
Parallel Folding;Flexible VPP
精度
MXFP8
FP8-Blockwise
内存
Memory-efficient permutation;细粒度重计算;FP8 primary weight;低精度优化器状态;优化器状态 offload
Memory-efficient permutation;细粒度重计算;FP8 primary weight;低精度优化器状态
通信
HybridEP
DeepEP;EP 通信 overlap
计算效率
Kernel fusion;CUDA Graphs;CPU 侧优化
Kernel fusion
9.2.2 优化配置的解剖
为什么是这个并行布局?
Et*3|(tt|)*29m|L 把 embedding 与 3 个 transformer 层放进第一个 stage、把 MTP 放进一个独立 stage、并把 loss 分开;GB200 上,Et*4|(tttt|)*14tmL 把层分布到 16 个虚拟 stage,以在更短的流水线上实现计算均衡。这些并行配置为后续所有优化奠定了基础。内存墙
通信墙
计算效率墙
9.2.3 经验教训
10 强化学习场景下的 Megatron-Core MoE
10.1 RL 后训练的挑战
10.2 Megatron-Bridge
10.3 面向强化学习的 Megatron-Core 优化
11 结论
附录 A 符号表
符号
说明
模型参数
$E$
一个 MoE 层里的专家数
$K$
Top-$k$ 路由宽度(每 token 激活的专家数;$K\ll E$)
$h$
Hidden 维度
$N$
模型总参数量
$N_{\text{active}}$
每 token 激活的参数量(随 $K$ 增长)
$N_{\text{total}}$
含所有专家的总参数量(随 $E$ 增长)
训练维度
$T$
每 GPU 的局部 token 数
$B$
Batch size(每 batch 的 token 数)
$s$
序列长度
$L$
MoE 层数
并行
TP
Tensor Parallelism 度
PP
Pipeline Parallelism 度
CP
Context Parallelism 度
DP
Data Parallelism 度
EP
Expert Parallelism 度
ETP
Expert Tensor Parallelism 度(MoE 专属)
EDP
Expert Data Parallelism 度(MoE 专属)
VPP
Virtual Pipeline Parallelism(交错的流水线 stage)
Batch 配置
MBS
Micro-batch size
GBS
Global batch size
GA
梯度累积步数
精度格式
BF16
BFloat16(16 位 brain floating point)
FP8-BLK
带 blockwise scaling 量化的 FP8(Hopper)
MXFP8
有原生 Tensor Core 支持的 Microscaling FP8(Blackwell)
指标与缩写
MFU
Model FLOP Utilization
GEMM
General Matrix Multiplication
SDPA
Scaled Dot-Product Attention
MLA
Multi-Latent Attention
MTP
Multi-Token Prediction
附录 B Benchmark 配置详情
B.1 配置详情
模型
配置
超参
DeepSeek-V3
GB300,256 GPU,4k seqlen,MXFP8,1233 TF
TP1 PP4 CP1 EP64 VPP4 MBS1 GBS8192
DeepSeek-V3
GB200,256 GPU,4k seqlen,MXFP8,1048 TF
TP1 PP4 CP1 EP64 VPP4 MBS1 GBS8192
DeepSeek-V3
GB200,256 GPU,4k seqlen,BF16,857 TF
TP1 PP8 CP1 EP32 VPP4 MBS1 GBS4096
DeepSeek-V3
H100,1,024 GPU,4k seqlen,FP8-BLK,368 TF
TP2 PP4 CP1 EP64 VPP4 MBS1 GBS8192
Qwen3-235B
GB300,256 GPU,4k seqlen,MXFP8,974 TF
TP1 PP4 CP1 EP64 VPP6 MBS2 GBS3072
Qwen3-235B
GB200,256 GPU,4k seqlen,MXFP8,919 TF
TP1 PP4 CP1 EP64 VPP6 MBS3 GBS3072
Qwen3-235B
GB200,256 GPU,4k seqlen,BF16,750 TF
TP1 PP4 CP1 EP32 VPP12 MBS1 GBS8192
Qwen3-235B
H100,256 GPU,4k seqlen,BF16,320 TF
TP2 PP8 CP1 EP32 VPP4 MBS1 GBS2048
Qwen3-235B
GB300,128 GPU,128k seqlen,MXFP8,1,150 TF
TP4 PP4 CP4 EP32 VPP12 MBS1 GBS1024
up_proj, mlp,Table 17 则列的是 mlp, mla_up_proj, moe_act, layernorm,另外 §B.2 说 H100 上 CUDA Graphs 关闭、1F1B overlap 打开,与 Table 17 一致。原文这几处彼此不一致,此处照原样保留、不做调和。另外原始 LaTeX 里 Qwen3 那几行写的是 VPP\phantom{0}6 和 VPP\phantom{0}4(\phantom{0} 只为与 VPP12 对齐),实际值就是 VPP6 和 VPP4。B.2 关键优化
B.3 可复现性
--fake-init-process-group 用法说明, https://github.com/NVIDIA/Megatron-LM/pull/2254 ↩︎
Author houmin
Publish July 30, 2026
LastMod July 31, 2026
License 本作品采用 CC BY-NC-ND 4.0 许可协议进行许可,转载时请注明原文链接
如果你在浏览博客的过程中发现了任何问题,欢迎在对应文章下评论。如果你有其他事情想要咨询,可以通过邮件联系我。