在本机 FP4 硬件上使用 MXFP4 预训练 大语言模型

为什么 大语言模型 的全流程 FP4 训练经常出现分歧，即使前向激活和激活梯度保持稳定？我们通过在 Transformer 训练中对 MXFP4 量化进行受控研究来解决这个问题，逐步在前向传播 (Fprop)、激活梯度 (Dgrad) 和权重梯度 (Wgrad) 上启用 FP4，同时保持所有其他因素固定。在 C4 数据集上对 Llama 3.1-8B 进行全面预训练时，我们观察到量化 Wgrad 是收敛性下降的主要驱动因素，而 Fprop 和 Dgrad 中的 FP4 仅引入了适度的额外词元要求。为了解释这种行为，我们在受控实验环境下评估结构化和随机干预措施。我们发现，一旦 Wgrad 被量化，随机舍入和随机 Hadamard 旋转就无法稳定训练，而确定性 Hadamard 旋转始终能够恢复稳定的优化。这些结果表明 FP4 训练的不稳定性是由敏感梯度路径上的结构化微尺度误差驱动的，而不是由随机性不足驱动的。我们在 AMD Instinct MI355X GPU 上运行原生 MXFP4 支持的实验，从而能够在不依赖软件模拟的情况下对这些影响进行受控研究。