NVIDIA 博客:用 Transformer Engine 高效训练生物基础模型 MoE
原标题:Efficient MoE Training for Biological Foundation Models
AI 摘要
NVIDIA 发布技术博客,介绍如何利用 Transformer Engine(TE)和 BioNeMo MoE recipe 高效训练生物基础模型。文章指出 MoE 架构虽能更高效扩展模型容量,但面临专家计算碎片化、路由通信开销和内存压力等瓶颈。TE 通过 GroupedLinear、MXFP8 低精度训练和融合 GroupedMLP 内核(整合量化、SwiGLU 与路由权重缩放)来提升 GPU 效率,并给出具体代码示例与运行前提。
正文节选
As language models grow, scaling dense architectures becomes increasingly expensive. In a dense transformer, every token passes through every layer, so adding capabilities increases computation for both training and inference. Mixture-of-experts (MoE) architectures take a different approach to scaling by using many subnetworks, or experts, while activating only a small subset for each token. This tradeoff has made MoE architectures increasingly attractive to the large language model (LLM) commun