Efficient MoE Training for Biological Foundation Models

NVIDIA published a tutorial showing how its Transformer Engine (TE) and BioNeMo MoE recipe cut the overhead of training mixture-of-experts biological foundation models, using GroupedLinear to submit expert work as one grouped GEMM call instead of a Python loop that launches a separate kernel per expert. The recipe pairs GroupedLinear with MXFP8 quantization to reduce memory use and a fused GroupedMLP kernel that combines quantization, SwiGLU, and routing-weight scaling, and requires at least two GPUs for expert parallelism, with NVIDIA Blackwell GPUs needed for the fused MXFP8 GroupedMLP kernel.

As language models grow, scaling dense architectures becomes increasingly expensive. In a dense transformer, every token passes through every layer, so adding capabilities increases computation for both training and inference. Mixture-of-experts MoE architectures take a different approach to scaling by using many subnetworks, or experts, while activating only a small subset for each token. This tradeoff has made MoE architectures increasingly attractive to the large language model LLM community. They can scale model capacity more efficiently, but the benefits depend heavily on implementation. Fragmented expert computation can reduce GPU utilization. Routing adds communication overhead, and larger parameter footprints create memory and distributed-training challenges. NVIDIA Transformer Engine TE helps address these bottlenecks with optimized primitives for grouped expert computation, kernel fusion, and low-precision training. As biological foundation models grow in parameter count and sequence length, these primitives can improve GPU efficiency while expanding model capacity. This tutorial shows how to put these techniques into practice with the NVIDIA BioNeMo MoE recipe and TE. You’ll see how GroupedLinear improves expert computation, MXFP8 reduces memory use, and the GroupedMLP kernel fuses quantization, SwiGLU, and routing-weight scaling. Together, these capabilities provide a practical reference for efficiently training MoE-based biological foundation models. Prerequisites Before you begin, you need: - Familiarity with Python, PyTorch, and distributed training concepts - An NVIDIA CUDA-enabled environment—you can use the linked Dockerfile or install the recipe requirements - At least two GPUs for expert parallelism; NVIDIA Blackwell GPUs are required to use the fused MXFP8 GroupedMLP kernel Challenge 1: Fragmented expert kernels MoE models replace a single dense feed-forward block with multiple expert networks. A naive implementation, however, can trigger excessive kernel launches. For example, the Hugging Face baseline implementation iterates over all experts in a Python loop, with each expert triggering separate kernel launches. for expert idx, expert layer in enumerate self.experts : idx, top x = torch.where expert mask expert idx current state = hidden states None, top x .reshape -1, hidden dim current hidden = expert layer current state routing weights top x, idx, None final hidden states.index add 0, top x, current hidden Grouped execution preserves the individual expert matrices but submits their work together. TE’s GroupedLinear applies multiple linear transformations in one call by gathering the expert weights and input tokens. Because each expert can receive a different number of tokens, GroupedLinear accepts per-expert token counts split sizes . It submits the local experts through the TE grouped GEMM path instead of launching one PyTorch Linear operation per expert, reducing launch and scheduling overhead. Use GroupedLinear as follows. Each expert retains its own weight tensor weight0 , weight1 , and so on , and the call accepts the per-expert token counts as an additional positional argument: python from transformer engine.pytorch.ops import GroupedLinear experts gate up = GroupedLinear num groups=num local experts, in features=hidden size, out features=2 intermediate size, bias=False, dtype=torch.bfloat16, device="cuda", gate up output = experts gate up tokens, split sizes Compared with the Python loop, this approach submits the gate-up projections as one grouped operation instead of multiple separate calls. Hugging Face Transformers also provides grouped mm . However, TE can fuse GroupedLinear with MXFP8 quantization, activation, routing-weight scaling, and intermediate data movement into a GroupedMLP kernel, as shown in the later sections. Challenge 2: Large model size and activation memory MoE architectures increase total parameter capacity, and genomics workloads often use long sequences, which puts pressure on activation memory during training. BF16 uses 16 bits to represent each model weight and activation. The BioNeMo recipe uses TE to support FP8 and MXFP8 training, reducing memory use. Both formats represent weight and activation values with 8 bits instead of 16. The main difference between FP8 and MXFP8 is scaling granularity: MXFP8 assigns a scaling factor to each block of 32 consecutive values, helping preserve numerical range and accuracy. On NVIDIA Blackwell GPUs, MXFP8 is hardware-accelerated, enabling MXFP8 GEMMs to use specialized Tensor Core instructions. For details about MXFP8 and block scaling, see the Transformer Engine FP8 primer https://nvidia.github.io/TransformerEngine/examples/fp8 primer.html . Challenge 3: Quantization overhead in low-precision training Although most training computation uses 8-bit precision, the model retains its master weights in 16 bits. The training framework therefore adds quantization and dequantization steps to convert between formats. Quantization converts BF16 weights and activations to MXFP8 before the low-precision GEMM; dequantization converts the result back to the higher-precision format. A naive path performs these steps as separate operations, motivating the fused MLP path described next. fp8 recipe = te recipe.MXFP8BlockScaling model = TEMixtralMXFP8ForCausalLM config, fp8 recipe=fp8 recipe, dispatcher=dispatcher The TE autocast API enables MXFP8 precision for the model’s forward and backward passes: with te.autocast enabled=True, recipe=self. fp8 recipe : for decoder layer in self.layers: hidden states = decoder layer hidden states See the BioNeMo recipe https://github.com/NVIDIA-BioNeMo/bionemo-recipes/tree/main/recipes/mixtral native te for the complete code. To use the fused MLP, import the Transformer Engine Sequential API to chain together gate up , ScaledSwiGLU , and down . The API also folds dequantization into the fused path. ScaledSwiGLU combines the routing probabilities “scales” with the expert feed-forward network computations. python from transformer engine.pytorch.ops import GroupedLinear, ScaledSwiGLU, Sequential experts ffn = Sequential GroupedLinear gate up , ScaledSwiGLU , GroupedLinear down The TE Sequential API scans the operations and, when the pattern matches, replaces the GroupedLinear → ScaledSwiGLU → GroupedLinear sequence with a fused operation object: ForwardGroupedMLP CuTeGEMMSwiGLU MXFP8 for the forward pass and a matching fused backward operation. This reduces framework overhead, fuses the SwiGLU and probability-scaling work into the grouped MLP path, and avoids materializing some intermediate results. Results These are several of the optimizations in the BioNeMo recipe. In our training benchmark on eight NVIDIA B200 Tensor Core GPUs, the recipe delivered up to 2.21x the throughput of the Hugging Face baseline. Run the recipe Start with the two-GPU L0 sanity configuration to confirm that expert parallelism and the training environment work correctly: 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 Select BF16 or MXFP8 based on your GPU and memory requirements and set the data-parallel and expert-parallel sizes so their product equals the total GPU count. The recipe README https://github.com/NVIDIA-BioNeMo/bionemo-recipes/tree/main/recipes/mixtral native te includes launch, checkpoint, and benchmark commands. Try the Mixtral Native Transformer Engine recipe https://github.com/NVIDIA-BioNeMo/bionemo-recipes/tree/main/recipes/mixtral native te in BioNeMo Recipes and learn more about the optimized MoE kernels in the NVIDIA Transformer Engine documentation https://docs.nvidia.com/deeplearning/transformer-engine/ . Acknowledgments Sudhakar Singh US mailto:sudhakars@nvidia.com , Varun Thumbe US mailto:vthumbe@nvidia.com , Santosh Santosh US mailto:santosha@nvidia.com , Timur Rvachov US mailto:trvachov@nvidia.com , Chris Hoge US mailto:choge@nvidia.com ,