What Is an AI Factory? Power In, Tokens Out
An AI factory is a fixed power and cooling budget converted into model output. This explainer maps its four parts and the throughput-per-watt trade-offs that connect them.
Loading post…
NVIDIAMichelle HortonTutorial
Efficient MoE Training for Biological Foundation Models
Summary
Mixture-of-experts (MoE) architectures scale model capacity more efficiently than dense transformers by activating only a subset of experts per token, but naive implementations suffer from fragmented expert computation, routing overhead, and memory pressure. NVIDIA Transformer Engine (TE) addresses these bottlenecks with optimized primitives including GroupedLinear for batched expert GEMMs, MXFP8 low-precision training, and a fused GroupedMLP kernel that combines quantization, SwiGLU activation, and routing-weight scaling. The BioNeMo recipe implements these techniques for biological foundation models, using TE’s autocast API for MXFP8 precision and the Sequential API to fuse the gate-up, ScaledSwiGLU, and down projections into a single operation. On eight NVIDIA B200 GPUs, the recipe achieved up to 2.21x the training throughput of the Hugging Face baseline for a Mixtral-8x7B model. The tutorial provides prerequisites, code examples, and configuration commands to reproduce the optimizations.
Context
MoE architectures scale model capacity efficiently by activating only a subset of experts per token, but naive implementations suffer from fragmented expert computation, routing overhead, and memory pressure from large parameter footprints and long sequences in biological foundation models.
Approach / What changed
Use NVIDIA Transformer Engine’s GroupedLinear to batch expert GEMMs, MXFP8 low-precision training to reduce memory, and the fused GroupedMLP kernel via the Sequential API to combine quantization, SwiGLU activation, and routing-weight scaling into a single operation. The BioNeMo recipe implements these optimizations for biological foundation models.
Takeaways
GroupedLinear in Transformer Engine batches multiple expert linear transformations into one grouped GEMM call, reducing kernel launch overhead compared to iterating over experts in a Python loop.
MXFP8 training uses 8-bit values with block scaling (32 consecutive values per scaling factor) to reduce memory use while preserving numerical range, and is hardware-accelerated on NVIDIA Blackwell GPUs.
The fused GroupedMLP kernel in TE’s Sequential API combines gate-up projection, ScaledSwiGLU activation, and down projection into a single forward and backward operation, reducing framework overhead and intermediate materialization.
FAQ
How does GroupedLinear improve MoE training efficiency?
GroupedLinear accepts per-expert token counts and submits all local expert GEMMs as one grouped operation instead of launching separate PyTorch Linear calls per expert. This reduces kernel launch and scheduling overhead, improving GPU utilization.
What is the difference between FP8 and MXFP8 in Transformer Engine?
The main difference is scaling granularity: FP8 uses per-tensor scaling, while MXFP8 assigns a scaling factor to each block of 32 consecutive values. MXFP8 helps preserve numerical range and accuracy and is hardware-accelerated on NVIDIA Blackwell GPUs with specialized Tensor Core instructions.
What throughput improvement did the BioNeMo recipe achieve over the Hugging Face baseline?
In a training benchmark on eight NVIDIA B200 Tensor Core GPUs, the recipe delivered up to 2.21x the throughput of the Hugging Face baseline for a Mixtral-8x7B model.
Source and freshness
- Source
- developer.nvidia.com
- Published
- Added to Yomu
Summaries are generated from the original by a language model. Check the original before relying on it.
See the idea
Visual explainers
How Blackwell Serving Turns Peak FLOPs Into Cheap Tokens
Blackwell inference is a pipeline, not a chip upgrade. This explainer walks five stages — fleet validation, batching, KV caching, low-precision kernels, and speculative decoding — that turn peak FLOPs into cheap tokens.
Keep reading
Related reading
The Good Stuff
Paid picks we use ourselves. Links carry a referral code — a subscription through one earns us a commission at no extra cost to you.
The best engineering reads, once a week
One useful email. Unsubscribe anytime.