SPAM：针对 LLM 训练中梯度峰的 Spike-Aware Adam 带动量重置
机器学习
2025-03-03 v2 人工智能
计算与语言
摘要
大型语言模型 (LLM) 在各类任务中展现出卓越的性能，但其训练仍高度依赖资源且易受到诸如训练不稳定等关键挑战的影响。这种不稳定性主要源于梯度与损失峰值，这些峰值会扰乱学习过程，常导致昂贵的检查点恢复和实验重启，从而进一步放大效率损失。本文对 LLM 训练中观察到的梯度峰进行全面调查，揭示其在多种架构和数据集中的普遍性。我们的分析表明，这些峰值可达通常梯度的 ，显著恶化模型性能。为此，我们提出一种新型优化器 Spike-Aware Adam 带动量重置 (SPAM)，通过动量重置和峰值感知梯度裁剪来抵消梯度峰值的影响。广泛实验（包括预训练和微调）表明，SPAM 在各类任务中 consistently 优于 Adam 及其变体，包括 (1) 60M 至 1B 参数的 LLM 预训练，(2) 4-bit LLM 预训练，(3) 强化学习，以及 (4) 时间序列预测。此外，SPAM 通过启用稀疏动量，实现内存高效训练，即只维护和更新动量术语的子集。在受内存限制的情况下，SPAM 超越 GaLore 和 Adam-Mini 等当前最佳内存高效优化器。我们的工作凸显了在 LLM 训练中抑制梯度峰值的重要性，并引入一种提升训练稳定性和规模资源效率的有效优化策略。代码已公开于 https://github.com/TianjinYellow/SPAM-Optimizer.git
引用
@article{arxiv.2501.06842,
title = {SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training},
author = {Tianjin Huang and Ziquan Zhu and Gaojie Jin and Lu Liu and Zhangyang Wang and Shiwei Liu},
journal= {arXiv preprint arXiv:2501.06842},
year = {2025}
}