СПАМ: Алгоритм Spike-Aware Adam с сбросом импульса для стабильного обучения LLM

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

January 12, 2025

Авторы: Tianjin Huang, Ziquan Zhu, Gaojie Jin, Lu Liu, Zhangyang Wang, Shiwei Liu

cs.AI

Аннотация

Большие языковые модели (LLM) продемонстрировали исключительную производительность в различных задачах, однако их обучение остается крайне ресурсоемким и подверженным критическим вызовам, таким как нестабильность обучения. Основным источником этой нестабильности являются всплески градиента и потерь, которые нарушают процесс обучения, часто приводя к дорогостоящим вмешательствам, таким как восстановление контрольной точки и перезапуск эксперимента, что дополнительно усиливает неэффективность. В данной статье представлено всестороннее исследование всплесков градиента, наблюдаемых во время обучения LLM, раскрывающее их распространенность в различных архитектурах и наборах данных. Наш анализ показывает, что эти всплески могут быть в 1000 раз больше обычных градиентов, значительно ухудшая производительность модели. Для решения этой проблемы мы предлагаем Spike-Aware Adam с Momentum Reset SPAM, новый оптимизатор, разработанный для противодействия всплескам градиента путем сброса импульса и учета всплесков градиента. Обширные эксперименты, включая как предварительное обучение, так и дообучение, демонстрируют, что SPAM последовательно превосходит Adam и его варианты в различных задачах, включая (1) предварительное обучение LLM от 60M до 1B, (2) предварительное обучение LLM с 4-битным кодированием, (3) обучение с подкреплением и (4) прогнозирование временных рядов. Кроме того, SPAM облегчает эффективное использование памяти, позволяя разреженный импульс, при котором обновляются и поддерживаются только некоторые члены импульса. При работе в условиях ограниченной памяти SPAM превосходит оптимизаторы современных память-эффективных методов, такие как GaLore и Adam-Mini. Наша работа подчеркивает важность смягчения всплесков градиента в обучении LLM и представляет эффективную стратегию оптимизации, которая улучшает как стабильность обучения, так и эффективность использования ресурсов в масштабе. Код доступен по ссылке 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