SPAM: Spike-Aware Adam con Reset di Momentum per un Addestramento Stabile di LLM

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

January 12, 2025

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni (LLM) hanno dimostrato prestazioni eccezionali in una vasta gamma di compiti, tuttavia il loro addestramento rimane estremamente intensivo in termini di risorse ed è soggetto a sfide critiche come l'instabilità dell'addestramento. Una fonte predominante di questa instabilità deriva da picchi nei gradienti e nella perdita, che interrompono il processo di apprendimento, portando spesso a interventi costosi come il ripristino dei checkpoint e il riavvio degli esperimenti, amplificando ulteriormente le inefficienze. Questo articolo presenta un'indagine esaustiva sui picchi nei gradienti osservati durante l'addestramento dei LLM, rivelando la loro diffusione attraverso diverse architetture e set di dati. La nostra analisi mostra che questi picchi possono essere fino a 1000 volte più grandi dei gradienti tipici, deteriorando notevolmente le prestazioni del modello. Per affrontare questo problema, proponiamo Spike-Aware Adam with Momentum Reset SPAM, un nuovo ottimizzatore progettato per contrastare i picchi nei gradienti attraverso il reset del momento e il ritaglio dei gradienti consapevole dei picchi. Esperimenti estesi, inclusi sia il pre-addestramento che il raffinamento, dimostrano che SPAM supera costantemente Adam e le sue varianti in vari compiti, tra cui (1) pre-addestramento dei LLM da 60M a 1B, (2) pre-addestramento dei LLM a 4 bit, (3) apprendimento per rinforzo e (4) previsione delle serie temporali. Inoltre, SPAM facilita l'addestramento efficiente in termini di memoria consentendo un momento sparso, in cui solo un sottoinsieme dei termini di momento viene mantenuto e aggiornato. Quando opera sotto vincoli di memoria, SPAM supera ottimizzatori efficienti in termini di memoria all'avanguardia come GaLore e Adam-Mini. Il nostro lavoro sottolinea l'importanza di mitigare i picchi nei gradienti nell'addestramento dei LLM e introduce una strategia di ottimizzazione efficace che migliora sia la stabilità dell'addestramento che l'efficienza delle risorse su larga scala. Il codice è disponibile su 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