Accelerating MoE Training for Biological Foundation Models with NVIDIA Transformer Engine
NVIDIA Transformer Engine 加速生物基礎模型 MoE 訓練：吞吐量提升達 2.21 倍
As biological foundation models scale, Mixture-of-Experts (MoE) architectures help expand capacity but suffer from GPU kernel launch overhead and high memory footprints. NVIDIA addresses this via Transformer Engine (TE) by using GroupedLinear to batch expert GEMMs, utilizing MXFP8 block scaling to save memory, and employing the Sequential API to fuse linear projections, SwiGLU, and routing scaling into a single GroupedMLP kernel, yielding up to 2.21x training throughput on B200 GPUs.
Key points
GroupedLinear Reduces Launch Overhead
Replaces naive Python loops with GroupedLinear to execute GEMM operations for multiple experts in a single call, minimizing GPU scheduling overhead.
MXFP8 Low-Precision Block Scaling
Enables hardware-accelerated MXFP8 training on Blackwell GPUs, scaling values in blocks of 32 to preserve numerical accuracy while cutting memory usage.
Fused GroupedMLP Kernel
Fuses quantization, SwiGLU activation, and routing scaling using the TE Sequential API, avoiding intermediate data materialization for peak efficiency.
Significant Throughput Gains
Delivers up to 2.21x the training throughput compared to the Hugging Face baseline in Mixtral-8x7B benchmarks on eight B200 GPUs.
How it works
| Hugging Face Baseline | NVIDIA Transformer Engine (TE) | |
|---|---|---|
| Expert Execution | Python 迴圈依序啟動 (Naive loop) | GroupedLinear 批次提交 (Grouped GEMM) |
| Data Precision | BF16 (16位元) | MXFP8 (8位元區塊縮放，適用於 Blackwell) |
| Kernel Optimization | 無融合，產生大量過渡資料 | 融合 GroupedMLP (融合量化、SwiGLU 與路由) |
| Training Throughput | 1.0x (基準點) | 高達 2.21x (在 8 顆 B200 GPU 上) |
Why it matters
Biological foundation models for genomics and proteomics require long sequences and massive parameter scales, making dense training prohibitively expensive. While MoE architectures save compute, they suffer from GPU underutilization due to expert routing overhead. NVIDIA's optimized pipeline allows life sciences researchers to train massive models on Blackwell GPUs efficiently, accelerating genomic research and drug discovery timelines.
Who it affects
- AI Developer
- AI Researcher
- Enterprise Leader
How to use it
- 1Training large-scale biological foundation models for genomics or proteomics.
- 2Training multi-billion parameter Mixtral MoE models using Expert Parallelism (EP) and FSDP under memory constraints.
Limitations & caveats
- Utilizing the fused MXFP8 GroupedMLP kernel requires specific NVIDIA Blackwell GPUs to leverage hardware acceleration.
- Training requires at least two GPUs to implement expert parallelism (EP), making it unsuitable for single-GPU development setups.
Related
Minimally Invasive Steering of LMs: Optimizing Rewards Without Quality Degradation
微創型語言模型導向技術：利用 MISVO 在不損害生成品質下優化輸出
This paper introduces MISVO, a minimally invasive steering method that uses local KL geometry to optimize LLM outputs for test-time rewards without parameter updates or quality degradation.
LFM2.5-VL-DSpark: Accelerating Vision-Language Models with Minimal Overhead
LFM2.5-VL-DSpark：以超低開銷將多模態模型推論速度提升達 3 倍
Liquid AI has released a 280M parameter DSpark draft model for LFM2.5-VL-3B, boosting decoding speeds up to 3.13 on-device with only 8.9% parameter overhead.
Reproducing OLMo 3 7B Pre-training in MaxText: A Case Study of Large-Scale Training on TPUs
以 MaxText 重現 OLMo 3 7B 預訓練：Google Cloud TPU 大規模訓練實戰指南
This case study details the successful reproduction of AI2's OLMo 3 7B pre-training and mid-training on Cloud TPUs using MaxText, detailing key performance tuning and debugging insights.