通过转置不变块量化进行稳定的 FP4 训练

降低训练精度是提高 大语言模型 (LLM) 训练效率的关键杠杆，但由于优化过程中的不稳定，超越 FP8 到 4 位浮点 (FP4) 仍然具有挑战性。我们确定了现有微尺度方法中这种不稳定性的根本根源：张量转置引起的尺度不一致。在传统的一维块量化中，前向和后向传递在转置后将不同的缩放因子分配给相同的值，导致梯度更新有偏差和不稳定。为了解决这个问题，我们提出了一种基于 2D 块 FP4 量化的低精度训练框架，该框架强制转置不变缩放并保持前向和后向计算之间的一致性。我们进一步将其与无截断缩放和随机舍入相结合，以控制量化误差并保持无偏梯度。为了处理注意力机制的敏感性，我们对查询和关键投影采用 MXFP8 量化，从而产生实用的混合精度设计。我们在最多 7B 个参数的密集 LLM 和一个 30B 专家混合模型上评估我们的方法，并在最多 100B 个词元上进行训练。在所有设置中，我们的方法实现了稳定的端到端 FP4 训练，并且与 BF16 性能紧密匹配，困惑度和下游精度下降不到 1.3%。这些结果表明，强制向前向后扩展一致性足以实现大规模的实际 FP4 训练，为更高效的 LLM 训练提供简单有效的途径。