SPAM: Spike-Aware Adam con Reinicio de Momento para un Entrenamiento Estable de LLM

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

January 12, 2025

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

cs.AI

Resumen

Los Modelos de Lenguaje de Gran Tamaño (LLMs) han demostrado un rendimiento excepcional en diversas tareas, sin embargo, su entrenamiento sigue siendo altamente intensivo en recursos y susceptible a desafíos críticos como la inestabilidad durante el entrenamiento. Una fuente predominante de esta inestabilidad proviene de los picos de gradiente y pérdida, que interrumpen el proceso de aprendizaje, a menudo resultando en intervenciones costosas como la recuperación de puntos de control y reinicios de experimentos, lo que amplifica aún más las ineficiencias. Este artículo presenta una investigación exhaustiva sobre los picos de gradiente observados durante el entrenamiento de LLM, revelando su prevalencia en múltiples arquitecturas y conjuntos de datos. Nuestro análisis muestra que estos picos pueden ser hasta 1000 veces más grandes que los gradientes típicos, deteriorando sustancialmente el rendimiento del modelo. Para abordar este problema, proponemos Spike-Aware Adam con Momentum Reset SPAM, un optimizador novedoso diseñado para contrarrestar los picos de gradiente mediante el reinicio de momento y la limitación de gradiente consciente de los picos. Experimentos extensos, que incluyen tanto pre-entrenamiento como ajuste fino, demuestran que SPAM supera consistentemente a Adam y sus variantes en diversas tareas, incluido (1) pre-entrenamiento de LLM de 60M a 1B, (2) pre-entrenamiento de LLM de 4 bits, (3) aprendizaje por refuerzo y (4) Pronóstico de Series Temporales. Además, SPAM facilita un entrenamiento eficiente en memoria al permitir un momento disperso, donde solo se mantienen y actualizan un subconjunto de términos de momento. Cuando opera bajo restricciones de memoria, SPAM supera a optimizadores eficientes en memoria de última generación como GaLore y Adam-Mini. Nuestro trabajo subraya la importancia de mitigar los picos de gradiente en el entrenamiento de LLM e introduce una estrategia de optimización efectiva que mejora tanto la estabilidad del entrenamiento como la eficiencia de recursos a gran escala. El código está disponible en 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