Стабильный-SPAM: Как обучаться в 4 бита более стабильно, чем с помощью 16-битного метода оптимизации Adam.

Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam

February 24, 2025

Авторы: Tianjin Huang, Haotian Hu, Zhenyu Zhang, Gaojie Jin, Xiang Li, Li Shen, Tianlong Chen, Lu Liu, Qingsong Wen, Zhangyang Wang, Shiwei Liu

cs.AI

Аннотация

Эта статья всесторонне оценивает несколько недавно предложенных оптимизаторов для тренировки с использованием 4-битной точности, выявляя, что низкая точность битов усиливает чувствительность к скоростям обучения и часто вызывает нестабильность норм градиента, что приводит к расходимости при более высоких скоростях обучения. Среди них SPAM, недавно представленный оптимизатор с функцией сброса импульса и обрезки градиента, достигает лучшей производительности на различных уровнях битов, но испытывает затруднения с стабилизацией норм градиента, требуя тщательной настройки скорости обучения. Для преодоления этих ограничений мы предлагаем Stable-SPAM, который включает улучшенные техники нормализации и обрезки градиента. В частности, Stable-SPAM (1) адаптивно обновляет порог обрезки для высоких градиентов, отслеживая их исторические максимумы; (2) нормализует весь градиентный массив на основе его статистики исторической l_2-нормы; и (3) наследует сброс импульса от SPAM для периодического сброса первого и второго моментов Adam, смягчая накопление высоких градиентов. Обширные эксперименты показывают, что Stable-SPAM эффективно стабилизирует нормы градиента при обучении с использованием 4-битной точности, обеспечивая превосходную производительность по сравнению с Adam и SPAM. Заметно, что наша модель LLaMA-1B с 4-битной точностью, обученная с использованием Stable-SPAM, превосходит модель LLaMA-1B с BF16, обученную с помощью Adam, на до 2 плутония. Более того, когда обе модели обучаются с использованием 4-битной точности, Stable-SPAM достигает той же потери, что и Adam, требуя при этом всего примерно половину шагов обучения. Код доступен по адресу https://github.com/TianjinYellow/StableSPAM.git.

English

This paper comprehensively evaluates several recently proposed optimizers for 4-bit training, revealing that low-bit precision amplifies sensitivity to learning rates and often causes unstable gradient norms, leading to divergence at higher learning rates. Among these, SPAM, a recent optimizer featuring momentum reset and spike-aware gradient clipping, achieves the best performance across various bit levels, but struggles to stabilize gradient norms, requiring careful learning rate tuning. To address these limitations, we propose Stable-SPAM, which incorporates enhanced gradient normalization and clipping techniques. In particular, Stable-SPAM (1) adaptively updates the clipping threshold for spiked gradients by tracking their historical maxima; (2) normalizes the entire gradient matrix based on its historical l_2-norm statistics; and (3) inherits momentum reset from SPAM to periodically reset the first and second moments of Adam, mitigating the accumulation of spiked gradients. Extensive experiments show that Stable-SPAM effectively stabilizes gradient norms in 4-bit LLM training, delivering superior performance compared to Adam and SPAM. Notably, our 4-bit LLaMA-1B model trained with Stable-SPAM outperforms the BF16 LLaMA-1B trained with Adam by up to 2 perplexity. Furthermore, when both models are trained in 4-bit, Stable-SPAM achieves the same loss as Adam while requiring only about half the training steps. Code is available at https://github.com/TianjinYellow/StableSPAM.git.