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

要約

大規模な言語モデル（LLMS）は、多様なタスク全体で並外れたパフォーマンスを実証していますが、トレーニングは非常にリソース集中的であり、トレーニングの不安定性などの重要な課題を受けやすいままです。
この不安定性の主な原因は、勾配と損失のスパイクに由来し、学習プロセスを混乱させ、チェックポイントの回復や実験の再開などの費用のかかる介入につながり、非効率性をさらに増幅します。
このホワイトペーパーでは、LLMトレーニング中に観察されたグラデーションスパイクに関する包括的な調査を提示し、複数のアーキテクチャとデータセットにわたってそれらの有病率を明らかにします。
私たちの分析は、これらのスパイクが典型的な勾配よりも最大1000ドルの時間$より大きく、モデルのパフォーマンスが大幅に悪化する可能性があることを示しています。
この問題に対処するために、Momentum Reset SpamであるSpike-Aware Adamを提案します。これは、Momentum ResetおよびSpike-Aware Gradientクリッピングを介してグラデーションスパイクに対抗するように設計された新しいオプティマイザーです。
トレーニング前と微調整の両方を含む広範な実験は、スパムが（1）60Mから1BまでのLLMプレトレーニング、（2）4ビットLLMプレトレーニング、（3）補強学習、（4）時系列を含むさまざまなタスクで一貫してアダムとそのバリアントを上回ることを示しています。
さらに、スパムは、運動量のサブセットのみが維持および更新されるまばらな運動量を有効にすることにより、メモリ効率の高いトレーニングを促進します。
メモリの制約の下で動作する場合、SpamはGaloreやAdam-Miniなどの最先端のメモリ効率の高いオプティマーよりも優れています。
私たちの仕事は、LLMトレーニングにおけるグラデーションスパイクを軽減することの重要性を強調し、大規模なトレーニングの安定性とリソース効率の両方を高める効果的な最適化戦略を導入します。
コードはhttps://github.com/tianjinyellow/spam-optimizer.gitで入手できます

要約(オリジナル)

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

arxiv情報

提供元, 利用サービス

arxiv.jp, Google