Megatron Core MoE Updates - 2025 H2

颜子杰, 陈楷文 | NVIDIA GPU加速计算专家团队 | Nov 07, 2025

目录

概览

该报告概述了 Megatron Core (MCore) 中混合专家模型 (MoE) 的更新，重点介绍了在 NVIDIA Hopper 和 Blackwell 架构上实现的 DeepSeek-V3 SOTA 训练性能。

核心性能指标：

  • H100: 370 TFLOPS
  • B200: 730 TFLOPS
  • GB200: 970 TFLOPS

下图展示了 MCore DeepSeek-V3 在不同硬件平台上的性能演进历史，从2025年初的170 TFLOPS逐步提升至2025年9月的970 TFLOPS。

MCore DeepSeek-V3 性能历史图
MCore DeepSeek-V3 性能历史图

实现这些性能的关键优化技术包括：
- 并行化 (Parallelism): 精细折叠 (Fine-grained Folding)
- 内存 (Memory): 精细重计算 (Fine-grained Recomputations)、内存高效置换 (Mem-Efficient Permutation)
- 通信 (Communication): DeepEP, HybridEP, 1F1B A2A 重叠
- CPU 开销 (CPU Overhead): 算子融合 (Fusions), cudaGraph
- FP8 (MXFP8/DeepSeek Recipe): 应用 FP8 精度进行训练
- 长上下文 (Long Context): 扩展至 512K 上下文长度

MCore MoE 开发分支 (Dev Branch)

为了支持新兴架构并实现最佳性能，MCore MoE 的开发在一个专用的开发分支上进行，其特点如下：
- 快速迭代与审查 (FAST ITERATION & REVIEW)
- SOTA MoE 特性 (SOTA MOE FEATURES)
- 全测试覆盖 (FULL TEST COVERAGE)
- 与主分支定期同步 (PERIODIC SYNC WITH MAIN)
- GitHub-First 开发模式 (GITHUB-FIRST DEVELOPMENT)
- 对社区贡献开放 (OPEN FOR CONTRIBUTIONS)

开发分支地址：https://github.com/NVIDIA/Megatron-LM/tree/dev

MCore MoE Dev Branch 开发理念
MCore MoE Dev Branch 开发理念

开发分支功能亮点

该开发分支已经合并了多项前沿功能：
- Qwen3-Next 和 GDN (Gated Delta Net) 支持 (已合并)
- 流水线并行下的激活卸载 (Activation Offloading) (已合并)
- Muon 和逐层分布 (layer-wise distribution) (已合并)
- 针对 NVL72 的 HybridEP (已合并)
- FSDP + EP 支持 (已合并)
- 重构和优化 MoE 的 cuda_graph_scope (已合并)
- MuonClip 支持
- 开发分支的混合上下文并行 (Hybrid-ContextParallel) 特性
- 融合的线性层和交叉熵 (Fused Linear and Cross Entropy) 操作

MCore MoE 路线图

以下是 MCore MoE 正在进行和未来的工作规划，涵盖模型支持、核心功能、并行化、优化、精度支持、通信和开发者体验等多个方面。

MCore MoE 路线图详情
MCore MoE 路线图详情

下一版本路线图 (MCore v0.16)

性能与内存增强
- [ ] 支持将 MTP 层放入独立的流水线阶段
- [ ] 融合线性和交叉熵操作

高级功能
- [ ] 增强 cuda_graph_scope 以支持 MoE 和 Mamba
- 更细粒度的图范围，如 MoE 路由器和调度预处理
- 微创式实现

  • [ ] MuonClip 支持（非拆分版本）
  • [ ] 为更简单的实现添加上下文并行支持

通信优化
- [ ] HybridEP MNNVL 优化
- [ ] HybridEP for NVL8+IB

Bug 修复
- [x] 修复 DeepSeek 和 Qwen HF tokenizer 的兼容性问题

持续进行的长期特性
- 针对 DeepSeek-V3, Qwen-3 等 fine-grained MoE 的端到端性能优化
- 无同步和全迭代 cudaGraph MoE 训练
- 设备发起的 HybridEP 和 GroupedGEMM
- MoE ECHO Dispatcher
- 针对 Blackwell 性能的 CPU 开销优化
- MLA CP 2.0 - 增强 MLA CP 以支持更长的序列训练
- 动态上下文并行以支持不平衡的长序列训练
- 探索 CPU 卸载 - Hybrid CPU-GPU expert routing
- Megatron FSDP MoE 训练的性能优化

性能分析

DeepSeek-V3 端到端性能 - Hopper

下表展示了在 Hopper 架构上，通过逐步应用各项优化技术，DeepSeek-V3 的性能提升过程。基线性能为 147 TFLOPS，最终通过 MTP Standalone 优化达到了 359 TFLOPS。

Hopper 平台 DeepSeek-V3 性能优化步骤及效果
Hopper 平台 DeepSeek-V3 性能优化步骤及效果

DeepSeek-V3 端到端性能 - Blackwell

下表详细列出了在 Blackwell 架构上，从基线 494.46 TFLOPS 到最终 970.01 TFLOPS 的各项优化特性及其带来的性能增益。关键优化包括 CUDA Graphs、HybridEP 和 Fuse quantization to normalization 等。

Blackwell 平台 DeepSeek-V3 性能优化特性列表及效果
Blackwell 平台 DeepSeek-V3 性能优化特性列表及效果

长上下文训练性能 - Hopper

下表展示了在 Hopper 平台上进行长上下文训练的性能数据，涵盖了不同 GPU 数量、序列长度和并行策略下的 TFLOPS 表现。

Hopper 平台长上下文训练性能表
Hopper 平台长上下文训练性能表

进行中和未来的工作重点

  • Qwen3-Next 和 DeepSeek-V3.2 模型支持
  • GB200 端到端性能交付
  • 无同步和全迭代的 cudaGraph MoE 训练
  • 设备发起的 HybridEP 和 GroupedGEMM
  • MoE ECHO Dispatcher
  • 全迭代 cudaGraph
  • 更多算子融合
  • 长上下文训练 (>512K)
  • HybridEP
  • 全面的 cudaGraph 支持
  • Moun 和其他新兴优化器
  • 利用 NVLINK-C2C 进行卸载

MCore MoE 优化之旅

MoE 并行折叠 (MoE Parallel Folding)

这是一个针对目标版本 v0.10 的技术。
- 核心思想: 解耦 MoE 层与 Attention 层的并行映射。允许在 Attention 和 MoE 层中使用不同的并行策略。
- 概念来源: MoE 并行折叠和 DLsim 中的异构并行 (Heterogeneous (UVXY) Parallelism) 本质上是相同的。我们默认将 MoE 中的 TPxEPxDP 从 Dense 和 TPxEPxDP 中折叠出来，并且没有为 CEP/TCP/TEP/CEP 提供显式设置。
- 接口设计:
- Attention 层使用传统的 TP/CP/DP。
- MoE 层可以指定任意的 TP/EP/DP。

  • 用例:

    • Attention: TP8CP8DP1

    • MoE: 可以使用任意并行映射，例如：

      • TP1EP64DP1 (之前的扩展 EP)
      • TP64EP1DP1 (之前的扩展 TP)
      • TP2EP16DP2
      • EP8TP8DP1 (带重排序)

    • 之前的扩展 TP 和扩展 EP 只是 EP 折叠的两种特例。

MoE 并行折叠示意图
MoE 并行折叠示意图

DeepSeek-V3 在 H100 上的内存优化

  • 基线使用完全重计算。内存优化旨在避免高开销的完全重计算。

  • 细粒度重计算 (Fine-Grained Recomputing):

    • 以低开销减少内存使用。
    • 使用 "mla_up_proj" + "layernorm" + "moe_act" 的重计算策略。
    • 这可以在仅增加 3% 端到端开销的情况下，节省 33% 的内存。

  • 内存高效置换 (Memory-Efficient Permutation):

    • 通过改变概率乘法的顺序，无需为反向传播保存 unpermutation 的输入。
    • 这可以在可忽略的开销下节省 20% 的内存。

  • 我们还使用 TP2 进一步减少内存使用。

H100 内存优化性能数据
H100 内存优化性能数据

内存优化：细粒度激活卸载 (Fine-grained Activation Offloading)

  • 在模块/算子粒度上卸载激活。
  • 与细粒度重计算协同工作，以尽可能减少总激活内存。
  • 支持 PP>1 & VPP>1。

  • 与最新的优化和模型结构兼容：

    • MXFP8/NVFP4
    • 流水线并行布局
    • 1F1B A2A 重叠

  • CUDA Graph (尚不能捕获卸载的模块)

  • 混合模型 (尚未发布)

下图展示了在 Grace-Blackwell Trays 上运行 DeepSeek-V3 的性能，通过卸载 moe_act，吞吐量基本保持不变，但最大内存占用显著降低。

细粒度激活卸载流程图及性能数据
细粒度激活卸载流程图及性能数据

DeepSeek-V3 在 H100 上的算子融合 (Fusions)

  • Kernel Fusions:

    • CE Loss fusion

    • MLA RoPE Fusion (用于前向和后向传播)

      • 前向: 955us -> 155us
      • 后向: 990us -> 150us

    • 我们还有其他融合，如置换融合 (permute fusion) 和路由融合 (router fusion)，与 Shopee 合作开发，但未包含在此次消融实验中。

下图展示了 MLA RoPE 融合前后的性能对比，以及融合操作对整体性能的提升。

H100 算子融合性能数据
H100 算子融合性能数据

DeepSeek-V3 on H100: DeepEP

DeepEP的优势包括：

  • 减少专家并行（EP）的通信开销。
  • 将全局排列（global permutation）与EP通信融合。

在使用Alltoall Dispatcher时，EP通信大约占据了50%的单步时间。DeepEP可以将EP通信时间大致减半。

下图展示了Alltoall Dispatcher和采用DeepEP的Flex Dispatcher的执行时间线对比，后者显著减少了通信（蓝色/紫色部分）的耗时。

Page 16
Page 16

性能数据

Step Time (s) Per GPU TFLOPs
DeepEP 32.7 261 (+54)

DeepSeek-V3 on H100: Blockwise FP8

  • 本节介绍DeepSeek的分块FP8（Blockwise FP8）方案及所有与FP8相关的优化。
  • FP8训练可以加速GEMM（通用矩阵乘法）计算并减小激活张量的大小。
  • FP8主权重（primary weight）和BF16优化器状态可以极大地减少内存使用。
  • 节省的内存对于后续的优化至关重要。

性能数据

Step Time (s) Per GPU TFLOPs
FP8 30.8 276
Page 17
Page 17

流水线并行优化：灵活的非对称流水线并行（PP）与自定义流水线布局

pipeline_model_parallel_layout 是一个用于定义流水线并行分区的灵活API，这对于非均衡模型的负载均衡至关重要。例如，对于DeepSeek-V3（61个解码器层 + 1个MTP层），我们可以通过设置 PP8VPP4 将其分区，具体命令如下：
--pipeline-model-parallel-layout="Et|(tt)*30|mL"

下图展示了如何将模型层（Embedding, Decoder, Loss）映射到不同的流水线阶段（Stage）以及不同的虚拟流水线并行（VPP）和物理流水线并行（PP）排名上，以实现负载均衡。

Page 18
Page 18

交叉流水线并行（Interleaved PP）与A2A（All-to-All）重叠

通过交叉执行不同微批次（micro-batch）的前向和后向计算，可以有效地将All-to-All通信操作与计算过程重叠，从而隐藏通信延迟。

  • 下图中的块大小是根据nsys报告中提取的芯片执行时间确定的。
  • 此工作与XHS合作完成。
Page 19
Page 19

DeepSeek-V3 on Hopper: EP A2A重叠与灵活VPP

  • 启用专家并行（EP）的All-to-All（A2A）重叠，以隐藏EP通信。
  • 使用灵活的虚拟流水线并行（VPP）以启用VPP4，从而减少流水线气泡（PP bubbles）。

下图的时间线展示了计算与通信的重叠情况，显著提升了GPU利用率。

Page 20
Page 20

性能数据

Step Time (s) Per GPU TFLOPs
EP Overlap 24.5 348 (+72)

<small>*仅供技术讨论和参考，性能可能因产品组合不同而异。</small>

DeepSeek-V3 on Hopper: MTP独立部署（MTP Standalone）

  • 启用MTP Standalone优化。
  • 在PP8VPP4配置下，总共有32个虚拟阶段。
  • MTP层包含了embedding、decoder和loss的计算。将MTP和loss放在一起会使最后一个阶段的计算负载过重。通过优化层的放置，可以使不同PP排名上的工作负载更加均衡。

下表对比了有无MTP Standalone优化时各阶段的层分配情况：
Stage allocation table from Page 21

通过将MTP和loss分离到不同的虚拟阶段，实现了更好的负载均衡。

性能数据

Step Time (s) Per GPU TFLOPs
MTP Standalone 23.8 359 (+11)
  • 在不强制平衡的真实数据下，性能为：单步时间24.8秒，吞吐量345 TFLOPs。

<small>*仅供技术讨论和参考，性能可能因产品组合不同而异。</small>

MCore中的MoE CUDA图（MoE CUDA Graph）

通过cuda_graph_scope可以对CUDA图生效的范围进行细粒度控制。
* 可控范围包括：[attn, mlp, moe, moe_router, moe_preprocess]

例如，设置--cuda-graph-scope attn mlp moe_router moe_preprocess可以捕获整个密集层（dense layer）以及MoE层中的attn+router+preprocess部分。

下图左侧是MoE层的流程图，标出了不同部分的scope；右侧是性能分析器（profiler）的截图，显示了被CUDA图捕获（Graphed）的区域。

Page 22
Page 22

DeepSeek-V3性能最佳实践：一份配置清单

以下是一组推荐的命令行参数配置，用于实现DeepSeek-V3的最佳性能。

Page 23
Page 23
参数 功能
--moe-token-dispatcher-type flex --moe-enable-deepep 启用DeepEP
--moe-router-dtype fp32 --moe-permute-fusion --moe-grouped-gemm --moe-router-fusion 启用MoE相关优化
--pipeline-model-parallel-layout "Et*3|(tt)*29m|L" 使用灵活的PP布局
--recompute-granularity selective --recompute-modules mla_up_proj mlp moe_act 启用细粒度的重计算
--no-rope-fusion: false 启用MLA Rope Fusion
--cross-entropy-loss-fusion --cross-entropy-fusion-impl native 启用交叉熵损失函数融合
--manual-gc --manual-gc-interval 10 启用手动垃圾回收以避免迭代内的意外GC
--use-precision-aware-optimizer --exp-avg-dtype bf16 --exp-avg-sq-dtype bf16 启用BF16优化器状态
--overlap-moe-expert-parallel-comm --delay-wgrad-compute 启用1F1B重叠
--fp8-recipe blockwise --fp8-format e4m3 启用DeepSeek FP8方案
--moe-router-padding-for-fp8 启用MoE Router Padding优化
--fp8-param-gather 启用FP8主权重
--external-cuda-graph --cuda-graph-scope attn mlp moe_router moe_preprocess 启用CudaGraph（可选）

MCore MoE模型库 (Model Zoo)

为了方便客户复现我们的性能，我们提供了一个MCore MoE模型库。

  • 包含性能复现所需的一切。
  • 为典型的MoE模型（如DeepSeek, Qwen, Mixtral）提供预定义的端到端训练脚本。
  • 预定义的模型参数。
  • 经过调优的并行化和性能参数。
  • 提供Dockerfiles。
  • 链接: https://github.com/yanring/Megatron-MoE-ModelZoo
Page 24
Page 24

Megatron-FSDP

Page 25
Page 25

完全分片数据并行（FSDP）

为什么选择FSDP？

  • 与3D并行相比，使用更简单。
  • 与流水线并行（PP）及其衍生技术更兼容。
  • 性能表现优异。

下图展示了在Blackwell平台上训练LM3 405B模型时，不同并行策略的性能对比。Megatron-FSDP (M-FSDP) 取得了最高的2020 TFLOPS/GPU。

Page 26
Page 26

Megatron-FSDP分片结构

FSDP2 vs Megatron-FSDP

  • Torch FSDP2 均匀地对参数进行分片。
  • Megatron-FSDP 非均匀地对模块进行分片。
  • 两者都利用可扩展的按模块/按组的通信。
  • FSDP2需要将按参数的分片（per-parameter shards）置换到按模块的通信缓冲区中，这在LM3 405B训练中导致了约10%的性能开销。

下图对比了两种分片方式的结构差异。

Page 27
Page 27

Megatron-FSDP缓冲区管理

FSDP双缓冲区（Double-Buffer）

  • 基线FSDP：重复创建集合通信缓冲区。

    • 当剩余内存较少时，需要进行内存分配/释放，这会降低性能并增加内存碎片。
    • 会产生CPU开销。

  • Megatron-FSDP：使用持久化的通信缓冲区。

    • 创建持久的双缓冲区，并在集合通信操作之间重用。
    • 在Blackwell平台上训练LM3 405B时，带来了1.5%的速度提升。

下图展示了Megatron-FSDP中使用预分配双缓冲区的流程。

Page 28
Page 28

用户缓冲区注册（User-Buffer-Registration）NCCL通信

  • 基线集合通信核心（collective kernels）的缺点：

    • 占用8-32个流式多处理器（SMs），在与计算重叠时会拖慢计算核心。
    • 在核心和用户内存之间双向复制数据会带来额外的延迟。

  • 注册持久化NCCL用户缓冲区的好处：

    • 启用零拷贝（zero-copy），消除拷贝延迟。
    • 能够利用NVIDIA NVLink和InfiniBand的特性。
    • 多播（Multi-cast）仅需1-4个SM，从而加速计算核心。
    • SHARP将归约（reduction）操作从GPU SM卸载到网络交换机。

下图展示了该技术所利用的硬件架构，包括GPU、NVSwitch和IB-Switch。

Page 29
Page 29

计算与通信的重叠

调整并行配置以隐藏通信

  • 在计算中隐藏通信是FSDP性能的关键。
  • 扩展计算量（例如，通过增加微批次大小）会导致更高的内存使用，但可以实现更好的重叠。

下图对比了基线反向传播和完全隐藏通信的两种情况。通过将微批次大小加倍（2X micro-batch size），计算时间足以完全覆盖通信时间，从而消除了暴露的通信开销。

Page 30
Page 30

Megatron-FSDP 分布式环境

Megatron-FSDP 的并行性由 DeviceMesh 定义和管理。

  • Megatron-FSDP 的分布式状态与5个并行维度进行交互：

    • DP-Outer (节点间/Inter-Node)
    • DP-Shard (节点内/Intra-Node)
    • CP (上下文并行/Context Parallel)
    • TP (张量并行/Tensor Parallel)
    • EP (专家并行/Expert Parallel)，通过一个辅助的 "EDTP" DeviceMesh 支持。

  • Megatron-FSDP 解耦了非 EP 模块和 EP 模块的分布式状态：

    • 对非 EP 模块强制使用 DeviceMesh。
    • 对 EP 模块使用辅助的 Expert DeviceMesh。

  • Megatron-FSDP 与基于 Megatron 的并行性兼容。

Page 31
Page 31

Megatron-FSDP 检查点 (Checkpointing)

支持 Torch 分布式检查点 (DCP)

  • 重构了检查点实现，以使用基于 DTensor 的 torch 分布式检查点。
  • 支持将 3D 并行检查点转换为 Megatron-FSDP 检查点。

在 DeepSeek Proxy 模型上进行检查点重载测试

下图显示，从 Megatron-FSDP 和 3D 并行模式的检查点重新加载后，训练曲线与基线完全一致，证明了检查点加载的正确性和无缝衔接。

Page 32
Page 32

Megatron-FSDP 基准测试

Blackwell 平台

下表展示了在 Blackwell 平台上对 DeepSeek-V3 模型的基准测试结果。

Page 33
Page 33

  • 启用 Megatron-FSDP 的标志 (flags)：

    • --use-megatron-fsdp
    • --data-parallel-sharding-strategy optim_grads_optim
    • --no-gradient-accumulation-fusion
    • --use-distributed-optimizer
    • --calculate-per-token-loss
    • --init-model-with-meta-device
    • --ckpt-format fsdp_dtensor

  • 更高性能的选项：

    • --grad-reduce-in-bf16
    • --fsdp-double-buffer
    • --use-nccl-ub

  • 未来可能的优化方向：

    • HybridEP (混合专家并行)
    • Activation offloading (激活卸载)
    • ...

PyTorch Conference 2025 演讲

该工作已在 PyTorch Conference 2025 上展示。

Page 34
Page 34