无需硬件支持,实现大规模MoE模型的高效FP4训练。
Practical FP4 Training for Large-Scale MoE Models on Hopper GPUs
- 通过直接FP8转FP4量化,避免精度转换开销。
- 671B参数模型下内存降14.8%,训练吞吐提升12.5%。
- 适合追求高性价比训练的大模型研究者。
大规模混合专家(MoE)模型训练受限于激活内存和专家并行通信开销,而无原生MXFP4或NVFP4支持时,FP4训练在Hopper架构上仍不实用。本文提出一种训练方案,使无原生4比特计算支持的Hopper GPU可实现高效的MXFP4训练。核心挑战在于将FP4融入现有BF16/FP8混合训练流程,避免代价高昂的精度转换(如FP4↔BF16↔FP8)。我们引入直接的FP8到FP4量化与反量化机制,结合感知缩放的行到列转换,实现低开销的FP4激活与专家并行通信。核心MoE计算仍以FP8执行,激活与通信采用MXFP4压缩,在不损害收敛性的前提下大幅节省内存与带宽。在671B参数规模下,该方法达到与强基线FP8相当的端到端训练性能,峰值激活内存减少14.8%(11.8 GB),训练吞吐从1157提升至1302 tokens/GPU/秒。结果表明,通过软硬件协同设计,即使无原生支持,也能在大规模MoE训练中实现实际的FP4效率。
原文摘要 · Abstract (English)
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 $\leftrightarrow$ BF16 $\leftrightarrow$ 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。