NVIDIA team's BioNeMo MoE recipe with Transformer Engine lifts biological foundation model training throughput to 2.21x the Hugging Face baseline on eight B200 GPUs

The tutorial gives a concrete way to replace the per-expert Python loop with GroupedLinear: each expert keeps its own weight tensor, and the call takes per-expert token counts (split_sizes) as an extra positional argument so TE's grouped GEMM path submits the local experts' gate-up projections in one call. Against the Hugging Face baseline shown in the text, which iterates with for expert_idx, expert_layer in enumerate(self.experts) and triggers separate kernel launches per expert, this submits multiple linear transformations as one grouped operation, reducing launch and scheduling overhead. The text includes runnable code for constructing and calling GroupedLinear alongside the baseline loop, so this is an implementation-level comparison rather than a standalone kernel-level timing measurement.

The tutorial describes how the BioNeMo recipe uses TE to support FP8 and MXFP8 training, representing weights and activations with 8 bits instead of BF16's 16 bits, where MXFP8 assigns a scaling factor to each block of 32 consecutive values to help preserve numerical range and accuracy and is hardware-accelerated on NVIDIA Blackwell GPUs using specialized Tensor Core instructions. Relative to BF16 representation, this path targets the activation-memory pressure created by larger MoE parameter counts and the long sequences common in genomics workloads; the text also notes that master weights stay in 16 bits, so the framework adds quantization and dequantization steps. The text shows the MXFP8BlockScaling recipe object and te.autocast usage, and notes that quantization/dequantization are separate operations in a naive path; the memory benefit is explained mechanistically without specific memory figures.

The tutorial shows the TE Sequential API scanning operations and, on a pattern match, replacing the GroupedLinear to ScaledSwiGLU to GroupedLinear sequence with the fused ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 forward op and a matching fused backward op, folding SwiGLU, routing-probability scaling, and dequantization into the grouped MLP path and avoiding materializing some intermediate results. Relative to grouped_mm, which Hugging Face Transformers also provides, the text emphasizes that TE can fuse GroupedLinear with MXFP8 quantization, activation, routing-weight scaling, and intermediate data movement into one GroupedMLP kernel. The text provides the Sequential(GroupedLinear(gate_up), ScaledSwiGLU(), GroupedLinear(down)) code and the fused op names, which is a framework-behavior description rather than a per-operator performance breakdown.

In the training benchmark on eight NVIDIA B200 Tensor Core GPUs, the recipe delivered up to 2.21x the throughput of the Hugging Face baseline; the tutorial also gives launch commands from the two-GPU L0_sanity configuration that confirms expert parallelism to the eight-GPU Mixtral-8x7B configuration with EP=8 and MXFP8 precision. Rather than only describing optimization principles, this pairs an end-to-end throughput multiple with a reproducible launch path and the reminder that data-parallel and expert-parallel sizes must multiply to the total GPU count. The throughput figure comes from the authors' own eight-GPU B200 benchmark; the text does not give baseline configuration details, sequence length, or error ranges, and the fused MXFP8 GroupedMLP kernel requires Blackwell GPUs.

This tutorial is aimed at engineering readers already familiar with Python, PyTorch, and distributed training concepts who have an NVIDIA CUDA-enabled environment; expert parallelism needs at least two GPUs, and using the fused MXFP8 GroupedMLP kernel requires NVIDIA Blackwell GPUs. It offers a reproducible path: first confirm expert parallelism and the training environment with the two-GPU L0_sanity configuration, then scale to the eight-GPU Mixtral-8x7B configuration with EP=8 and MXFP8 precision, choosing BF16 or MXFP8 based on GPU and memory requirements while keeping the product of data-parallel and expert-parallel sizes equal to the total GPU count. The intended setting is MoE-based biological foundation model training, especially long-sequence workloads such as genomics where activation memory is under pressure.

The 2.21x throughput figure comes from the authors' own eight-GPU B200 benchmark, and the text does not state baseline configuration, sequence length, model-scale details, or measurement error ranges, so how well that number reproduces on other hardware and workloads remains an open question. The effect of MXFP8's one scaling factor per 32 consecutive values on training convergence and final accuracy is only described as helping preserve numerical range and accuracy, with no convergence curves or downstream task results given. The tutorial is also primarily mechanistic and code-oriented, without a per-operator performance breakdown, so readers wanting to attribute gains to grouped GEMM, low precision, or the fused kernel would need to run their own ablations.