Stable-SPAM: Cómo entrenar en 4 bits con mayor estabilidad que Adam en 16 bits

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

February 24, 2025

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

cs.AI

Resumen

Este artículo evalúa exhaustivamente varios optimizadores propuestos recientemente para entrenamiento de 4 bits, revelando que la precisión de bajo bit amplifica la sensibilidad a las tasas de aprendizaje y a menudo causa inestabilidad en las normas de los gradientes, lo que lleva a la divergencia con tasas de aprendizaje más altas. Entre estos, SPAM, un optimizador reciente que incluye reinicio de momento y recorte de gradientes consciente de picos, logra el mejor rendimiento en varios niveles de bits, pero tiene dificultades para estabilizar las normas de los gradientes, requiriendo un ajuste cuidadoso de la tasa de aprendizaje. Para abordar estas limitaciones, proponemos Stable-SPAM, que incorpora técnicas mejoradas de normalización y recorte de gradientes. En particular, Stable-SPAM (1) actualiza adaptativamente el umbral de recorte para gradientes con picos al rastrear sus máximos históricos; (2) normaliza toda la matriz de gradientes basándose en las estadísticas históricas de su norma l_2; y (3) hereda el reinicio de momento de SPAM para restablecer periódicamente los primeros y segundos momentos de Adam, mitigando la acumulación de gradientes con picos. Experimentos extensos muestran que Stable-SPAM estabiliza efectivamente las normas de los gradientes en el entrenamiento de modelos de lenguaje grandes (LLM) de 4 bits, ofreciendo un rendimiento superior en comparación con Adam y SPAM. Notablemente, nuestro modelo LLaMA-1B de 4 bits entrenado con Stable-SPAM supera al LLaMA-1B entrenado con Adam en precisión BF16 por hasta 2 puntos de perplejidad. Además, cuando ambos modelos se entrenan en 4 bits, Stable-SPAM alcanza la misma pérdida que Adam mientras requiere aproximadamente la mitad de los pasos de entrenamiento. El código está disponible en 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.