用MXFP4训练大模型，速度提升1.7倍且损失极小。
Training LLMs with MXFP4
- 通过随机舍入和哈达玛变换降低梯度方差
- 67亿参数GPT模型训练误差几乎与混合精度BF16相当
- 在支持硬件上比FP8快2倍，反向传播加速超1.7倍
低精度数据类型如MXFP4可加速矩阵乘法（GEMMs）并降低训练成本。然而，直接在训练中使用MXFP4替代BF16会显著降低模型质量。本文提出首个接近无损的MXFP4训练方案，其GEMMs运算速度比支持硬件上的FP8快2倍。核心思路是采用随机舍入（SR）计算无偏梯度估计，实现更精准的模型更新。但直接对MXFP4应用SR会产生块级异常值带来的高方差，影响收敛。为此，我们引入随机哈达玛变换，理论证明可有效控制SR方差。我们在高达67亿参数的GPT模型上进行了训练，结果表明该方法与混合精度BF16训练相比仅有微小性能下降。本方案在训练中超过一半的FLOPs以MXFP4完成，使反向传播速度较FP8提升超1.3倍、较BF16提升超1.7倍。
原文摘要 · Abstract (English)
Low precision (LP) datatypes such as MXFP4 can accelerate matrix multiplications (GEMMs) and reduce training costs. However, directly using MXFP4 instead of BF16 during training significantly degrades model quality. In this work, we present the first near-lossless training recipe that uses MXFP4 GEMMs, which are $2\times$ faster than FP8 on supported hardware. Our key insight is to compute unbiased gradient estimates with stochastic rounding (SR), resulting in more accurate model updates. However, directly applying SR to MXFP4 can result in high variance from block-level outliers, harming convergence. To overcome this, we use the random Hadamard tranform to theoretically bound the variance of SR. We train GPT models up to 6.7B parameters and find that our method induces minimal degradation over mixed-precision BF16 training. Our recipe computes $>1/2$ the training FLOPs in MXFP4, enabling an estimated speedup of $>1.3\times$ over FP8 and $>1.7\times$ over BF16 during backpropagation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。