{"slug": "efficient-moe-training-for-biological-foundation-models", "title": "Efficient MoE Training for Biological Foundation Models", "summary": "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.", "body_md": "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.\n\nMixture-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.\n\nThis 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.\n\nThis 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.\n\n## Prerequisites\n\nBefore you begin, you need:\n\n- Familiarity with Python, PyTorch, and distributed training concepts\n- An NVIDIA CUDA-enabled environment—you can use the linked Dockerfile or install the recipe requirements\n- At least two GPUs for expert parallelism; NVIDIA Blackwell GPUs are required to use the fused MXFP8 GroupedMLP kernel\n\n## Challenge 1: Fragmented expert kernels\n\nMoE 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.\n\n```\nfor expert_idx, expert_layer in enumerate(self.experts):\n    idx, top_x = torch.where(expert_mask[expert_idx])\n    current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)\n    current_hidden = expert_layer(current_state) * routing_weights[top_x, idx, None]\n    final_hidden_states.index_add_(0, top_x, current_hidden)\n```\n\nGrouped 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.\n\nUse 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:\n\n``` python\nfrom transformer_engine.pytorch.ops import GroupedLinear\n\nexperts_gate_up = GroupedLinear(\n    num_groups=num_local_experts,\n    in_features=hidden_size,\n    out_features=2 * intermediate_size,\n    bias=False,\n    dtype=torch.bfloat16,\n    device=\"cuda\",\n)\n\ngate_up_output = experts_gate_up(tokens, split_sizes)\n```\n\nCompared with the Python loop, this approach submits the gate-up projections as one grouped operation instead of multiple separate calls.\n\nHugging 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.\n\n## Challenge 2: Large model size and activation memory\n\nMoE 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.\n\nThe 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).\n\n## Challenge 3: Quantization overhead in low-precision training\n\nAlthough 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.\n\n```\nfp8_recipe = te_recipe.MXFP8BlockScaling()\nmodel = TEMixtralMXFP8ForCausalLM(config, fp8_recipe=fp8_recipe, dispatcher=dispatcher)\n```\n\nThe TE autocast API enables MXFP8 precision for the model’s forward and backward passes:\n\n```\nwith te.autocast(enabled=True, recipe=self._fp8_recipe):\n    for decoder_layer in self.layers:\n        hidden_states = decoder_layer(hidden_states)\n```\n\nSee the [BioNeMo recipe](https://github.com/NVIDIA-BioNeMo/bionemo-recipes/tree/main/recipes/mixtral_native_te) for the complete code.\n\nTo 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.\n\n``` python\nfrom transformer_engine.pytorch.ops import GroupedLinear, ScaledSwiGLU, Sequential\n\nexperts_ffn = Sequential(GroupedLinear(gate_up), ScaledSwiGLU(), GroupedLinear(down))\n```\n\nThe 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.\n\n## Results\n\nThese 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.\n\n## Run the recipe\n\nStart with the two-GPU `L0_sanity` configuration to confirm that expert parallelism and the training environment work correctly:\n\n```\ntorchrun --nproc_per_node=2 train_fsdp2_ep.py --config-name L0_sanity\n```\n\nAfter validation, scale to the Mixtral-8x7B configuration with expert parallelism (EP=8) and MXFP8 precision across eight GPUs:\n\n```\ntorchrun --nproc_per_node=8 train_fsdp2_ep.py --config-name L1_8x7B_ep checkpoint.ckpt_dir=/path/to/ckpt\n```\n\nSelect 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.\n\nTry 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/).\n\n## Acknowledgments \n\n[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),", "url": "https://wpnews.pro/news/efficient-moe-training-for-biological-foundation-models", "canonical_source": "https://developer.nvidia.com/blog/efficient-moe-training-for-biological-foundation-models/", "published_at": "2026-09-24 15:00:00+00:00", "updated_at": "2026-09-24 15:31:24.188926+00:00", "lang": "en", "topics": ["machine-learning", "ai-research", "ai-infrastructure", "large-language-models", "ai-tools"], "entities": ["NVIDIA", "NVIDIA Transformer Engine", "NVIDIA BioNeMo", "GroupedLinear", "GroupedMLP", "MXFP8", "Hugging Face", "PyTorch"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/efficient-moe-training-for-biological-foundation-models", "markdown": "https://wpnews.pro/news/efficient-moe-training-for-biological-foundation-models.md", "text": "https://wpnews.pro/news/efficient-moe-training-for-biological-foundation-models.txt", "jsonld": "https://wpnews.pro/news/efficient-moe-training-for-biological-foundation-models.jsonld"}}