博客

实现最高 41% 的预训练提速:在 B200 上使用 TorchTitan 和 MXFP8 以及 DeepEP 进行 DeepSeek-V3 训练

要点速览

通过 PyTorch 和 Nebius 的联合协作,我们成功在搭载 256 个 NVIDIA B200 GPU 的集群上,使用 TorchTitan 实现了 DeepSeek-V3 混合专家 (MoE) 模型(16B 和 671B)的训练。我们在 BF16 基准之上评估了两种正交优化方案:MXFP8 训练(通过 TorchAO)和 DeepEP 通信加速(通过 DeepEP)。主要亮点如下:

  • DeepSeek-V3 671B:仅使用 DeepEP 即可获得 859 token/秒(+32%) 的性能,优于 BF16 基准(651 token/秒)。在分组 GEMM 上添加 MXFP8 并结合 DeepEP,性能进一步提升至 918 token/秒,总吞吐量提升 +41%
  • DeepSeek-V3 16B MoE:通过 1,500 步的损失收敛实验证实,MXFP8 训练与 BF16 等效(收敛表现无劣化)。

所有实验均在 Nebius Cloud 上使用开源的 PyTorch 原生工具运行,且完全可复现。请参阅最后一部分(可复现性)以获取所有方案。

为什么进行此项实验

训练前沿规模的 MoE 模型既需要软件的成熟度,也需要系统级的效率。随着 NVIDIA Blackwell (B200) GPU 的到来及其对 MXFP8 张量核心的原生支持,我们有机会超越单纯的“训练提速”,向显著优化性价比迈进,特别是对于计算和 GPU 间通信均为瓶颈的 MoE 架构而言。

我们将 TorchTitan 作为预训练框架,旨在回答以下问题:

  1. MXFP8 能将计算加速多少? Blackwell 的第五代张量核心原生支持 MXFP8,使得符合条件的 GEMM 运算在峰值 TFLOPS 上比 BF16 高出 2 倍。我们希望衡量在 TorchTitan 内将 MXFP8(通过 TorchAO)应用于 DeepSeek-V3 的路由专家(分组 GEMM)和线性层时,实际的端到端加速比。
  2. DeepEP 能将通信加速多少? MoE 模型在每一层需要两次全对全(all-to-all)交换,以将 token 分派给专家并组合其输出。由于传输大小和目的地在每一步都由路由器动态确定,标准的集合通信效率低下(通常为预先知道固定大小的传输而设计)。随着专家并行(EP)规模的增加,该问题愈演愈烈,all-to-all 成为巨大瓶颈。DeepEP 使用专门构建的 NVLink 和 RDMA 内核替换了标准的 all-to-all 后端,通过允许 GPU 直接发送权重来减少 CPU 参与,从而降低延迟。这更适合此类可变工作负载。我们希望量化通过减少 TorchTitan 专家并行流水线中的通信瓶颈所获得的吞吐量增益。
  3. 这些计算和通信增益能否叠加? MXFP8 针对计算(GEMM),而 DeepEP 针对通信(all-to-all)。我们希望验证在 TorchTitan 中同时应用这两种优化是否能产生比单一优化更大的累积加速,并展示在这种组合配置下大规模端到端预训练的稳定性。

背景

在深入实验之前,本节简要概述我们评估的两种关键技术:MXFP8 混合精度训练(包括我们测试的不同方案)以及 DeepEP 的优化专家并行通信。

MXFP8:通过 TorchAO 实现微缩(Microscaling)FP8

MXFP8(微缩 FP8)是由 OCP 微缩格式规范 定义的一种低精度数值格式。与每个张量或每行使用单个缩放因子的标准 float8 不同,MXFP8 为每 32 个元素的块分配一个 共享指数(E8M0 缩放)。这种更细粒度的缩放既保留了数值保真度,又支持硬件的 FP8 张量核心。

NVIDIA Blackwell GPU 通过 tcgen05.mma 张量核心指令为 MXFP8 提供 原生硬件支持,这意味着 MXFP8 GEMM 以完整的 FP8 吞吐量运行,无需仿真开销。这使得 B200 成为 MXFP8 训练的天然硬件目标。

在实践中,MXFP8 是通过 TorchAO 应用的,它提供了:

  • 针对线性层的 MXFP8:将 nn.Linear 层转换为在每一层中动态地将输入量化为 MXFP8,用于所有三个 GEMM(前向输出、输入梯度、权重梯度),并累加回 BF16。
  • 针对分组 GEMM 的 MXFP8:将 torch._grouped_mm 操作转换为使用 TorchAO 的 _to_mxfp8_then_scaled_grouped_mm 自定义自动求导函数,动态地将所有三个分组 GEMM 的输入量化为 MXFP8,从而在主导 MoE 专家层的分组矩阵乘法上提供净加速。

在我们的实验中,我们使用了两种 MXFP8 配置

配置 描述
所有分组 GEMM 的 MXFP8 应用于分组 GEMM 的 MXFP8

MXFP8 GMM 方案 在每个分组 GEMM 之前动态地量化输入。当 GEMM 较大时,此方案效果显著;计算开销按 O(GMNK) 增长,而量化开销仅按输入的 O(MK) 或权重的 O(GNK) 增长。因此,对于大维度,开销可以忽略不计,净加速显著。然而,当 GEMM 维度较小时(如在较小的 MoE 模型中),量化开销可能等于或超过 MXFP8 的加速比,导致性能持平甚至下降。

图 1:从路由器输入到 all-to-all 组合输出的 MXFP8 默认过程

DeepEP:优化的专家并行通信

DeepEP 是由 DeepSeek 开发的 GPU 通信库,专门为 MoE 专家并行中的 all-to-all token 分派和组合操作而构建。在标准专家并行中,每个 GPU 可能需要根据路由器的决策向其他所有 GPU 发送 token(并从其接收 token)。标准集合通信库(例如 NCCL、RCCL)提供的通用 all-to-all 原语无法利用 MoE 工作负载的结构。

DeepEP 通过以下方式解决了这一问题:

  • 流水线式节点间和节点内通信:数据首先通过节点内高带宽 NVLink 移动,然后通过 InfiniBand 的 RDMA 跨节点传输。这利用了节点内(~900 GB/s NVLink)和节点间互连之间的带宽不对称性。
  • 基于 NVSHMEM 的 GPU 发起 RDMA: DeepEP 构建在 NVSHMEM 之上,它公开了一个分区全局地址空间(PGAS),每个 GPU 将对称内存区域映射为所有节点可寻址。结合 InfiniBand GPUDirect Async (IBGDA),GPU 可以直接从 CUDA 内核驱动 NIC 操作,无需任何 CPU 参与,消除了小型动态传输的关键延迟源。
  • 融合元数据传输:将 token 嵌入、专家索引和路由权重打包到单个通信操作中,减少了内核启动和同步开销。
  • 可配置的 SM 使用率:可以调节通信内核以使用特定数量的流式多处理器(SM),从而释放其余部分用于重叠计算。

在我们拥有 32 个节点(EP=32)的设置中,DeepEP 替换了标准的 PyTorch all-to-all 后端,用于 MoE 分派/组合操作。

硬件和集群环境

Nebius B200 集群

所有实验均在具有以下规格的 Nebius Cloud 集群上运行

组件 规格
节点数 32
每个节点 GPU 数 8 x NVIDIA B200
总 GPU 数 256
节点内互连 NVLink / NVSwitch
节点间互连 InfiniBand
调度 Soperator (Slurm + Kubernetes)

软件栈

组件 详情
训练框架 TorchTitan
混合精度 TorchAO (MXFP8)
专家通信 DeepEP
编译 torch.compile (模型 + 损失)
数据集 C4

在基础设施方面,我们运行在 Nebius Soperator 上,它将 Slurm 风格的调度和多节点语义引入 Kubernetes,而无需将 Kubernetes 的操作复杂性暴露给机器学习工程师。该环境开箱即用,已达到生产就绪状态:一个预先配置、性能优化的集群,此次运行只需极少的设置。两项能力最为关键:

  • 首先Soperator 通过主动检查(包括 NCCL all-reduce 基准)和被动信号(例如 GPU XID 错误)持续验证 GPU 和互连的健康状况,并自动排除并替换低于性能阈值的节点。相同的遥测数据会反馈到可观测性仪表板和警报中,这加快了故障排除速度并减少了查找根本原因的时间。
  • 其次扩缩容非常简单:调整集群大小只需一条命令,新节点会自动继承完整的运行时环境。这让我们能够专注于优化训练方案,而不是集群搭建、滞后节点缓解或重新运行失败的作业。

实验 1:DeepSeek-V3 671B

训练配置

参数
模型 DeepSeek-V3 671B
节点/GPU 32 / 256
序列长度 8192
本地批大小 64
张量并行 (TP) 2
流水线并行 (PP) 2
数据并行 (DP) 1
FSDP -1 (自动)
上下文并行 (CP) 1
专家并行 (EP) 32
专家张量并行 (ETP) 1
激活检查点 全量
编译 模型, 损失
学习率 1e-4
学习率预热步数 2000
数据集 C4

所有 671B 配置的完整 Slurm 脚本和 TorchTitan 训练配置可在 Nebius 的 ML-CookBook 仓库中获取。

测试的配置

我们评估了三种配置,改变了精度模式和专家并行通信后端

配置 精度 专家通信
BF16 + 标准 EP BF16 标准 all-to-all
BF16 + DeepEP BF16 DeepEP
MXFP8 GMM + DeepEP 所有分组 GEMM 上的 MXFP8 DeepEP

吞吐量结果

图 2:每种配置的吞吐量(tokens/秒)。两种不同的 EP 后端 x 两种不同的精度。该图表展示了每项优化如何对总吞吐量增益做出贡献。从 BF16 基准(651 TPS)开始,DeepEP 在 BF16 的基础上增加了 +32%。DeepEP 和 MXFP 的组合增加了 +41%,总计达到 918 TPS。

关键观察

  • 仅 DeepEP 就提供了 +32.0% 的增益。这是单一最大的个人增益。对于此 EP=32 跨 32 个节点的 671B 模型,节点间 all-to-all 通信是一个主要瓶颈,而 DeepEP 的优化 RDMA+NVLink 转发极大地缓解了它。
  • MXFP8 GMM + DeepEP 结合良好,合计达到 +41.0%。这些增益是相加的,因为它们针对不同的瓶颈:MXFP8 加速计算(GEMM),而 DeepEP 加速通信(all-to-all)。

实验 2:DeepSeek-V3 16B MoE 损失收敛验证

训练配置

参数
模型 DeepSeek-V3 16B MoE
节点/GPU 32 / 256
序列长度 8192
本地批大小 16
张量并行 (TP) 1
流水线并行 (PP) 1
数据并行 (DP) 1
FSDP -1 (自动)
上下文并行 (CP) 1
专家并行 (EP) 32
专家张量并行 (ETP) 1
激活检查点 全量
编译 模型, 损失
学习率 1e-4
学习率预热步数 2000
数据集 C4

所有 16B 配置的完整 Slurm 脚本和 TorchTitan 训练配置可在 Nebius ML-CookBook 仓库中获取。

测试的配置

所有 16B 实验均使用标准 EP(无 DeepEP)

配置 描述
BF16 基准
MXFP8 (MXFP8 a2a EP 后端) 仅在分组 GEMM 上使用 MXFP8(专家并行 + GMM)

损失收敛

我们运行了每种 16B 配置 1,500 步训练,以验证 MXFP8 不会降低收敛性。

图 3:BF16 和 MXFP8 在 1,500 步内的训练损失。两条曲线几乎相同。插图显示了放大后的最后 10% 训练过程。

图 4:训练最后 30% 步数的放大视图。两种配置紧密重合,证实了等效的收敛行为。

损失曲线证实,MXFP8 训练在此训练窗口内与 BF16 基准表现出相同的收敛性。MXFP8 与 BF16 之间没有出现有意义的背离。

结果摘要

模型 配置 TPS 加速比
671B BF16 + 标准 EP 651 基准
671B BF16 + DeepEP 859 +32.0%
671B MXFP8 GMM + DeepEP 918 +41.0%

经验教训

  • DeepEP 是 671B 的单一最大杠杆。 在 EP=32 跨 32 个节点时,节点间 all-to-all 通信主导了步长时间。DeepEP 的分层 NVLink+RDMA 转发大幅削减了此瓶颈。
  • 这两种优化是互补的。 MXFP8 针对计算(GEMM),DeepEP 针对通信(all-to-all)。由于它们解决了独立的瓶颈,它们的增益是可以叠加的。合计 +41% 的增益接近个人改进的总和。

MXFP8 的收敛性与 BF16 等效。 在 16B 模型上进行 1,500 步训练后,MXFP8 和 BF16 的损失曲线无法区分,这与 MoE 架构上其他 MXFP8 训练实验的结果一致。

未来工作

此处展示的结果尚未完全体现 MXFP8 在 671B 模型上的全部潜力。我们目前的 671B 实验仅在分组 GEMM(路由专家)上应用了 MXFP8,因为 torch.compile 在与张量并行 (TP) 结合使用时,尚不支持 MXFP8 线性层。由于 671B 配置需要 TP=2,我们无法在注意力层和共享专家线性层上启用 MXFP8。

一旦 TP + torch.compile + MXFP8 线性层支持登陆 TorchTitan,我们预计 671B 的吞吐量将进一步提升。

可复现性

所有实验均使用开源工具,并可在 Nebius Cloud B200 集群上复现。更多信息请参阅 Nebius 仓库中的 README 文件。

组件 参考
TorchTitan https://github.com/pytorch/torchtitan 
TorchAO https://github.com/pytorch/ao
DeepEP https://github.com/deepseek-ai/DeepEP 
671B 训练配置 train_671b.sh
16B 训练配置 train_16b.sh
环境设置 setup.sh, requirement.txt, README