微缩放格式下训练不稳定性的表征与缓解

训练大型语言模型是一个昂贵且受计算限制的过程，随着模型规模扩大、算法改进和新数据收集，必须重复进行。为解决此问题，下一代硬件加速器越来越多地支持更低精度的算术格式，例如 NVIDIA Blackwell 架构中引入的微缩放（MX）格式。这些格式在参数块内使用共享缩放因子来扩展可表示范围，并以降低的精度执行前向/后向 GEMM 操作以提高效率。在这项工作中，我们研究了块缩放精度格式在模型训练过程中的挑战和可行性。在从头训练的近一千个语言模型中——跨越从 $2 \times 10^{17}$ 到 $4.8 \times 10^{19}$ FLOPs 的计算预算，并覆盖广泛的权重-激活精度组合——我们一致观察到，以 MX 格式训练表现出损失函数中尖锐的、随机的失稳现象，尤其是在较大的计算规模下。为了解释这一现象，我们在一个较小的代理模型上进行了受控实验和消融研究，该模型表现出与语言模型相似的行为，并遍历了架构设置、超参数和精度格式。这些实验启发了一个简单模型，其中由层归一化仿射参数和一小部分激活值的量化引入的乘性梯度偏差可能触发失控性发散。通过在代理模型上进行原位干预实验，我们证明可以通过在训练中途修改精度方案来避免或延迟失稳。在这些发现的指导下，我们在 LLM 设置中评估了稳定策略，并表明某些混合配置可以恢复与全精度训练相竞争的性能。我们在 https://github.com/Hither1/systems-scaling 发布我们的代码。