提出新优化器SPAM，解决大模型训练中梯度突增导致的不稳定问题。
SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training
- 通过动量重置和突增感知裁剪，动态应对梯度突增
- 在60M到1B参数模型上训练更稳定，性能优于Adam等变体
- 支持稀疏动量存储，内存受限时表现超越GaLore等先进方法
大型语言模型（LLMs）在各类任务中表现出色，但其训练过程仍高度依赖资源且易受训练不稳定性影响。本文深入分析了大模型训练中的梯度与损失突增现象，发现其可高达典型梯度的1000倍，严重损害模型性能。为此，我们提出新型优化器SPAM：Spike-Aware Adam with Momentum Reset，通过动量重置与突增感知梯度裁剪有效缓解该问题。大量实验表明，SPAM在多种场景下均优于Adam及其变体，涵盖从60M到1B参数的LLM预训练、4比特低精度预训练、强化学习及时间序列预测。此外，SPAM支持稀疏动量机制，在内存受限条件下表现优于当前最优的内存高效优化器如GaLore和Adam-Mini。本工作强调了抑制梯度突增的重要性，并提供了一种兼顾训练稳定性与资源效率的优化策略。代码已开源：https://github.com/TianjinYellow/SPAM-Optimizer.git
原文摘要 · Abstract (English)
Large Language Models (LLMs) have demonstrated exceptional performance across diverse tasks, yet their training remains highly resource-intensive and susceptible to critical challenges such as training instability. A predominant source of this instability stems from gradient and loss spikes, which disrupt the learning process, often leading to costly interventions like checkpoint recovery and experiment restarts, further amplifying inefficiencies. This paper presents a comprehensive investigation into gradient spikes observed during LLM training, revealing their prevalence across multiple architectures and datasets. Our analysis shows that these spikes can be up to $1000\times$ larger than typical gradients, substantially deteriorating model performance. To address this issue, we propose Spike-Aware Adam with Momentum Reset SPAM, a novel optimizer designed to counteract gradient spikes through momentum reset and spike-aware gradient clipping. Extensive experiments, including both pre-training and fine-tuning, demonstrate that SPAM consistently surpasses Adam and its variants across various tasks, including (1) LLM pre-training from 60M to 1B, (2) 4-bit LLM pre-training,(3) reinforcement learning, and (4) Time Series Forecasting. Additionally, SPAM facilitates memory-efficient training by enabling sparse momentum, where only a subset of momentum terms are maintained and updated. When operating under memory constraints, SPAM outperforms state-of-the-art memory-efficient optimizers such as GaLore and Adam-Mini. Our work underscores the importance of mitigating gradient spikes in LLM training and introduces an effective optimization strategy that enhances both training stability and resource efficiency at scale. Code is available at https://github.com/TianjinYellow/SPAM-Optimizer.git
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。