重温 LLM 预训练 中的 Adam-SGD 差距：大有效学习率的作用

人们普遍认为，随机梯度下降 (SGD) 的性能明显低于 预训练 大语言模型 (LLM) 中的 Adam 等自适应优化器。然而造成这种差距的根本原因仍不清楚。在这项工作中，我们将很大一部分差异归因于 SGD 无法维持与 Adam 更大的有效学习率相当的学习率。通过对 LLM 预训练 动力学的实证和理论分析，我们发现训练的特点是梯度范数小和权重梯度比大，这种效应在 预训练 中典型的较大批量大小中变得更加明显，需要如此大的有效学习率。然而，我们发现输出层梯度大小在不同的标记类别中变得非常不均匀，并且在训练期间经常出现大的梯度尖峰。总之，这些影响严重限制了 SGD 可接受的学习率。在这种理解的指导下，我们证明了在大学习率下稳定 SGD 的简单裁剪机制使其能够恢复 Adam 的大部分性能。在我们的大规模实验中，当 预训练 具有 1M 词元批量大小的 1B 参数 LLaMA 模型时，大学习率 SGD 和 Adam 之间的验证损失差距从 50% 以上缩小到仅 3.5% 左右。