MCore MoE in 2025 - DeepSeek-V3 and Beyond

Zijie Yan and Hongbin Liu
NVIDIA

目录

• 亮点总结

• 模型架构与挑战

  • 模型架构
  • DeepSeek-V3 训练的挑战：功能性
  • DeepSeek-V3 训练的挑战：内存
  • DeepSeek-V3 训练的挑战：效率

• MCore DeepSeek-V3 优化详解

  • 功能完备性
  • 通信优化
  • 内存优化
  • 计算优化
  • 流水线并行优化

• 性能与实践

  • DeepSeek-V3 端到端性能
  • DeepSeek-V3 性能最佳实践
  • MCore MoE 模型库

• 专题：在 1F1B 调度中重叠 MoE AllToAll 通信

  • 背景
  • 两种重叠策略
  • 合并 FWD-BWD 与 AllToAll 重叠
  • 交叉流水线并行（Interleaved PP）与 AllToAll 重叠
  • 基准测试结果
  • 性能分析
  • 当前状态

• 总结与未来展望

  • 总结
  • MCore MoE 路线图：正在进行和未来的工作

亮点总结

• 整体状态: MCore MoE现已完全支持DeepSeekV3训练，包括FP8，是开源社区中唯一实现此功能的框架。使用FP8优化后，实现了349 TFLOPS的性能。
• 功能性 (✓)

• 内存优化 (✓):

  • 内存高效的排列（Permutation）
  • 细粒度的重计算

• 通信优化:

  • DeepEP集成 (✓)
  • FP8 A2A (开发中, WIP)
  • 带自定义流水线布局的灵活非对称流水线并行 (✓)
  • 1F1B A2A 重叠 (开发中，准备早期试用)
  • 上下文并行 (✓)

• 计算优化:

  • 一系列算子融合 (✓)
  • cuDNN MLA 优化 (✓)

• 提供端到端的可复现示例

模型架构与挑战

模型架构

幻灯片对比了三个大型语言模型：LLAMA 70B，Mixtral-8x7B和DeepSeek-V3。

• LLAMA 70B:

  • 模型架构: 隐藏层维度8192，80层。
  • 基础设施/框架: TP + PP + Zero1。

• Mixtral-8x7B 47B-A13B:

  • 模型架构: 隐藏层维度4096，32层，Top-2门控，8个专家。
  • 基础设施/框架: EP8 + PP8。

• DeepSeek-V3 671B-A37B:

  • 模型架构: 隐藏层维度7168，61层，Top-8门控，细粒度MoE（256个专家/专家很小）。
  • 基础设施/框架: TP=1（专家非常小，TP非必需），EP 64。A2A（All-to-All）通信占主导，引入了DualPipe重叠和DeepEP。EP通信成为性能瓶颈。PP = 8/16。

关键点:

• Mixtral-8x7B: 第一个开源的MoE模型，约47B参数。
• DeepSeek-V3: 细粒度的MoE模型，含671B参数。
• EP A2A通信（带重叠）:
  • Mixtral-8x7B: NVLink内部的A2A，端到端（E2E）通信开销约3-5%。
  • DeepSeek-V3: 跨NVL和IB的A2A，端到端通信开销约50%。

DeepSeek-V3 (DSV3) 模型超参数:

DeepSeek-V3 训练的挑战：功能性

DeepSeek-V3模型引入了多项改变，均已在MCore中得到支持：

• 多令牌预测 (Multi-Token Prediction, MTP)
• 多层级注意力 (Multi-latent Attention, MLA)
• 混合模型架构 (前3层为密集层)
• 节点限制路由 (Node-limited routing)
• 带专家偏置和序列辅助损失的负载均衡
• 非均衡模型架构 (前3层为密集层，最后一层为MTP)

DeepSeek-V3 训练的挑战：内存

• 静态内存与激活内存:

  • 总静态内存: DeepSeek-V3 (11,445 GB) 远超 LLAMA 70B (1,157 GB) 和 Mixtral 8x7B (783 GB)。
  • 激活内存: DeepSeek-V3 (118 GB) 同样远高于 LLAMA 70B (21 GB) 和 Mixtral 8x7B (30 GB)。

• 内存瓶颈: 即使在1024个GPU上使用3D并行，DSV3训练时每个GPU的激活内存仍然超过80GB。

• 基准测试所需资源:
  • LLAMA 70B: 64x80GB
  • Mixtral 8x7B: 64x80GB
  • DeepSeek-V3: 1024x80GB

DeepSeek-V3 训练的挑战：效率

• 密集的通信:

  • DSV3的跨节点通信量达到每次迭代10TB，而之前的模型仅为100GB级别。
  • 在旧模型中，通常只有数据并行（Data Parallel）和流水线并行（Pipeline Parallel）需要跨节点通信。而DSV3的大规模专家并行（EP）需要大量的跨节点通信。

• 稀疏计算:

  • DeepSeek V3从256个专家中激活8个（1/32），这使得计算高度稀疏。
  • 更高的稀疏度通常意味着更小的GEMM（通用矩阵乘法）形状和更低的GEMM效率。
  • 相比之下，Mixtral激活8个专家中的2个（1/4），而Llama3是密集模型（1/1）。

MCore DeepSeek-V3 优化详解

功能完备性

使用MTP进行收敛性测试

• 在MCore v0.12中，已为使用MTP的DeepSeek-V3训练做好了全面的功能准备。
• 新增功能:
  • 多令牌预测 (MTP): 在单次前向传播中预测多个未来令牌，有助于模型收敛。
  • 节点限制路由: 指定每个令牌被发送到的节点数量，从而减少通信量。
  • 多层级注意力 (MLA): 减少推理过程中的KV缓存。

DeepSeek-V3 微调的损失曲线:

通信优化

DeepEP

• 传统令牌分发: 结合NCCL通信和置换操作，节点间发送的令牌存在重复。
• DeepEP:
  • 将置换和通信融合在单个内核中。
  • 通过两阶段通信，对令牌通信执行节点级的去重。
  • NVLink和IB通信可以重叠进行。
  • 在DeepEP发布后，MCore实现了Day-0支持。

上图展示了从传统的AlltoAll令牌分发器到简化的DeepEP分发器的演进，分为两个阶段：跨节点通信和节点内通信。右侧表格数据显示，DeepEP显著降低了不同EP规模下的MoE层延迟。

MoE Parallel Folding

• 核心思想：将 MoE 层的并行映射与密集层的并行映射解耦。

  • Attention 和 MLP 是两种具有不同计算和通信特性的工作负载。并行折叠（Parallel folding）允许它们各自拥有不同的、最优化的并行配置。

• 通信组布局:

  • 密集层 (Dense Layer) 的并行维度可以表示为：[PP x [ DP x CP x TP ]]
  • MoE 层 (MoE Layer) 的并行维度则为：[PP x [ ExpertDP x EP x ExpertTP* ]]
    • 其中，并行维度被划分为节点间（Inter Node）和节点内（Intra Node）通信。
    • 注：在 MCore v0.10 之后，可以设置任意的 ExpertTP 大小，默认情况下会继承 TP 的大小。

• 接口设计：

  • Attention 层使用传统的张量并行（TP）/上下文并行（CP）/数据并行（DP）映射。
  • MoE 层可以指定任意的张量并行（TP）/专家并行（EP）/数据并行（DP）映射。

• DSV3（DeepSeek-V3）用例：

  • Attention 层使用 TP2DP64。
  • MoE 层使用 TP1EP64DP2。

• 图示：下图展示了在4个Rank上，Attention层采用4路张量并行（TP4），而MoE层则采用4路专家并行（EP4），体现了并行策略的解耦。

内存优化

细粒度重计算

• 新的输出丢弃式检查点 (output-discarding checkpointing):

  • 传统的标准检查点方法需要保存函数的输入和输出。
  • 新方法可以节省被设置检查点的函数的输出内存。

• 应用范围:

  • MLA上投影线性和RoPE应用。
  • MoE SwiGLU激活函数（目前为BF16）。
  • 注意力层和MoE层之前的Pre RMSNorm模块（目前为BF16）。

• 效果: 这些模块的重计算可以显著减少内存使用，且对性能影响有限。

• 统一接口: 所有细粒度重计算可通过以下命令行参数启用：
  • --recompute-granularity selective --recompute-modules [submodules1, submodules2, ...]
  • 可选模块包括: [core_attn, mlp, moe, moe_act, layernorm, mla_up_proj]

内存高效的排列

• 问题: 在MoE层中，中间张量的大小与路由器的topk值成线性关系。之前，需要保存输入张量以用于反向传播中的unpermute操作。对于DeepSeek模型（TopK=8），这个被保存的张量占用了大约25 GB。
• 解决方案:
  1. 重写计算流程，将probs（概率）的乘法操作放入MLP的激活函数中。
  2. 这样，用于unpermute的已保存张量可以被释放，并在反向传播中重新计算。
  3. 数值与原始计算流程保持一致。

计算优化

核函数优化

• MLA Yarn RoPE 融合
• 融合的核心注意力 (Fused Core Attention)
• 路由器融合 (Router Fusion)
• 交叉熵融合 (Cross Entropy fusion)
• 用于DeepEP分发器的 indices_to_multi_hot 融合

MLA Yarn RoPE融合

• 性能提升:

  • 与原生PyTorch实现相比，速度提升了3-15倍。
  • 在DSV3中，端到端性能提升约10%。

• 格式支持: 支持SBHD和THD格式。

上图右侧展示了Q和KV路径上的RoPE操作被融合成一个单一的核函数。左侧表格量化了融合前后的延迟，显示出显著的加速比。

FusedAttention

• 挑战: MLA中的头维度（qk_dim != v_dim）与通用多头注意力（MHA，qk_dim = v_dim）不同。在cuDNN 9.7中，其性能比DeepSeek的内部核函数要慢。
• 合作成果: 与cuDNN团队合作后，获得了比DeepSeek原有实现更好的性能。

上表比较了DeepSeek的性能分析数据、cuDNN 9.7和cuDNN 9.11.0.28在Fused Attention上的延迟。结果显示，最新版本的cuDNN性能最佳，实现了1.03倍至1.07倍的性能提升。

流水线并行优化

带有自定义流水线布局的灵活非对称流水线并行

• pipeline_model_parallel_layout 是一个灵活的 API，用于定义流水线并行分区，这对于非均衡模型的负载均衡至关重要。
• 示例：对于 DeepSeek-V3（61个解码器层 + 1个 MTP 层）这样层数不均衡的模型，在使用 PP8VPP4（8路流水线并行，4路虚拟流水线并行）时，可以设置如下布局：
  --pipeline-model-parallel-layout="Et*3|(tt|)*29m|L"

• 布局解析：

  • Stage 0: embedding + 3个 decoder 层
  • Stage 1-29: 每个Stage包含2个 decoder 层
  • Stage 30: mtp 层
  • Stage 31: loss 计算

• 该图表展示了如何将不同数量的层非对称地分配到不同的流水线阶段和GPU上，以实现负载均衡。

性能与实践

DeepSeek-V3 端到端性能

以下表格展示了在H100 GPU上，通过一系列优化逐步提升DeepSeek-V3训练性能的过程。

注：基准测试中临时使用了强制路由；仅供技术讨论。

DeepSeek-V3 性能最佳实践

这张幻灯片提供了一系列用于优化 DeepSeek-V3 训练性能的推荐参数配置。

MCore MoE 模型库

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

• 特点：

  • 包含性能复现所需的所有内容。
  • 为典型的 MoE 模型提供了预定义的端到端（E2E）训练脚本。
  • 预定义了包括 DeepSeek、Qwen 和 Mixtral 在内的模型参数。
  • 包含了调优后的并行化和性能参数。
  • 提供 Dockerfiles。
  • 提供从 HuggingFace (HF) 模型到 MCore 格式的转换脚本。

• 链接: https://github.com/yanring/Megatron-MoE-ModelZoo

专题：在 1F1B 调度中重叠 MoE AllToAll 通信

背景

• 问题: 专家并行（EP）中的 AllToAll 通信耗时较长，尤其是在细粒度 MoE 模型中。因此，需要通过与计算重叠来隐藏这部分通信开销。

• 现有方法:

  1. 操作级重叠 (OP-level overlap): 将 AllToAll 与专家 MLP 计算重叠。

     • 例如，在 Mixtral-8x7B 中重叠 a2a 和 FC1 等2到3个连续操作。
     • 缺点: 要求计算时间必须大于或等于被重叠的通信时间。

  2. DualPipe (批次级重叠, Batch-level overlap): 重叠不同微批次（micro-batches）的计算和通信。

     • 缺点: 调度策略过于复杂，需要在 MCore 中进行大规模代码重构。

• 模型与策略: 下图流程展示了不同模型（如 Deepseek-v3, snowflake arctic, GPT4/Grok-2, Mixtral-8x7B）适用的不同重叠策略。

两种重叠策略

• 策略一（上图）:

  • 关键特性: 峰值激活内存占用为2倍。合并的前向传播（FWD-FWD）可能为 AllToAll 重叠留下的空间较小。

• 策略二（下图，合并的 FWD-BWD）:

  • 关键特性:

    • 合并的前向-后向传播（Merged FWD-BWD）等效于 DualPipe。
    • 无额外的内存开销。
    • 无法对单个前向和后向传播（1 FWD & 1 BWD）应用 AllToAll 隐藏。

  • 该策略因其内存效率而被采纳（由绿色对勾标记）。

合并 FWD-BWD 与 AllToAll 重叠

• 基本实现: 将前向传播（fprop）和后向传播（bprop）的计算和通信进行拆分，并交替放入计算流（Compute stream）和通信流（Comm stream）。
• W/D 拆分: 将梯度计算（wgrad）和数据梯度（dgrad）拆分，以打破后向传播（B/dispatch）对后向 MLP（B/mlp）的依赖，从而减少计算流中的空闲气泡（bubbles）。
• 图示: 下图对比了基线方案和采用 MLP W/D 拆分（等同于 DSV3）方案的执行流。后者通过更细粒度的任务划分，实现了计算和通信更好的重叠，减少了流水线停顿。

交叉流水线并行（Interleaved PP）与 AllToAll 重叠

• 设置: 流水线并行大小=4，虚拟流水线并行大小=3，梯度累积=8。
• 核心思想: 在 1F1B（one forward, one backward）调度开始前，额外执行一个微批次（micro-batch）的前向传播。
• 效果: 相邻微批次的前向传播（Fprop）和后向传播（bprop）之间没有依赖关系，从而可以实现更好的流水线调度和重叠。

基准测试结果

PP=1

• 设置:

  • 8层 MoE，无 MTP，约等于 DSv3 模型的1/8。
  • 配置为 TP2EP64PP1_mbs1gbs2048，在 128 个 H100 GPU上运行。

• 性能数据:

  

• 结论:

  • AllToAll 重叠带来了 1.21 倍的加速。
  • AllToAll 重叠不会引入额外的内存分配。
  • AllToAll 重叠在一定程度上会增加 PyTorch 的缓存，原因可能是某个流上的内存段无法被另一个流重用，但这尚未完全验证。

PP=1, 内存分配

• 基线情况 (左图): 在基线中，前向传播的内存分配之后是后向传播的内存释放，因此内存使用模式呈现出明显的“峰谷”状。
• AllToAll 重叠情况 (右图): 在 AllToAll 重叠的情况下，前向分配与后向释放并行发生。这不会导致新的内存分配，而是使内存得到更充分的利用，内存使用模式更加平稳和密集。

DeepSeek-V3-no-MTP

• 设置:

  • DeepSeek-V3 模型，无 MTP。
  • 配置为 TP2EP64PP8VPP4_mbs1gbs8192，在 1024 个 H100 GPU上运行。
  • 未使用 fp8 AllToAll。

• 结果:

  • 实现了 20% 的端到端（e2e）加速，与在代理模型上的先前结果和预测相符。
  • 在 FP8 训练中获得了更高的加速比，这是因为 AllToAll 通信在总时间中的占比较大。

• 图表: A2A 重叠在 DeepSeek-V3-no-MTP 模型上的加速效果。

  • bf16: 加速比为 1.22。
  • fp8: 加速比为 1.29。

性能分析

1F1B中的AllToAll重叠

该图示展示了在1F1B流水线调度中，AllToAll通信与计算的重叠情况。时间线被分为四个部分：

1. Bubble Not overlapped (气泡，未重叠)
2. Overlapped (重叠部分)
3. Not overlapped (未重叠部分)
4. Bubble (气泡)

延迟计算公式如下：

• 气泡 (Bubble): (p-1)/v * (F + B)
• 最后一个流水线并行(PP) rank中的未重叠部分: (p(v-1)+1)/v * (F + B)
• 最后一个PP rank中的重叠部分: (mv-(p(v-1)+1))/v * (F&B)
• 总延迟 (Total latency): (p-1)/v * (F + B) + (p(v-1)+1)/v * (F + B) + (mv-(p(v-1)+1))/v * (F&B) = p(F + B) + (m - p)(F&B) + (p-1)/v * (F&B)

其中：
* p: 流水线并行大小 (pipeline parallel size)
* v: 虚拟流水线并行大小 (virtual pipeline parallel size)
* F: 流水线并行rank的每个微批次的前向传递时间 (forward time per micro batch of one pipeline parallel rank)
* B: 流水线并行rank的每个微批次的后向传递时间 (backward time per micro batch of one pipeline parallel rank)
* m: 微批次数量 (number of micro batches)

DualPipeV中的AllToAll重叠

该图示分析了DualPipeV流水线调度中的AllToAll重叠。时间线跨越4个设备，包含以下阶段：
* 预热阶段 (Warmup stage)
* 关键路径上的重叠部分 (Overlapped in critical path)
* 刷新阶段 (Flush stage)

延迟计算公式如下：

• 前向气泡 (Bubble in forward): (p-1)/2 * F
• 最后一个rank的预热阶段 (Warmup stage of last rank): (p + p/2) * F + B/2
• 关键路径上的重叠部分 (Overlapped part in the critical path): (m - (p + p/2)) * (F&B) = (m - (p/2 + 1)) * (F&B)
• 刷新阶段 (Flush stage): B/2 + (p/2) * (B-W) + W/2 = (p+1)/2 * B 如果 W = B/2
• 总延迟 (Total latency): (p-1)/2 * F + (p + p/2) * F + (p+1)/2 * B + (m - (p/2 + 1)) * (F&B) = (3p)/2 * F + (p+1)/2 * B + (m - (p/2 + 1)) * (F&B)

其中：
* p: 流水线并行大小 (pipeline parallel size)
* v: 虚拟流水线并行大小 (virtual pipeline parallel size)
* F: 流水线并行rank的每个微批次的前向传递时间
* B: 流水线并行rank的每个微批次的后向传递时间
* W: 流水线并行rank的每个微批次的wgrad时间 (wgrad time per micro batch of one pipeline parallel rank)
* m: 微批次数量 (number of micro batches)

1F1B与DualPipeV总延迟比较

• 1F1B-总时间: p(F+B) + (m-p)(F&B) + (p-1)/v * (F&B)
• DualPipeV-总时间: (3p)/2 * F + (p+1)/2 * B + (m - (p/2 + 1)) * (F&B)

假设没有a2a重叠，F&B=F+B, 且 B=2F
* T_1f1b = m(F+B) + (p-1)/v * (F+B) = 3mF + 3/v * (p-1)F
* T_dualpipe = m(F+B) + (p-1)F = 3mF + (p-1)F
* 当 v>3 时，1F1B比DualPipeV更快。

假设F被B完全隐藏，F&B=B, 且 B=2F
* T_dualpipe - T_1f1b = (p/2 - (2(p-1))/v) * F
* 当 v=2, T_dualpipe - T_1f1b = (1 - p/2) * F < 0, 当 p > 2
* 当 v=4, T_dualpipe - T_1f1b = (p/2 - (p-1)/2) * F > 0
* 当 v=8, T_dualpipe - T_1f1b = (p/2 - (p-1)/4) * F > 0

• 当 v=4, T_dualpipe/T_1f1b = (F/2)/( (3p)/2 * F + m*2F) = 1 / (3p+4m)，当p和m较大时，延迟差异很小。

1F1B与DualPipeV中AllToAll重叠的比较：
* 对于较大的虚拟流水线并行大小，1F1B速度更快。
* 在PP size较大和微批次数量较多的情况下，两种策略之间的差距很小。
* 对于像DeepSeek-V3这样的混合模型，我们需要仔细设计以平衡各个PP阶段的工作负载。

AllToAll重叠带来的端到端加速的限制

端到端加速 (Speedup) 公式:
Speedup = θ_1f1b-batches * θ_a2a-in-1f1b * γ_gemm-eff * min(t_comp, t_a2a)

公式各部分解读：

• 重叠批次的比例 (Proportion of overlapped batches):

  • 公式: (num_micro_batches - pp_size) / num_micro_batches
  • 适用于PP=1和交错式PP (interleaved PP)。
  • 微批次越多越好。

• 1f1b阶段中AllToAll通信的比例 (Proportion of AllToAll comm in an 1f1b stage, prominent at):

  • 节点间AllToAll通信 (Inter-node AllToAll)。
  • 稀疏模型/细粒度MoE模型 (Sparse models/Fine-grained MoE models)。

• SM carveout对gemm效率的影响 (Impact of SM carveout on gemm efficiency):

  • 在DSv3训练中，DeepEP占用20个SM。
  • 对gemm效率造成约20%的开销。
  • SM高效版的DeepEP正在开发中（WIP），目标是使用4-8个SM。

• 暴露的AllToAll通信时间 (Exposed AllToAll communication time):

  • 超过93%的all2all可以被很好地重叠。
  • 剩余部分也可以被优化，但在P1阶段。

注：详细数据来自对DeepSeek-V3代理模型的性能分析。

当前状态

• EA (Early Access) 代码:

  • https://github.com/hb8125/Megatron-LM/tree/hongbin/1f1b_overlap_ea?tab=readme-ov-file#1f1b-with-a2a-overlap
  • 注意：这是预览版本，可能与最终版本有显著差异。

• 特性 (Features):

  • 隐藏AllToAll通信，支持批处理级别重叠 (batch-level overlapping)
  • 与交错式1F1B流水线并行兼容 (interleaved 1F1B pipeline parallel)
  • 支持DeepEP
  • 支持FP8训练
  • 支持MTP
  • 支持dW&dX拆分以实现更好的重叠 (dW&dX split for better overlapping)

• 进行中的工作 (WIP):

  • 代码审查 (Code Review)
  • 性能调优 (Performance tuning)

• 贡献者 (Contributors):

  • NVIDIA: Hongbin, Pingtian, Youngeun, Shunkang, Sangkug, Zijie, Gao
  • Xiaohongshu: Zhenhai

• 推荐配置:

  • 以下环境变量有助于提升性能：

CUDA_DEVICE_MAX_CONNECTIONS: 32
TORCH_NCCL_AVOID_RECORD_STREAMS: 1
NVTE_ALLOW_NONDETERMINISTIC_ALGO: 1
PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
NCCL_NVLS_ENABLE: 0
NVTE_FWD_LAYERNORM_SM_MARGIN: 8
NVTE_BWD_LAYERNORM_SM_MARGIN: 8 # 24 for DeepEP

* 将以下标志添加到您的训练脚本中：

--combined-1f1b
--combined-1f1b-recipe ep,a2a
#[optional] only works with TE v2.3.0
--delay-wgrad-compute

总结与未来展望

总结

• 批处理级别的重叠可以在不增加明显内存占用的情况下，很好地隐藏MoE的AllToAll通信，该功能目前已在MCore EA repo中可用。
• 在MCore中，MoE AllToAll重叠可以为类似DeepSeek-V3的模型带来超过20%的端到端（E2E）加速。
• 在MCore中，MoE AllToAll重叠可以实现与DualPipeV相当的延迟。
• 大量的微批次数量以及AllToAll通信与计算的比例可以最大化加速效果。

MCore MoE 路线图：正在进行和未来的工作

• 将当前的优化进行产品化。
• 最小化 FP8 的 CPU 开销。
• FP8 的内存优化。
• 细粒度的重叠（Overlapping）。
• 支持长序列（Long Sequence）。
• 为强化学习（RL）训练引入 DSV3。