📖 AI 总结

文章介绍了如何利用NVIDIA Transformer Engine高效训练基于混合专家（MoE）架构的生物基础模型。MoE通过每次仅激活部分专家子网络，比密集Transformer更高效地扩展模型容量，但其性能高度依赖实现方式，专家计算的碎片化、路由通信开销以及更大的参数规模都会带来挑战。为此，文章提出三项关键技术：GroupedLinear将多个专家GEMM合并为一次分组操作，减少内核启动开销；MXFP8块缩放8位精度相比BF16降低内存占用，并在Blackwell GPU上获得硬件加速；Sequential API将GroupedLinear、ScaledSwiGLU与路由权重缩放融合为单一内核，避免中间结果物化。在八块B200 GPU的基准测试中，BioNeMo方案吞吐量最高达到Hugging Face基线的2.21倍。这表明优化后的MoE训练方案能显著提升生物基础模型的训练效率。