返回 文章 build CMS 文章

NVIDIA Transformer Engine 优化 MoE 生物基础模型训练：吞吐量提升 2.21 倍

用 NVIDIA Transformer Engine 的 GroupedLinear、MXFP8 和融合 GroupedMLP 内核，把 MoE 生物基础模型训练吞吐量提升到 Hugging Face 基线的 2.21 倍。

NVIDIA
MoETransformer Engine生物基础模型MXFP8
成长分 / 100 78 综合收获、行动、留存与影响

NVIDIA Transformer Engine 优化 MoE 生物基础模型训练：吞吐量提升 2.21 倍
为什么值得读了解 MoE 架构在扩展模型容量时面临的三大工程瓶颈：专家内核碎片化、激活内存压力和低精度量化开销。

掌握 Transformer Engine 提供的具体优化原语（GroupedLinear、MXFP8、ScaledSwiGLU、Sequential API）及其代码用法。

关键洞察
  1. 稠密 Transformer 中每个 token 经过每一层，增加能力会同时增加训练和推理计算量；MoE 通过每个 token 只激活少量专家来更高效地扩展模型容量。
  2. 朴素 MoE 实现（如 Hugging Face 基线）在 Python 循环中逐个遍历专家，每个专家触发单独内核启动，导致启动和调度开销过大。
  3. TE 的 GroupedLinear 通过收集专家权重和输入 token，在一次调用中应用多个线性变换，接受每个专家的 token 计数（split_sizes）作为额外位置参数。
转成行动

深入阅读

正文与原文对照

原文保真覆盖：全文原文字符：10720

随着语言模型规模的增长，扩展稠密架构变得越来越昂贵。在稠密 Transformer 中，每个 token 都会经过每一层，因此增加能力会同时增加训练和推理的计算量。

混合专家（MoE）架构采用不同的扩展方法，它使用许多子网络（即专家），但每个 token 只激活其中一小部分。

这种权衡使得 MoE 架构对大型语言模型（LLM）社区越来越有吸引力。它们可以更高效地扩展模型容量，但收益在很大程度上取决于实现方式。碎片化的专家计算会降低 GPU 利用率。路由会带来通信开销，而更大的参数占用会带来内存和分布式训练方面的挑战。NVIDIA Transformer Engine（TE）通过针对分组专家计算、内核融合和低精度训练的优化原语，帮助解决这些瓶颈。随着生物基础模型的参数数量和序列长度不断增长，这些原语可以在扩展模型容量的同时提高 GPU 效率。

两个并排的框图。稠密 Transformer 块将每个 token 送入完整的前馈网络，而稀疏 MoE 块使用 top-2 路由器将每个 token 发送到八个专家中的两个，并将它们的输出组合为加权和。

图 1. 稠密 Transformer 块（左）与稀疏 MoE Transformer 块（右）的对比

本教程展示如何通过 NVIDIA BioNeMo MoE 配方和 TE 将这些技术付诸实践。你将看到 GroupedLinear 如何改进专家计算，MXFP8 如何减少内存使用，以及 GroupedMLP 内核如何融合量化、SwiGLU 和路由权重缩放。这些能力共同为高效训练基于 MoE 的生物基础模型提供了实用参考。

先决条件

在开始之前，你需要：

  • 熟悉 Python、PyTorch 和分布式训练概念
  • 一个支持 NVIDIA CUDA 的环境——你可以使用链接的 Dockerfile 或安装配方依赖项
  • 至少两个 GPU 用于专家并行；使用融合 MXFP8 GroupedMLP 内核需要 NVIDIA Blackwell GPU

挑战 1：碎片化的专家内核

MoE 模型用多个专家网络替换单个稠密前馈块。然而，朴素的实现可能会触发过多的内核启动。例如，Hugging Face 基线实现会在 Python 循环中遍历所有专家，每个专家都会触发单独的内核启动。

for expert_idx, expert_layer in enumerate(self.experts):
idx, top_x = torch.where(expert_mask[expert_idx])
current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
current_hidden = expert_layer(current_state) * routing_weights[top_x, idx, None]
final_hidden_states.index_add_(0, top_x, current_hidden)

分组执行保留了各个专家矩阵，但将其工作一起提交。TE 的 GroupedLinear 通过收集专家权重和输入 token，在一次调用中应用多个线性变换。由于每个专家可能接收到不同数量的 token，GroupedLinear 接受每个专家的 token 计数（split_sizes

）。它通过 TE 分组 GEMM 路径提交本地专家，而不是为每个专家启动一个 PyTorch Linear 操作，从而减少了启动和调度开销。

按如下方式使用 GroupedLinear。每个专家保留自己的权重张量（weight0

、weight1

等），该调用接受每个专家的 token 计数作为额外的位置参数：

from transformer_engine.pytorch.ops import GroupedLinear
experts_gate_up = GroupedLinear(
num_groups=num_local_experts,
in_features=hidden_size,
out_features=2 * intermediate_size,
bias=False,
dtype=torch.bfloat16,
device="cuda",
)
gate_up_output = experts_gate_up(tokens, split_sizes)

与 Python 循环相比，这种方法将门控上投影作为一次分组操作提交，而不是多次单独的调用。

Hugging Face Transformers 也提供了 grouped_mm

。然而，TE 可以将 GroupedLinear

与 MXFP8 量化、激活、路由权重缩放以及中间数据移动融合到一个 GroupedMLP 内核中，如后续章节所示。

两条执行时间线。Python 循环逐个运行八个专家，因此时间随专家数量增长。分组 GEMM 在单个 GroupedLinear 操作中运行全部八个专家。

图 2. Hugging Face 的 modeling_mixtral.py 在 Python 循环中一次运行一个专家，而 TE 将所有专家 GEMM 批处理为一次分组操作

挑战 2：大模型规模和激活内存

MoE 架构增加了总参数容量，而基因组学工作负载通常使用长序列，这给训练期间的激活内存带来了压力。BF16 使用 16 位来表示每个模型权重和激活值。

BioNeMo 配方使用 TE 来支持 FP8 和 MXFP8 训练，从而减少内存使用。两种格式都用 8 位而不是 16 位来表示权重和激活值。FP8 和 MXFP8 的主要区别在于缩放粒度：MXFP8 为每 32 个连续值组成的块分配一个缩放因子，有助于保持数值范围和精度。在 NVIDIA Blackwell GPU 上，MXFP8 受硬件加速，使 MXFP8 GEMM 能够使用专门的 Tensor Core 指令。有关 MXFP8 和块缩放的详细信息，请参阅 Transformer Engine FP8 入门指南。

位布局比较。BF16 每个值使用 16 位：1 位符号、8 位指数、7 位尾数。MXFP8 使用 E4M3 格式的 8 位，外加每 32 个元素一个共享的 E8M0 缩放因子。

图 3. BF16 每个值使用 16 位，而 MXFP8 使用 8 位。来源：Transformer Engine FP8 入门指南

挑战 3：低精度训练中的量化开销

尽管大多数训练计算使用 8 位精度，模型仍以 16 位保留其主权重。因此，训练框架会添加量化和反量化步骤以在格式之间转换。量化在低精度 GEMM 之前将 BF16 权重和激活值转换为 MXFP8；反量化将结果转换回更高精度的格式。朴素路径将这些步骤作为单独操作执行，这促使了接下来描述的融合 MLP 路径。

fp8_recipe = te_recipe.MXFP8BlockScaling()
model = TEMixtralMXFP8ForCausalLM(config, fp8_recipe=fp8_recipe, dispatcher=dispatcher)

TE 自动转换 API 为模型的前向和反向传播启用 MXFP8 精度：

with te.autocast(enabled=True, recipe=self._fp8_recipe):
for decoder_layer in self.layers:
hidden_states = decoder_layer(hidden_states)

完整代码请参见 BioNeMo recipe。

要使用融合 MLP，请导入 Transformer Engine Sequential API，将 gate_up

、ScaledSwiGLU

和 down

串联起来。该 API 还会将反量化折叠进融合路径中。ScaledSwiGLU

将路由概率（“scales”）与专家前馈网络计算相结合。

from transformer_engine.pytorch.ops import GroupedLinear, ScaledSwiGLU, Sequential
experts_ffn = Sequential(GroupedLinear(gate_up), ScaledSwiGLU(), GroupedLinear(down))

TE Sequential API 会扫描操作，当模式匹配时，将 GroupedLinear

→ ScaledSwiGLU

→ GroupedLinear

序列替换为一个融合操作对象：前向传播使用 ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8，并匹配一个融合反向操作。这减少了框架开销，将 SwiGLU 和概率缩放工作融合到分组 MLP 路径中，并避免了一些中间结果的物化。

两个垂直流程图。左侧标记为“未融合 MLP”：五个顺序框，分别读取量化、门控上投影、SwiGLU、反量化和门控下投影。右侧标记为“融合 MLP”：三个框，分别读取量化、一个高亮的融合分组 MLP 框（替换了门控上投影、SwiGLU 和反量化步骤）以及门控下投影。

图 4. MXFP8 路径在下投影之前将多个操作融合到一个内核中

结果

这些是 BioNeMo 配方中的几项优化。在八块 NVIDIA B200 Tensor Core GPU 上的训练基准测试中，该配方实现了高达 Hugging Face 基线 2.21 倍的吞吐量。

一个标题为“Mixtral-8x7B 在 8×B200 上的训练吞吐量”的条形图，测量每 GPU 每秒处理的 token 数。三个条形：Hugging Face 使用 BF16 和 FSDP8 为 4,096；Transformer Engine 使用 BF16 和 FSDP8 为 4,447；Transformer Engine 使用专家并行和 MXFP8 为 9,050，标记为 2.21 倍。

图 5. Mixtral-8x7B 在八块 NVIDIA B200 Tensor Core GPU 上的训练吞吐量，显示高达 Hugging Face 基线 2.21 倍的吞吐量

运行配方

从双 GPU 的 L0_sanity 配置开始，以确认专家并行和训练环境正常工作：

torchrun --nproc_per_node=2 train_fsdp2_ep.py --config-name L0_sanity

验证之后，扩展到 Mixtral-8x7B 配置，在八块 GPU 上使用专家并行（EP=8）和 MXFP8 精度：

torchrun --nproc_per_node=8 train_fsdp2_ep.py --config-name L1_8x7B_ep checkpoint.ckpt_dir=/path/to/ckpt

根据你的 GPU 和内存需求选择 BF16 或 MXFP8，并设置数据并行和专家并行的大小，使其乘积等于 GPU 总数。 配方 README 包含启动、检查点和基准测试命令。

在 BioNeMo Recipes 中试用 Mixtral Native Transformer Engine 配方，并在 NVIDIA Transformer Engine 文档中进一步了解优化的 MoE 内核。

致谢

Sudhakar Singh US、Varun Thumbe US、Santosh Santosh US、Timur Rvachov US、Chris Hoge US，