BioNeMo Uses Grouped Kernels and MXFP8 to...

Scaling dense transformers increases computation because every token traverses every layer. Mixture-of-experts (MoE) models instead expand capacity through multiple expert subnetworks while selecting only a few for each token. This makes larger models more computationally efficient, provided their implementation handles expert execution, routing and memory effectively.

Fragmented expert workloads can leave GPUs underused, routing introduces communication costs, and additional parameters complicate memory management and distributed training. NVIDIA Transformer Engine (TE) addresses these issues through grouped computation, fused operations and low-precision training. These techniques are relevant to biological foundation models as both parameter counts and sequence lengths grow.

Two block diagrams side by side. A dense Transformer block sends every token through the full feed-forward network, while a sparse MoE block uses a top-2 router to send each token to two of eight experts and combines their outputs as a weighted sum. — Figure 1. Dense transformer block (left) compared with a sparse MoE transformer block (right)

The dense and sparse MoE transformer blocks illustrate the architectural difference. The BioNeMo MoE recipe provides a practical implementation using GroupedLinear for expert execution, MXFP8 for lower memory consumption and a GroupedMLP path that combines quantization, SwiGLU and routing-weight scaling.

Requirements

Working knowledge of Python, PyTorch and distributed training.
A CUDA-enabled NVIDIA environment, configured with the recipe’s Dockerfile or requirements.
At least two GPUs to exercise expert parallelism. The fused MXFP8 GroupedMLP kernel requires NVIDIA Blackwell GPUs.

Grouping fragmented expert workloads

A straightforward MoE implementation replaces one dense feed-forward network with several experts but executes them individually. The Hugging Face baseline uses a Python loop to select each expert’s tokens, run its network, apply routing weights and accumulate the output. That approach generates separate kernel launches for each expert.

TE’s GroupedLinear keeps expert weight matrices separate while submitting their linear transformations together. Its split_sizes argument describes the token count assigned to each expert, accommodating uneven routing. Local experts then run through TE’s grouped GEMM path rather than through individual PyTorch Linear calls.

For a gate-up projection, the recipe configures GroupedLinear with num_groups=num_local_experts, in_features=hidden_size, out_features=2 * intermediate_size, bias=False, dtype=torch.bfloat16 and device="cuda". Each expert retains a weight tensor such as weight0 or weight1, and execution uses experts_gate_up(tokens, split_sizes).

This groups the gate-up projections into one submission, reducing launch and scheduling costs. Hugging Face Transformers also offers grouped_mm; TE additionally supports integration with quantization, activation, routing-weight scaling and intermediate data movement through its fused GroupedMLP path.

Two execution timelines. A Python loop runs eight experts one after another, so time grows with the expert count. A grouped GEMM runs all eight in a single GroupedLinear operation. — Figure 2. Hugging Face’s modeling_mixtral.py runs each expert one at a time in a Python loop, while TE batches all expert GEMMs into one grouped operation

The execution comparison contrasts the per-expert Python loop in Hugging Face’s modeling_mixtral.py with TE’s grouped expert GEMMs.

Reducing memory pressure with MXFP8

MoE increases total parameter capacity, while long genomics sequences can make training activations particularly demanding. The BioNeMo recipe supports FP8 and MXFP8 through TE, using 8-bit representations for weights and activations instead of BF16’s 16-bit values.

MXFP8 differs from FP8 in its scaling granularity: it applies a separate scale to every block of 32 consecutive values to help retain numerical range and accuracy. NVIDIA Blackwell GPUs accelerate MXFP8 GEMMs with specialized Tensor Core instructions.

Bit-layout comparison. BF16 uses 16 bits per value: 1 sign, 8 exponent, 7 mantissa. MXFP8 uses 8 bits in E4M3 format plus one shared E8M0 scale per 32 elements. — Figure 3. BF16 uses 16 bits per value, while MXFP8 uses 8 bits. Source: Transformer Engine FP8 primer

The precision comparison shows the 16-bit BF16 representation alongside MXFP8’s 8-bit values.

Fusing precision conversion and expert computation

Low-precision computation still requires higher-precision master weights. In this recipe, those weights remain 16-bit, so training must convert BF16 weights and activations to MXFP8 before low-precision GEMMs and convert results back afterward. Performing those conversions separately adds overhead.

The recipe selects block scaling with te_recipe.MXFP8BlockScaling() and passes it, together with the configuration and dispatcher, to TEMixtralMXFP8ForCausalLM. A te.autocast(enabled=True, recipe=self._fp8_recipe) context enables MXFP8 for forward and backward computation.

TE’s Sequential API assembles the expert feed-forward network from a gate-up GroupedLinear, ScaledSwiGLU and a down-projection GroupedLinear. ScaledSwiGLU incorporates the routing probabilities into the expert computation, while the fused path also incorporates dequantization.

When TE recognizes this operation sequence, it substitutes ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 for the forward path and a corresponding fused backward operation. This reduces framework overhead, combines activation and probability scaling with grouped computation, and avoids writing some intermediate results.

Two vertical flowcharts. Left, labeled “Unfused MLP”: five sequential boxes reading Quantize, Gate Up, SwiGLU, Dequantize, and Gate Down. Right, labeled “Fused MLP”: three boxes reading Quantize, a highlighted Fused Group MLP box that replaces the gate-up, SwiGLU, and dequantize steps, and Gate Down. — Figure 4. The MXFP8 path fuses multiple operations into one kernel before the down projection

The MXFP8 execution diagram shows operations combined into a kernel ahead of the down projection.

Training benchmark

In a Mixtral-8x7B training benchmark using eight NVIDIA B200 Tensor Core GPUs, the BioNeMo recipe achieved up to 2.21x the throughput of the Hugging Face baseline. Grouped execution, low precision and fusion are among the recipe’s optimizations.

A bar chart titled “Mixtral-8x7B Training Throughput on 8×B200,” measuring tokens per second per GPU. Three bars: Hugging Face with BF16 and FSDP8 at 4,096; Transformer Engine with BF16 and FSDP8 at 4,447; and Transformer Engine with expert parallelism and MXFP8 at 9,050, labeled 2.21x. — Figure 5. Mixtral-8x7B training throughput on eight NVIDIA B200 Tensor Core GPUs, showing up to 2.21x the throughput of the Hugging Face baseline

The throughput comparison reports the eight-GPU Mixtral-8x7B result.

Launching the recipe

First validate the environment and expert parallelism with the two-GPU L0_sanity configuration:

torchrun --nproc_per_node=2 train_fsdp2_ep.py --config-name L0_sanity

After that check, launch Mixtral-8x7B across eight GPUs with MXFP8 and expert parallelism set to EP=8:

torchrun --nproc_per_node=8 train_fsdp2_ep.py --config-name L1_8x7B_ep checkpoint.ckpt_dir=/path/to/ckpt

Choose BF16 or MXFP8 according to available hardware and memory needs. The data-parallel size multiplied by the expert-parallel size must equal the GPU count. The Mixtral Native Transformer Engine recipe in BioNeMo Recipes includes README instructions for launches, checkpoints and benchmarks; TE documentation covers the optimized MoE kernels.

Contributors

Acknowledged contributors include Sudhakar Singh US, Varun Thumbe US, Santosh Santosh US, Timur Rvachov US and Chris Hoge US.

Faradawn Yang works on NVIDIA’s AI platform software team, focusing on inference products. He earned a master’s in computer science and a bachelor’s in mathematics and computer science at the University of Chicago, and previously worked in data engineering at a marketing measurement company.

Peter St John is a machine learning engineer on NVIDIA’s BioNeMo team, accelerating protein and genomics workflows. Before joining NVIDIA in 2022, he spent 7 years as a project lead at the National Renewable Energy Lab researching fuel chemistries and microbial biomass conversion. He earned his chemical engineering PhD at UCSB in 2015.

Kyle Tretina leads product marketing work at NVIDIA for digital biology and drug discovery, including BioNeMo and BioPharma initiatives. With a PhD in molecular microbiology and immunology, his work connects AI, chemistry and biology with platforms for molecular and protein design.

Zoey Zhang manages AI training products for Digital Biology at NVIDIA. Her experience includes software engineering and machine learning research. She studied Biomedical Engineering at the University of Waterloo, specializing in Medical AI and Computing, and focuses on applying AI and accelerated computing to scientific discovery and treatment development.