NVIDIA 发布用于生物基础模型的高效 MoE 训练配方

用途与变化

NVIDIA 开发者博客于 2026 年 9 月 24 日发布教程，介绍如何用 NVIDIA BioNeMo MoE 配方与 Transformer Engine（TE）训练基于 MoE 的生物基础模型。文章围绕三个训练瓶颈给出对应做法：

• 专家计算碎片化：传统实现按专家在 Python 循环中逐个调用，产生大量内核启动。TE 的 GroupedLinear 将多个专家的线性变换合并为一次分组 GEMM 调用，接受每个专家的 token 数（split_sizes），保留各自权重张量。
• 显存压力：配方支持 FP8 与 MXFP8，权重和激活用 8 位而非 BF16 的 16 位；MXFP8 对每 32 个连续值分配一个缩放因子，Blackwell GPU 对其有硬件加速。
• 量化开销：TE 的 Sequential API 将 GroupedLinear、ScaledSwiGLU 与路由权重缩放组合，匹配时替换为融合算子 ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 及对应反向算子，减少中间结果物化。

适用场景

面向需要训练 MoE 架构生物基础模型（如蛋白质、基因组方向）的开发者。先决条件包括 Python、PyTorch 与分布式训练基础、CUDA 环境，专家并行至少需要两块 GPU，融合 MXFP8 GroupedMLP 内核要求 Blackwell GPU。

结果与限制

文章称在八块 NVIDIA B200 Tensor Core GPU 上，该配方训练 Mixtral-8x7B 的吞吐最高为 Hugging Face 基线的 2.21 倍。该数据来自厂商自述基准，未在本次整理中复现；实际收益会因硬件、并行配置和模型设置而异。文中提到的启动命令、配置文件名称和精度选项需以配方 README 与 TE 文档为准，本文不代为验证。