NVIDIA Shows an Efficient MoE Training Path for Biological Foundation Models

NVIDIA's developer blog published a tutorial on training mixture-of-experts (MoE) biological foundation models more efficiently. It is credited to Faradawn Yang, Peter St John, Kyle Tretina and Zoey Zhang.

Dense transformers send every token through every layer, so extra capacity is costly in training and inference alike. MoE splits the feed-forward block into many experts and activates only a few per token. The savings depend on implementation: fragmented expert computation lowers GPU utilization, routing adds communication overhead, and bigger parameter footprints strain memory and distributed training.

Grouped experts instead of a Python loop

The first bottleneck is kernel fragmentation. In the Hugging Face baseline, experts run one after another in a Python loop, each triggering separate kernel launches. TE's GroupedLinear gathers expert weights and tokens and applies all linear transformations in one grouped GEMM call, taking per-expert token counts through a split_sizes argument. NVIDIA's stack can fuse that projection with MXFP8 quantization and routing-weight scaling in a GroupedMLP kernel.

Memory and quantization overhead

MoE checkpoints carry more parameters, and genomics workloads often use long sequences, which presses on activation memory. BF16 stores each value in 16 bits; FP8 and MXFP8 use 8. MXFP8 adds a scaling factor per block of 32 consecutive values to protect numerical range, and on Blackwell GPUs those GEMMs run on specialized Tensor Core instructions.

Low-precision training still keeps master weights in 16 bits, so quantization and dequantization add work. Chaining GroupedLinear, ScaledSwiGLU and GroupedLinear through TE's Sequential API allows the library to swap the pattern for one fused operation, ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 forward and a matching backward pass.

2.21x throughput on eight B200 GPUs

On a benchmark with Mixtral-8x7B across eight NVIDIA B200 Tensor Core GPUs, the recipe reached 2.21x the throughput of the Hugging Face baseline. Fused MXFP8 GroupedMLP kernels need Blackwell hardware, and expert parallelism needs at least two GPUs.

The tutorial includes launch commands: a two-GPU L0_sanity run to validate the environment, then L1_8x7B_ep with expert parallelism (EP=8) and MXFP8 precision across eight GPUs.