Stable-SPAM : Comment s'entraîner en 4 bits avec plus de stabilité qu'avec Adam en 16 bits

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

February 24, 2025

Auteurs: Tianjin Huang, Haotian Hu, Zhenyu Zhang, Gaojie Jin, Xiang Li, Li Shen, Tianlong Chen, Lu Liu, Qingsong Wen, Zhangyang Wang, Shiwei Liu

cs.AI

Résumé

Cet article évalue de manière exhaustive plusieurs optimiseurs récemment proposés pour l'entraînement en 4 bits, révélant que la faible précision amplifie la sensibilité aux taux d'apprentissage et provoque souvent une instabilité des normes de gradient, conduisant à une divergence à des taux d'apprentissage plus élevés. Parmi ceux-ci, SPAM, un optimiseur récent intégrant une réinitialisation de l'impulsion et un découpage de gradient prenant en compte les pics, obtient les meilleures performances à différents niveaux de bits, mais peine à stabiliser les normes de gradient, nécessitant un réglage minutieux du taux d'apprentissage. Pour pallier ces limitations, nous proposons Stable-SPAM, qui intègre des techniques améliorées de normalisation et de découpage du gradient. Plus précisément, Stable-SPAM (1) met à jour de manière adaptative le seuil de découpage pour les gradients en pic en suivant leurs maxima historiques ; (2) normalise l'ensemble de la matrice de gradient en fonction de ses statistiques historiques de norme l_2 ; et (3) hérite de la réinitialisation de l'impulsion de SPAM pour réinitialiser périodiquement les premier et second moments d'Adam, atténuant ainsi l'accumulation de gradients en pic. Des expériences approfondies montrent que Stable-SPAM stabilise efficacement les normes de gradient lors de l'entraînement de modèles de langage en 4 bits, offrant des performances supérieures à celles d'Adam et de SPAM. Notamment, notre modèle LLaMA-1B en 4 bits entraîné avec Stable-SPAM surpasse le LLaMA-1B en BF16 entraîné avec Adam jusqu'à 2 points de perplexité. De plus, lorsque les deux modèles sont entraînés en 4 bits, Stable-SPAM atteint la même perte qu'Adam tout en nécessitant environ la moitié des étapes d'entraînement. Le code est disponible à l'adresse 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.