SPAM : Spike-Aware Adam avec Réinitialisation du Momentum pour un Entraînement Stable de LLM

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

Résumé

Les grands modèles de langage (LLMs) ont démontré des performances exceptionnelles dans diverses tâches, mais leur entraînement reste très gourmand en ressources et sujet à des défis critiques tels que l'instabilité de l'entraînement. Une source prédominante de cette instabilité provient des pics de gradient et de perte, qui perturbent le processus d'apprentissage, entraînant souvent des interventions coûteuses telles que la récupération de points de contrôle et le redémarrage des expériences, amplifiant ainsi les inefficacités. Cet article présente une enquête approfondie sur les pics de gradient observés lors de l'entraînement des LLM, révélant leur prévalence à travers plusieurs architectures et ensembles de données. Notre analyse montre que ces pics peuvent être jusqu'à 1000 fois plus grands que les gradients typiques, détériorant considérablement les performances du modèle. Pour résoudre ce problème, nous proposons Spike-Aware Adam avec Momentum Reset SPAM, un nouvel optimiseur conçu pour contrer les pics de gradient grâce à une réinitialisation de la quantité de mouvement et un rognage de gradient conscient des pics. Des expériences approfondies, comprenant à la fois la pré-formation et le fine-tuning, démontrent que SPAM surpasse de manière constante Adam et ses variantes dans diverses tâches, y compris (1) la pré-formation des LLM de 60M à 1B, (2) la pré-formation des LLM sur 4 bits, (3) l'apprentissage par renforcement et (4) la prévision de séries temporelles. De plus, SPAM facilite l'entraînement efficace en mémoire en permettant un moment clairsemé, où seul un sous-ensemble de termes de moment est conservé et mis à jour. Lorsqu'il fonctionne sous contraintes de mémoire, SPAM surpasse les optimiseurs efficaces en mémoire de pointe tels que GaLore et Adam-Mini. Notre travail souligne l'importance d'atténuer les pics de gradient dans l'entraînement des LLM et introduit une stratégie d'optimisation efficace qui améliore à la fois la stabilité de l'entraînement et l'efficacité des ressources à grande échelle. Le code est disponible sur 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