稠密 Transformer 的每个 token 都要经过全部层，继续扩大能力的代价同时落在训练和推理两侧。混合专家（MoE）架构换了一种扩展方式：保留大量专家子网络，但每个 token 只激活其中一小部分。这让容量扩展更经济，但收益高度依赖实现质量——如果专家计算被拆成大量小内核、量化与反量化步骤零散分布，框架开销会把稀疏带来的好处抵消掉。
NVIDIA 在 BioNeMo 的 MoE 配方中给出的第一项优化是 Transformer Engine 的 GroupedLinear。朴素实现会用 Python 循环逐个遍历专家，每个专家触发一次内核启动；GroupedLinear 则把专家权重与输入 token 聚合后一次提交，由于各专家接收的 token 数不同，接口额外接受 per-expert 的 token 计数（split_sizes）作为参数。结果是多次专家 GEMM 被合并成一次分组操作，内核启动开销显著下降。
第二项优化是精度。该配方支持 FP8 与 MXFP8 训练，两者都用 8 位表示权重和激活，区别在缩放粒度：MXFP8 为每 32 个连续值分配一个缩放因子，更能保住数值范围和精度，并且在 Blackwell GPU 上由硬件加速。由于训练框架仍以 16 位保存主权重，量化与反量化不可避免，朴素路径会把这些转换拆成多个独立步骤。
第三项优化针对上述碎片化。Transformer Engine 的 Sequential API 会扫描算子序列，一旦匹配到 GroupedLinear → ScaledSwiGLU → GroupedLinear 的模式，就替换为融合算子：前向为 ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8，反向也有对应的融合实现。这样 SwiGLU 与路由权重缩放被并入分组 MLP 路径，中间结果不再落到显存里，框架层开销也随之减少。
基准测试在 8 张 NVIDIA B200 Tensor Core GPU 上进行，模型配置为 Mixtral-8x7B，专家并行度设为 8 并启用 MXFP8。结果显示，BioNeMo 配方的训练吞吐最高达到 Hugging Face 基线的 2.21 倍。需要注意的是，融合的 MXFP8 GroupedMLP 内核要求 Blackwell 架构，且专家并行至少需要两张 GPU。
配方并非只能照搬。NVIDIA 建议按 GPU 与显存条件在 BF16 和 MXFP8 之间选择，并让数据并行度与专家并行度的乘积等于 GPU 总数；仓库 README 中提供了启动、检查点与基准测试命令。对生物基础模型团队来说，这套组合的意义在于把 MoE 的稀疏优势真正落到吞吐上，而不只是在参数规模上做文章。
WeChat
Profile