Efficient MoE Training for Biological Foundation Models | NVIDIA Technical Blog

This post details the implementation of NVIDIA Transformer Engine (TE) optimizations for training Mixture-of-Experts (MoE) models, specifically within the BioNeMo recipe for biological foundation models. Key contributions include: 1. Using TE's GroupedLinear to reduce kernel launch overhead by submitting multiple expert GEMMs as a single grouped operation, contrasting with naive Python loop-based expert iteration. 2. Leveraging MXFP8 precision on NVIDIA Blackwell GPUs to reduce memory usage compared to BF16, with hardware acceleration for MXFP8 GEMMs. 3. Introducing a fused GroupedMLP kernel via the TE Sequential API that combines GroupedLinear, ScaledSwiGLU, and routing-weight scaling into a single operation, avoiding intermediate materialization and reducing framework overhead. Benchmarks show up to 2.21x throughput improvement over a Hugging Face baseline.