Efficient MoE Training for Biological Foundation Models

• The Hugging Face baseline MoE layer loops over experts in Python, with separate kernel launches per expert. Transformer Engine GroupedLinear takes per-expert token counts (split_sizes) and submits all local experts as one grouped GEMM, cutting launch overhead.
• MXFP8 stores weights and activations in 8 bits like FP8, but assigns a scaling factor to every block of 32 consecutive values, which NVIDIA says helps preserve numerical range and accuracy. It is hardware-accelerated on Blackwell Tensor Cores.
• Chaining GroupedLinear, ScaledSwiGLU and GroupedLinear through TE's Sequential API lets it swap in a single fused MXFP8 kernel that also folds in routing-weight scaling and dequantization, avoiding some intermediate tensors.
• NVIDIA reports the BioNeMo recipe reached up to 2.21x the throughput of the Hugging Face baseline on eight B200 GPUs. The fused MXFP8 GroupedMLP kernel requires Blackwell GPUs, and expert parallelism needs at least two GPUs.
• The recipe suggests validating first with a two-GPU L0_sanity config, then scaling to Mixtral-8x7B with expert parallelism of 8, keeping data-parallel times expert-parallel size equal to the total GPU count.

• transformer — The neural network architecture behind modern AI models, built on attention — letting every word directly consider every other word in parallel.
• mixture-of-experts — A model built from many specialist sub-networks where only a few activate per token, giving big-model capability at small-model running cost.
• quantization — Shrinking a model by storing its numbers less precisely — like rounding — so it runs faster and fits on smaller hardware, at a small quality cost.
• token — The chunk of text a model reads and writes in — roughly three-quarters of a word — and the unit AI usage is billed in.

NVIDIA's recipe fuses grouped expert GEMMs, SwiGLU, and MXFP8 quantization into a single kernel on Blackwell GPUs, delivering 2.21x training throughput over a naive Hugging Face MoE baseline, a directly applicable technique for anyone training mixture-of-experts models at scale.

• articleValidate GPU Cluster Readiness Before AI Workloads Land
• articleBuilding an Adaptive Agentic Cybersecurity System with NVIDIA Nemotron
• articleRun NVIDIA BioNeMo NIM Microservices for Protein Structure Prediction in Claude Science
• articleScale AV Perception Across Vehicle Platforms with NVIDIA Omniverse NuRec

• articleAccelerating Dropless MoE Training in JAX with NVIDIA Transformer EngineTanya Lenz
• articleDense vs. MoE Models: Active Parameters, Throughput, and When to Choose EachElizabeth Goodman
• articleEMO: Pretraining mixture of experts for emergent modularityallenai.org

Checking sign-in…

Loading comments…