SPAM: 安定したLLMトレーニングのためのスパイク感知アダムとモーメンタムリセット

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

January 12, 2025

著者: Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, Shiwei Liu

cs.AI

要旨

大規模言語モデル（LLMs）は、さまざまなタスクで優れたパフォーマンスを示していますが、そのトレーニングは非常にリソース集約型であり、トレーニングの不安定性などの重要な課題に影響を受けやすいです。この不安定性の主要な原因の1つは、勾配と損失の急上昇であり、これらは学習プロセスを妨げ、しばしばコストのかかる介入（チェックポイントの回復や実験の再開始など）を引き起こし、効率をさらに損ないます。本論文では、LLMのトレーニング中に観察される勾配スパイクについて包括的な調査を行い、これらのスパイクが複数のアーキテクチャとデータセット全体にわたって広く存在していることを明らかにします。当社の分析によると、これらのスパイクは通常の勾配の1000倍にもなり、モデルのパフォーマンスを著しく低下させることがあります。この問題に対処するため、勾配スパイクに対抗するためにモーメンタムのリセットとスパイク認識勾配クリッピングを通じて勾配スパイクを緩和することを目的とした新しいオプティマイザ、Spike-Aware Adam with Momentum Reset SPAMを提案します。60Mから1BへのLLMの事前トレーニング、4ビットLLMの事前トレーニング、強化学習、および時系列予測を含むさまざまなタスクで、SPAMがAdamおよびその派生モデルを常に上回ることを示す包括的な実験を行いました。さらに、SPAMは、スパースモーメンタムを可能にすることでメモリ効率のトレーニングを促進し、メモリ制約下で動作する際には、GaLoreやAdam-Miniなどの最先端のメモリ効率の最適化手法を上回ります。私たちの研究は、LLMのトレーニング中の勾配スパイクを緩和することの重要性を強調し、規模の経済性とトレーニングの安定性の両方を高める効果的な最適化戦略を紹介しています。コードはhttps://github.com/TianjinYellow/SPAM-Optimizer.gitで入手可能です。

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 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