---
title: "Efficient MoE Training for Biological Foundation Models"
description: "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."
---

# Efficient MoE Training for Biological Foundation Models

[NVIDIA](https://yomu.fyi/company/nvidia-developer-blog) · Michelle Horton · Sep 24, 2026

**Type:** Tutorial

## 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.

**Tags:** [BioNeMo](https://yomu.fyi/topic/bionemo), [Mixture of Experts](https://yomu.fyi/topic/mixture-of-experts), [NVIDIA](https://yomu.fyi/topic/nvidia), [Transformer Engine](https://yomu.fyi/topic/transformer-engine)

- Source: [NVIDIA](https://developer.nvidia.com/blog/efficient-moe-training-for-biological-foundation-models)
- Source URL: https://developer.nvidia.com/blog/efficient-moe-training-for-biological-foundation-models
- Ingested by Yomu: 2026-09-24T18:00:14.037Z

[Read original post](https://developer.nvidia.com/blog/efficient-moe-training-for-biological-foundation-models)