SPAM: Spike-Bewuste Adam met Momentum Reset voor Stabiele LLM Training

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

January 12, 2025

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

cs.AI

Samenvatting

Grote Taalmodellen (LLM's) hebben uitzonderlijke prestaties getoond bij diverse taken, maar hun training blijft zeer resource-intensief en vatbaar voor kritieke uitdagingen zoals trainingsinstabiliteit. Een belangrijke bron van deze instabiliteit komt voort uit gradient- en verliespieken, die het leerproces verstoren en vaak leiden tot kostbare interventies zoals het herstellen van checkpoints en het opnieuw starten van experimenten, waardoor inefficiënties verder worden versterkt. Dit artikel presenteert een uitgebreid onderzoek naar gradientpieken die worden waargenomen tijdens de training van LLM's, waarbij hun voorkomen wordt onthuld in meerdere architecturen en datasets. Onze analyse toont aan dat deze pieken tot wel 1000 keer groter kunnen zijn dan typische gradients, waardoor de prestaties van het model aanzienlijk verslechteren. Om dit probleem aan te pakken, stellen we Spike-Aware Adam met Momentum Reset SPAM voor, een nieuw optimizer ontworpen om gradientpieken tegen te gaan door middel van momentumreset en spike-aware gradientclipping. Uitgebreide experimenten, inclusief zowel pre-training als fine-tuning, tonen aan dat SPAM consequent Adam en zijn varianten overtreft bij verschillende taken, waaronder (1) LLM pre-training van 60M tot 1B, (2) 4-bit LLM pre-training, (3) reinforcement learning, en (4) Time Series Forecasting. Daarnaast maakt SPAM geheugenefficiënte training mogelijk door het inschakelen van schaarse momentum, waarbij slechts een subset van momentumtermen wordt onderhouden en bijgewerkt. Bij het werken onder geheugenbeperkingen presteert SPAM beter dan geavanceerde geheugenefficiënte optimizers zoals GaLore en Adam-Mini. Ons werk benadrukt het belang van het verminderen van gradientpieken in LLM-training en introduceert een effectieve optimalisatiestrategie die zowel de trainingsstabiliteit als de resource-efficiëntie op schaal verbetert. De code is beschikbaar op 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