使用 MXFP4 训练 LLMs

诸如 MXFP4 等低精度（LP）数据类型可以加速矩阵乘法（GEMMs）并降低训练成本。然而，在训练期间直接使用 MXFP4 代替 BF16 会显著降低模型质量。在本研究中，我们提出了首个使用 MXFP4 GEMMs 的近无损训练方案，该方案在支持的硬件上比 FP8 快 $2\times$。我们的核心见解是使用随机舍入（SR）计算无偏梯度估计，从而得到更准确的模型更新。然而，直接将 SR 应用于 MXFP4 可能会因为块级异常值产生高方差，损害收敛性。为了克服这一点，我们使用随机 Hadamard 变换在理论上界定了 SR 的方差。我们训练了参数量高达 6.7B 的 GPT 模型，并发现我们的方法与混合精度 BF16 训练相比仅产生极小的性能下降。我们的方案在 MXFP4 中计算了 $>1/2$ 的训练 FLOPs，在反向传播期间估计比 FP8 加速 $>1.3\times$，比 BF16 加速 $>1.7\times$。