DeepSeek-V3 Pre-training Optimization on Grace Blackwell

姚鑫 | NVIDIA GPU加速计算专家团队 高级工程师 | NVIDIA AI Open Day, Nov. 7th, 2025

目录

• 软件栈 (Software Stack
• 方法论 (Methodology
• 基础设置 (Basic Settings
• 基线性能剖析 (Baseline Profile
• 优化项 (Optimization Items
• 1. 核函数融合与优化
• 2. 内存节省
• 3. CUDA Graphs
• 4. CPU 优化
• 5. HybridEP
• 优化历程总结 (Optimization Journey
• 讨论
• 未来工作
• 致谢

软件栈 (Software Stack)

该方案的软件栈构建于 Transformer Layer 之上，主要由 Megatron-Core 和 Transformer Engine 两部分组成。

Megatron-Core
* 模型定义
* 模型并行化
* 流水线并行调度 (Pipeline parallel schedule)
* EP alltoall 重叠 (EP alltoall overlap)
* MoE 并行折叠 (MoE Parallel folding)

• 优化器 (Optimizer)

  • 分布式优化器 (ZeRO-1)
  • FSDP
  • 精度感知优化器 (Precision-aware optimizer)
  • Muon 优化器

• 检查点 (Checkpoint)

  • 分布式检查点 (Distributed checkpoint)

Transformer Engine

• 低精度 (Low-precision)

  • 逐张量 FP8, 逐块 FP8, MXFP8 (Per-tensor FP8, blockwise FP8, MXFP8)
  • NVFP4

• GEMM 实现

  • 低精度 GEMMs
  • 分组 GEMM (Grouped GEMM)
  • 多流 cuBLAS (multi-stream cuBLAS)
  • CUTLASS
  • cuBLASLT 原生分组 GEMM (即将推出)

• 核函数优化 (Kernel optimizations)

  • Norm + quant 融合
  • Activation + quant 融合
  • Router/Permute 融合

• Attention 后端实现

  • cuDNN: 针对 Hopper 和 Blackwell 进行了优化
  • FA2 和 FA3

方法论 (Methodology)

性能优化的过程遵循一个迭代循环：

1. 运行实验: 在当前软件栈下，使用最佳配置运行实验。
2. 性能剖析与瓶颈定位: 对运行过程进行性能剖析，找出瓶颈所在。
3. 性能优化: 针对找到的瓶颈进行优化。
4. 返回第一步，继续迭代。

基础设置 (Basic Settings)

实验配置与Transformer层架构

1. 并行策略: TP1 PP8 VPP4 EP32 MBS1 GBS2048 on 256x GB200*
2. MXFP8 配方:
   a. 所有线性层使用 FP8，而 embedding/lm head/communication/optimizer 使用高精度。

下图展示了 Transformer 层的数据流和各模块的计算精度。其中，权重（Weights）使用 MXFP8 精度，而激活值（activations）在 BF16 和 MXFP8 之间转换。

初始状态与基线性能

除了上述配置外，还包括以下设置：
3. 端到端性能
4. 强制平衡，不使用 MTP
5. 开始性能研究时已具备的特性:
a. MXFP8 配方
b. Blackwell 上的 MLA attention kernels
c. MXFP8 分组 GEMM (grouped GEMM)
d. Yarn RoPE 融合/Permute 融合/交叉熵损失融合
e. 灵活的 PP 布局
f. Token分发器: NCCL AlltoAll based, DeepEP based (不支持 NVL72)
g. 主权重为 FP8
h. BF16 优化器状态

基线性能
在初始配置下，DeepSeek-V3 模型的基线性能如下表所示：

基线性能剖析 (Baseline Profile)

通过性能剖析工具，我们识别出以下几个主要瓶颈：
* [红色] 存在过多微小的 kernels，导致 PyTorch/Python 开销和 CUDA kernel 启动开销过大。
* [绿色] NCCL alltoall-based 分发器不是最优的，它要求在发送 tokens 之前进行显式的全局 permute 操作。
* [黄色] Recompute MLA up projection 过程受限于 CPU。

优化项 (Optimization Items)

优化的主要目标是减少 CPU 开销，因为 GPU 速度远超 CPU，但 CPU 无法跟上其速度。同时，使用 FP8 的细粒度 MoE 模型也带来了新的挑战，例如 Router 和预处理过程与 EP 通信的协调，以及量化和缩放因子相关的 kernels。

优化工作主要分为以下几个方面：

1. 核函数融合与优化 (Kernel fusion and optimizations)
   a. MXFP8 数据和缩放因子 (scales) 的批量内存分配
   b. Permute 核函数优化
   c. Router 融合
   d. 融合 MXFP8 swizzle 缩放因子
   e. 融合量化到 normalization
2. 内存节省以解锁更多优化 (Memory saving to unlock more optimizations)
   a. DeepEP
   b. 细粒度重计算 (Fine-grained recompute)
3. CUDA Graphs
4. CPU 优化 (CPU Optimization)
5. HybridEP

1. 核函数融合与优化

通过一系列核函数融合，显著减少了 kernel 数量和 CPU 开销。
* a. 批量内存分配: 针对 MXFP8 数据和 scales，默认自 TE 2.7 版本起。
* b. Permute 核优化: MCore 参数 --moe-permute-fusion，自 TE 2.6 起优化。
* c. Router 融合: MCore 参数 --moe-router-fusion，需要 TE 2.6+。
* d. 融合 MXFP8 swizzle 缩放因子: 默认自 TE 2.7 版本起。
* e. 融合量化到 normalization: 通过 TE 环境变量 NVTE_NORM_FWD_USE_CUDNN=1 NVTE_NORM_BWD_USE_CUDNN=1 配合 cuDNN 9.14+ 实现。

下图展示了融合 alloc zeros 和 swizzle 操作到 Routed experts 后的性能剖析对比。

这些融合带来了显著的性能提升，具体如下：
* a, b, c 项融合: +35 TFLOPS
* 将 kernel 数量从 72 个减少到 31 个。

• d. 融合 MXFP8 swizzle 缩放因子: +35.5 TFLOPS
• e. 融合量化到 normalization: +13.8 TFLOPS

下图展示了应用这些优化后的性能剖析结果。

2. 内存节省

DeepEP

• DeepEP (原型版本):
  a. 将全局 permute 操作融合到通信核中，降低了峰值内存消耗。
  b. 将预处理过程融合到 "cached_notify"，进一步降低了 CPU 开销。
  c. 对于通信本身，DeepEP 在多节点 NVLink (MNNVL) 上存在性能问题，并且比 NCCL 慢。

下图展示了通过 DeepEP 移除 recompute 并进一步减少 router kernels 后的效果。

细粒度重计算

• DeepEP (原型版本): +54.3 TFLOPS
• 细粒度重计算 (Fine-grained recompute): +44.7 TFLOPS
  a. 通过使用 "mlp_act" 重计算替代 "up_proj" 重计算来节省内存。

下图展示了 Transformer 前向传播层中如何通过保存特定中间结果（BF16 输出和 FP8 输入）并重计算激活函数来实现内存节省。

3. CUDA Graphs

• 部分 CUDA Graphs: 支持所有 FP8 配方和灵活的 PP 布局。
• 统一 CUDA Graphs 内存池: 在 PP 阶段之间复用输入/输出缓冲区。
• 捕获所有静态模块: 包括 attention、router、预处理和共享 experts。

该优化带来了 +84.8 TFLOPS 的性能提升。下图对比了使用 CUDA Graphs 前后的性能剖析，显示了 Attention、Router 和 experts 等模块被图化 (Graphed) 后的效果。

4. CPU 优化

• CPU 绑定: +70.6 TFLOPS
  a. 在启动每个进程前添加 bindpcie 脚本。
  b. 使迭代更加稳定。
• TE 中的 CPU 端优化: +12.7 TFLOPS

下图高亮了动态部分的 CPU 开销，并展示了优化后的通信带宽从 -400 GB/s 提升至 -190 GB/s 的变化。

5. HybridEP

• SOL EP 通信: +113.6 TFLOPS
• 完全兼容 NVL72
• 双向 HW dispatch 和 combine: 实现约 ~670 GB/s 的带宽。
• 融合预处理、permute 和 pad

下图展示了 HybridEP dispatch 操作在性能剖析中的表现。

优化历程总结 (Optimization Journey)

下表总结了从基线开始，逐步应用各项优化后，DeepSeek-V3 模型在单 GPU 上的 TFLOPS 性能演进过程。性能从最初的 494.46 TFLOPS/GPU 提升至最终的 970.01 TFLOPS/GPU。

讨论

• 为什么我们不在 GB200* 上使用 FP8 dispatch？

  • FP8 dispatch 并非没有开销。它需要用于反向传播的“反量化和再量化”核函数。
  • 需要评估 FP8 加速的 dispatch 与额外的量化成本。

• 为什么我们不在 GB200* 上使用 1F1B alltoall 重叠（NVIDIA 的批间重叠解决方案）？

  • 1F1B alltoall 重叠会在阶段之间引入同步点，这使得 CPU 开销更加严重。我们可能会在未来重新考虑这一点，例如，使用 CUDA Graphs。

• 与 B200* 的性能比较

  • 在相同设置下，B200 的性能为 780 TFLOPS，GB200 的性能为 970 TFLOPS，速度提升了 1.24倍。

• 与 H100* 的性能比较

  • 与 H100 估计的约 380 TFLOPS 相比，GB200 的速度快了 2.55倍*。

* 仅供技术讨论和参考。性能可能因不同的产品组合而异。

未来工作

• 列表中的优化项

  • 更激进的融合：激活 (act) + 量化 (quant)，GEMM + 激活 + 量化，dispatch + 置换 (permute) + 量化等。
  • 在 Blackwell 架构上改进 SDPA (Scaled Dot-Product Attention) 的反向传播。
  • 改进分组 GEMM (Grouped GEMM)。

• 隐藏 EP (Expert Parallelism) 通信： 细粒度重叠或迭代-批次 (iter-batch) 重叠。

• 通过全迭代 CUDA Graphs 完全移除 CPU 开销。
• 扩展到更多 GPU。
• CPU 卸载以利用 NVLink-C2C。

致谢

• 鸣谢：Xin Yao, Hongxiao Bai, Yaobin Zhang, Tong Liu, Fan Yu, Kunlun Li, Zhongbo Zhu, Zhenhuan Liu, Zijie Yan
• 以及 NVIDIA Compute Arch/TE/cuBLAS/cuDNN 团队。

---