高效 MoE 训练：生物基础模型
来自 NVIDIA Developer · 期号 2026-09-24
摘要
MoE 模型把多个专家网络放在一个 Python 循环里逐个调用，GPU 利用率很低。 NVIDIA Transformer Engine 的 GroupedLinear 把多个专家的矩阵乘法合并成一次分组操作，减少内核启动开销。MXFP8 块缩放 8 位精度比 BF16 省一半显存，在 Blackwell GPU 上有硬件加速。TE Sequential API 把分组线性、SwiGLU 激活和路由权重缩放融合成一个 ForwardGroupedMLP 内核，避免中间结果落盘。 在八卡 B200 上，BioNeMo 配方比 Hugging Face 基线吞吐量最高提升 2.21 倍。 本文是技术教程，附了完整的代码和启动命令。
深读
导读 MoE 模型在训练时面临专家计算碎片化、激活内存大和量化开销三个工程瓶颈。NVIDIA 这篇教程给出了 Transformer Engine 的完整解决方案：用 GroupedLinear 合并专家矩阵乘法、用 MXFP8 块缩放精度省一半显存、用 Sequential API 将多个操作融合成一个内核。读完能拿到一套可直接运行的 BioNeMo 训练配方，以及理解每个优化手段解决的具体问题。 关键信息 - 在八卡 NVIDIA B200 GPU 上，BioNeMo 配方相比 Hugging Face 基线实现了最高 2.21 倍的吞吐量提升。 - MXFP8 以每 32 个连续值为一个块分配缩放因子，在 Blackwell GPU 上有专用 Tensor Core 指令硬件加速。 - TE Sequential API 将 GroupedLinear、ScaledSwiGLU 和路由权重缩放融合成一个 ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 内核，避免中间结果落盘。 - Hugging Face 基线在 Python 循环中逐个调用专家，每个专家触发独立的内核启动；TE 的 GroupedLinear 将所有专家的 GEMM 提交为一次分组操作。 - 模型在训练中保留 BF16 主权重，量化步骤将权重和激活转为 MXFP8 执行低精度 GEMM，再反量化回高精度格式。 逐节讲解 MoE 的工程挑战与 Transformer Engine 的应对思路 MoE 架构通过为每个 token 只激活少量专家来扩展模型容量，比稠密 Transformer 更高效。但实现不当会导致三个问题：专家计算碎片化降低 GPU 利用率、路由引入通信开销、更大的参数量带来内存和分布式训练挑战。NVIDIA Transformer Engine 提供了分组专家计算、内核融合和低精度训练的原语来应对这些瓶颈。 挑战一：碎片化的专家内核 MoE 模型用一个 Python 循环逐个调用专家，每个专家触发独立的内核启动，开销很大。TE 的 GroupedLinear 将多个专家的线性变换合并为一次分组操作，接受每个专家的 token 数作为 split_sizes 参数，通过 TE 分组 GEMM 路径提交，减少启动和调度开销。Hugging Face Transformers 也有 grouped_mm，但 TE 能将 GroupedLinear 与 MXFP8 量化、激活函数、路由权重缩放融合成更完整的 GroupedMLP 内核。 挑战二：大模型与激活内存压力 MoE 增加了总参数量，基因组学任务常用长序列，进一步推高训练中的激活内存。BF16 用 16 位表示每个值，而 MXFP8 只用 8 位，内存减半。MXFP8 与 FP8 的区别在…