Practical FP4 Training for Large-Scale MoE Models on Hopper GPUs
Organizations: Zhejiang Lab, Hangzhou, China
Abstract
Training large-scale Mixture-of-Experts (MoE) models is bottlenecked by activation memory and expert-parallel communication, yet FP4 training remains impractical on Hopper-class GPUs without native MXFP4 or NVFP4 support. In this work, we present a training recipe that enables MXFP4 efficiency for MoE models on Hopper architectures without native 4-bit computation support. A central challenge is to integrate FP4 into an existing BF16/FP8 hybrid training pipeline without incurring costly precision round-trips (e.g., FP4 BF16 FP8). We address this challenge by introducing direct FP8-to-FP4 quantization and de-quantization, together with scaling-aware FP4 row-wise to column-wise conversion, enabling FP4 activations and expert-parallel communication with minimal overhead. Core MoE computations are executed in FP8, while activations and expert-parallel communication are compressed using MXFP4, achieving substantial memory and bandwidth savings without degrading convergence. At the 671B parameter scale, our method achieves end-to-end training performance comparable to strong FP8 baselines, while reducing peak activation memory by 14.8% (11.8 GB) and improving training throughput by 12.5%, from 1157 to 1302 tokens per GPU per second. These results show that FP4 efficiency can be practically realized for large-scale MoE training through careful software-hardware co-design, even without native FP4 Tensor Core support.
Figures & tables
| Recomputation Scope | Baseline (BF16) | Baseline (FP8) | Ours (MXFP4) | |||
|---|---|---|---|---|---|---|
| TGS | Mem | TGS | Mem | TGS | Mem | |
| Attn + LN + MoE Experts | 1122 | 68.29% | 1157 | 74.20% | 1156 | 59.40% |
| Attn + LN + MLP | OOM | - | OOM | - | 1248 | 60.62% |
| MLA Up Proj Only | OOM | - | OOM | - | 1302 | 70.11% |
| Recomputation Scope | Baseline (BF16) | Baseline (FP8) | Ours (MXFP4) | |||
|---|---|---|---|---|---|---|
| TGS | Mem | TGS | Mem | TGS | Mem | |
| Attn + LN + MoE Experts | 1725.16 | 53.22% | 1621.96 | 56.81% | 1608.68 | 49.96% |
| Attn + LN + MLP | 1838.56 | 60.80% | 1746.20 | 57.03% | 1616.90 | 49.85% |
| MLP Up Proj Only | OOM | - | 1831.20 | 66.13% | 1722.98 | 58.62% |