SPAM: Spike-Aware Adam com Reinício de Momento para Treinamento Estável 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

Resumo

Os Modelos de Linguagem de Grande Escala (LLMs) têm demonstrado um desempenho excepcional em diversas tarefas, no entanto, seu treinamento continua sendo altamente intensivo em recursos e suscetível a desafios críticos, como a instabilidade durante o treinamento. Uma fonte predominante dessa instabilidade advém de picos nos gradientes e perdas, que perturbam o processo de aprendizado, frequentemente resultando em intervenções custosas, como recuperação de pontos de verificação e reinícios de experimentos, amplificando ainda mais as ineficiências. Este artigo apresenta uma investigação abrangente sobre os picos nos gradientes observados durante o treinamento de LLM, revelando sua prevalência em várias arquiteturas e conjuntos de dados. Nossa análise mostra que esses picos podem ser até 1000 vezes maiores do que gradientes típicos, deteriorando substancialmente o desempenho do modelo. Para lidar com esse problema, propomos o Spike-Aware Adam with Momentum Reset SPAM, um otimizador inovador projetado para neutralizar os picos nos gradientes por meio de redefinição de momento e limitação de gradientes cientes dos picos. Experimentos extensivos, incluindo tanto pré-treinamento quanto ajuste fino, demonstram que o SPAM consistentemente supera o Adam e suas variantes em várias tarefas, incluindo (1) pré-treinamento de LLM de 60M a 1B, (2) pré-treinamento de LLM de 4 bits, (3) aprendizado por reforço e (4) Previsão de Séries Temporais. Além disso, o SPAM facilita o treinamento eficiente em memória, permitindo momento esparsos, onde apenas um subconjunto de termos de momento são mantidos e atualizados. Ao operar sob restrições de memória, o SPAM supera otimizadores eficientes em memória de ponta, como GaLore e Adam-Mini. Nosso trabalho destaca a importância de mitigar os picos nos gradientes no treinamento de LLM e introduz uma estratégia de otimização eficaz que aprimora tanto a estabilidade do treinamento quanto a eficiência de recursos em grande escala. O código está disponível em 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