NVIDIA describes how Transformer Engine (TE) primitives and the BioNeMo Mixture-of-Experts (MoE) recipe can be used to accelerate training of MoE-based biological foundation models. Techniques such as GroupedLinear, MXFP8 low-precision formats, and a fused GroupedMLP kernel reduce kernel-launch overhead, lower memory use, and improve GPU utilization.
Dense vs. MoE scaling
In a dense transformer every token passes through every layer, so increasing a model’s capabilities increases computation proportionally for both training and inference. MoE architectures instead maintain many subnetworks (experts) and activate only a small subset per token, enabling more efficient capacity scaling. Realizing those benefits in practice depends on implementation: fragmented expert computation can lower GPU utilization, routing adds communication overhead, and larger parameter footprints create memory and distributed-training challenges.
How Transformer Engine helps
Transformer Engine (TE) provides optimized primitives for grouped expert computation, kernel fusion, and low-precision training to address these bottlenecks. These primitives are particularly relevant for biological foundation models that increase in parameter count and sequence length, since they can improve GPU efficiency while allowing expanded model capacity.
Key technical approaches in the recipe
1) Fragmented expert kernels
Naive MoE implementations may iterate over experts in Python, issuing a separate kernel launch per expert. TE’s GroupedLinear primitive instead performs multiple linear transforms in a single call by gathering expert weights and input tokens. Because experts may receive different numbers of tokens, GroupedLinear accepts per-expert token counts (split_sizes) and submits local experts via a grouped GEMM path rather than starting a PyTorch Linear operation for each expert, reducing launch and scheduling overhead.
2) Model size and activation memory
MoE increases total parameter capacity and genomics workloads can use long sequences, which increases activation memory pressure. The BioNeMo recipe supports FP8 and MXFP8 training using TE to reduce memory usage: both formats use 8 bits per value instead of BF16’s 16 bits. MXFP8 differs by assigning a scaling factor to each block of 32 consecutive values, preserving numerical range and accuracy. On NVIDIA Blackwell GPUs, MXFP8 is hardware-accelerated so MXFP8 GEMMs can use specialized Tensor Core instructions.
3) Quantization overhead in low-precision training
Most computation can run at 8-bit precision while keeping master weights in 16-bit format, which requires quantization and dequantization steps. A naive implementation runs these steps as separate operations. TE’s fused MLP path collapses GroupedLinear → ScaledSwiGLU → GroupedLinear into a single fused operation (ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 for the forward pass and a corresponding fused backward op), which fuses the SwiGLU activation and routing-weight scaling and avoids materializing some intermediates.
Results
In the authors’ training benchmark on eight NVIDIA B200 Tensor Core GPUs, the BioNeMo recipe delivered up to 2.21× the throughput of a Hugging Face baseline for Mixtral-8x7B training.
How to run the recipe
Start with the two-GPU L0_sanity configuration to validate expert parallelism and the environment:
torchrun --nproc_per_node=2 train_fsdp2_ep.py --config-name L0_sanity
After validation, scale to the Mixtral-8x7B configuration with expert parallelism (EP=8) and MXFP8 precision across eight GPUs:
torchrun --nproc_per_node=8 train_fsdp2_ep.py --config-name L1_8x7B_ep checkpoint.ckpt_dir=/path/to/ckpt
Choose BF16 or MXFP8 depending on GPU and memory needs. Set data-parallel and expert-parallel sizes so their product equals the total GPU count. The recipe README includes launch, checkpoint, and benchmark commands.
Prerequisites
- Familiarity with Python, PyTorch, and distributed training concepts
- NVIDIA CUDA-enabled environment (the linked Dockerfile can be used or the recipe requirements installed)
- At least two GPUs for expert parallelism; NVIDIA Blackwell GPUs are required to use the fused MXFP8 GroupedMLP kernel
Acknowledgments
Sudhakar Singh US, Varun Thumbe US, Santosh Santosh US, Timur Rvachov US, Chris Hoge US
