title: SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training
summary: 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 1000times 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
paper_authors._id: 678646a9195de4bd4d6f75c2
paper_authors.name: Tianjin Huang
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-14T13:11:11.237Z
paper_authors._id: 64cd4743a785f2043b32915e
paper_authors.avatarUrl: /avatars/ba0b497a194dfea8449112d71fc67654.svg
paper_authors.fullname: Tianjin Huang
paper_authors.type: user
paper_authors.user: TianjinHuang
paper_authors._id: 678646a9195de4bd4d6f75c3
paper_authors.name: Ziquan Zhu
paper_authors._id: 678646a9195de4bd4d6f75c4
paper_authors.name: Gaojie Jin
paper_authors._id: 678646a9195de4bd4d6f75c5
paper_authors.name: Lu Liu
paper_authors._id: 678646a9195de4bd4d6f75c6
paper_authors.name: Zhangyang Wang
paper_authors._id: 678646a9195de4bd4d6f75c7
paper_authors.name: Shiwei Liu
paper_authors.status: admin_assigned
paper_authors.statusLastChangedAt: 2025-01-14T14:53:45.458Z
paper_authors._id: 65b04d2291e63920a7898c9e
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/65b04d2291e63920a7898c9e/iUHs235G4bqK-KnH_94ti.jpeg
paper_authors.fullname: Liu
paper_authors.type: user
paper_authors.user: Shiweiliuiiiiiii
paper_title: SPAM: Spike-Aware Adam with Momentum Reset for Stable LLM Training
paper_summary: 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 1000times 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
paper_ai_summary: SPAM optimizer addresses training instability in LLMs by mitigating gradient spikes, enhancing performance and efficiency across various tasks.
paper_ai_keywords: LLMs gradient spikes loss spikes Spike-Aware Adam with Momentum Reset SPAM momentum reset spike-aware gradient clipping memory-efficient training sparse momentum GaLore Adam-Mini