奇链 奇链
← 文章 / AI技术
NVIDIA 开发者博客 7小时前 · 2026-09-25 02:23:10 · 3 阅读

面向生物基础模型的高效 MoE 训练

随着语言模型规模不断扩大，稠密架构的扩展成本越来越高。在稠密 Transformer 中，每个 token 都要经过每一层，因此增加模型能力就意味着训练和推理的计算量同步增加。

Mixture-of-experts（MoE）架构提供了另一种扩展思路：它包含大量子网络（即专家），但每个 token 只激活其中一小部分。

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

Two block diagrams side by side. A dense Transformer block sends every token through the full feed-forward network, while a sparse MoE block uses a top-2 router to send each token to two of eight experts and combines their outputs as a weighted sum.
图 1. 稠密 Transformer block（左）与稀疏 MoE transformer block（右）对比

本教程演示如何借助 NVIDIA BioNeMo 的 MoE 方案和 TE 将这些技术付诸实践。你将了解 GroupedLinear 如何改进专家计算、MXFP8 如何降低内存占用，以及 GroupedMLP kernel 如何融合量化、SwiGLU 和路由权重缩放。这些能力共同为高效训练基于 MoE 的生物基础模型提供了实用的参考。

前提条件

开始之前，你需要：

  • 熟悉 Python、PyTorch 和分布式训练的基本概念
  • 一个支持 NVIDIA CUDA 的环境——可以使用附带的 Dockerfile 或按方案要求自行安装依赖
  • 至少两块 GPU 以支持专家并行；使用融合的 MXFP8 GroupedMLP kernel 则需要 NVIDIA Blackwell GPU

挑战一：碎片化的专家 kernel

MoE 模型用多组专家网络取代了单一的稠前馈块。但朴素的实现会引发过多的 kernel launch。例如，Hugging Face 的基准实现在 Python 循环中遍历所有专家，每个专家都会触发独立的 kernel launch。

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)

Grouped execution 保留了各专家独立的权重矩阵，但将它们的计算任务批量提交。TE 的 GroupedLinear 在一次调用中执行多个线性变换，方法是先聚合专家权重和输入 token。由于每个专家接收的 token 数量可能不同，GroupedLinear 接受每个专家的 token 数量（split_sizes）作为输入。它通过 TE 的 grouped GEMM 路径处理本地专家，而非为每个专家分别启动一次 PyTorch Linear 操作，从而降低了 launch 和调度开销。

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 循环，这种方法将 gate-up 投影作为一个 grouped 操作提交，而非多次独立的调用。

Hugging Face Transformers 也提供了 grouped_mm。不过，TE 可以将 GroupedLinear 与 MXFP8 量化、激活函数、路由权重缩放及中间数据搬运融合成一个 GroupedMLP kernel，具体展示见后续章节。

两条执行时间线。Python 循环逐个运行八个专家，耗时随专家数量增长；分组 GEMM 则通过单次 GroupedLinear 操作一次性处理全部八个专家。
图 2：Hugging Face 的 modeling_mixtral.py 在 Python 循环中逐个执行专家，而 TE 将所有专家的 GEMM 合并为一次分组操作

挑战 2：模型体积庞大与激活内存压力

MoE 架构扩大了总参数量，而生信任务常处理长序列，导致训练时的激活内存压力骤增。BF16 使用 16 位表示每个模型权重和激活值。

BioNeMo 方案利用 TE 支持 FP8 和 MXFP8 训练以压缩内存占用。这两种格式均以 8 位而非 16 位存储权重和激活值。两者的核心区别在于缩放粒度：MXFP8 为每 32 个连续数值分配一个缩放因子，从而更好地保持数值范围和精度。在 NVIDIA Blackwell GPU 上，MXFP8 得到硬件加速，可调用专用的 Tensor Core 指令执行 MXFP8 GEMM。有关 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 位的主权重。训练框架需额外引入量化与反量化步骤以在两种格式间转换。量化将 BF16 权重和激活值转为 MXFP8 供低精度 GEMM 使用；反量化则将计算结果还原为高精度格式。常规实现将这两步作为独立操作执行，这也促使我们提出下文描述的融合 MLP 路径。

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

TE 的 autocast 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 路径，并避免物化部分中间结果。

两张纵向流程图。左侧标注
图 4：MXFP8 路径在 down 投影之前将多个操作融合为一个 kernel

结果

以上是 BioNeMo recipe 中的几项优化。在八块 NVIDIA B200 Tensor Core GPU 上进行的训练基准测试中，该方案取得了最高 2.21 倍于 Hugging Face 基线的吞吐量。

一张条形图，标题为“Mixtral-8x7B 在 8×B200 上的训练吞吐量”，测量指标为每 GPU 每秒处理的 token 数。三个条形：使用 BF16 和 FSDP8 的 Hugging Face 为 4,096；使用 BF16 和 FSDP8 的 Transformer Engine 为 4,447；使用专家并行和 MXFP8 的 Transformer Engine 为 9,050，标注为 2.21x。
图 5. 在八块 NVIDIA B200 Tensor Core GPU 上运行 Mixtral-8x7B 的训练吞吐量，最高可达 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 原生 Transformer Engine 配方，并在NVIDIA Transformer Engine 文档中了解优化的 MoE 内核详情。

原始来源： NVIDIA 开发者博客

评论 (0)