SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training

大型语言模型（LLMs）在各种任务中表现出色，但其训练仍然非常耗费资源，并且容易受到诸如训练不稳定等关键挑战的影响。这种不稳定的主要来源是梯度和损失的突增，这些现象会扰乱学习过程，通常导致代价高昂的干预措施，如检查点恢复和实验重启，进一步加剧了效率低下。本文对LLM训练过程中观察到的梯度突增进行了全面的研究，揭示了这些突增在多个架构和数据集中的普遍存在。我们的分析表明，这些突增可能比典型的梯度大1000倍，严重损害模型性能。为了解决这一问题，我们提出了带有动量重置的抗突增Adam优化器（Spike-Aware Adam with Momentum Reset, SPAM），这是一种新型优化器，通过动量重置和抗突增梯度裁剪来对抗梯度突增。广泛的实验，包括预训练和微调，证明SPAM在各种任务中始终超越Adam及其变体，包括（1）从6000万到10亿参数的LLM预训练，（2）4位LLM预训练，（3）强化学习，以及（4）时间序列预测。此外，SPAM通过启用稀疏动量来促进内存高效的训练，在这种模式下，仅维护和更新一部分动量项。在内存受限的情况下，SPAM的表现优于最先进的内存高效优化器，如GaLore和Adam-Mini。我们的研究强调了在LLM训练中缓解梯度突增的重要性，并引入了一种有效的优化策略，该策略在大规模训练中增强了稳定性和资源利用效率。代码可在https://github.com/TianjinYellow/SPAM-Optimizer.git获取。

该论文旨在解决大型语言模型（LLM）训练过程中常见的梯度和损失尖峰问题，这些问题导致训练不稳定，增加了实验重启和检查点恢复的成本。尽管梯度尖峰是一个已知问题，但在LLM训练中的具体影响及其广泛性尚未得到充分研究。

论文提出了一种名为Spike-Aware Adam with Momentum Reset (SPAM)的新优化器，通过动量重置和尖峰感知梯度裁剪来对抗梯度尖峰。相比现有的优化器如Adam及其变体，SPAM在多个任务中表现出更好的性能，并且能够在资源受限的环境中实现更高效的训练。

论文通过广泛的实验验证了SPAM的有效性，涵盖了从60M到1B参数的LLM预训练、4-bit LLM预训练、强化学习和时间序列预测等任务。此外，SPAM支持稀疏动量，即只维护和更新部分动量项，从而在内存有限的情况下优于其他内存高效优化器如GaLore和Adam-Mini。作者还提供了开源代码，便于复现和进一步研究。

近年来，关于优化LLM训练稳定性和效率的研究层出不穷。例如，《On the Stability of Large Language Model Training》探讨了训练过程中的稳定性问题；《Gradient Clipping in Deep Learning: A Comprehensive Review》综述了梯度裁剪技术；《Memory-Efficient Optimizers for Large-Scale Models》则专注于内存高效的优化方法。