SPAM: Spike-bewusstes Adam mit Momentenrücksetzung für stabiles LLM-Training

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

January 12, 2025

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

cs.AI

Zusammenfassung

Große Sprachmodelle (LLMs) haben eine außergewöhnliche Leistungsfähigkeit über verschiedene Aufgaben hinweg gezeigt, doch ihr Training bleibt äußerst ressourcenintensiv und anfällig für kritische Herausforderungen wie Trainingsinstabilität. Eine vorherrschende Quelle dieser Instabilität liegt in Gradienten- und Verlustspitzen, die den Lernprozess stören und oft zu kostspieligen Eingriffen wie Checkpoint-Wiederherstellung und Experiment-Neustarts führen, was die Ineffizienz weiter verstärkt. Diese Arbeit präsentiert eine umfassende Untersuchung von Gradientenspitzen, die während des Trainings von LLMs beobachtet wurden, und zeigt deren Häufigkeit über verschiedene Architekturen und Datensätze hinweg auf. Unsere Analyse zeigt, dass diese Spitzen um das bis zu 1000-fache größer sein können als typische Gradienten und die Leistung des Modells erheblich beeinträchtigen. Um dieses Problem anzugehen, schlagen wir Spike-Aware Adam mit Momentum Reset SPAM vor, einen neuartigen Optimierer, der Gradientenspitzen durch Momentum-Reset und spike-aware Gradientenbeschränkung entgegenwirkt. Umfangreiche Experimente, einschließlich sowohl Pre-Training als auch Feinabstimmung, zeigen, dass SPAM Adam und seine Varianten über verschiedene Aufgaben hinweg kontinuierlich übertrifft, darunter (1) LLM Pre-Training von 60M bis 1B, (2) 4-Bit LLM Pre-Training, (3) Reinforcement Learning und (4) Zeitreihenvorhersage. Darüber hinaus ermöglicht SPAM ein speichereffizientes Training, indem es spärliches Momentum ermöglicht, bei dem nur eine Teilmenge von Momentum-Termen beibehalten und aktualisiert wird. Bei Betrieb unter Speicherbeschränkungen übertrifft SPAM optimierer mit speichereffizienten Optimierern wie GaLore und Adam-Mini. Unsere Arbeit betont die Bedeutung der Reduzierung von Gradientenspitzen im LLM-Training und führt eine effektive Optimierungsstrategie ein, die sowohl die Trainingsstabilität als auch die Ressourceneffizienz im großen Maßstab verbessert. Der Code ist verfügbar unter 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