Stable-SPAM: Como Treinar em 4-Bit com Mais Estabilidade do que o Adam em 16-Bit

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

Resumo

Este artigo avalia de forma abrangente vários otimizadores recentemente propostos para treinamento em 4 bits, revelando que a precisão de baixo bit amplifica a sensibilidade às taxas de aprendizado e frequentemente causa instabilidade nas normas dos gradientes, levando à divergência em taxas de aprendizado mais altas. Entre esses, o SPAM, um otimizador recente que apresenta reset de momentum e corte de gradiente com consciência de picos, alcança o melhor desempenho em vários níveis de bits, mas luta para estabilizar as normas dos gradientes, exigindo um ajuste cuidadoso da taxa de aprendizado. Para abordar essas limitações, propomos o Stable-SPAM, que incorpora técnicas aprimoradas de normalização e corte de gradiente. Em particular, o Stable-SPAM (1) atualiza adaptativamente o limite de corte para gradientes com picos, rastreando seus máximos históricos; (2) normaliza toda a matriz de gradiente com base em suas estatísticas históricas de norma l_2; e (3) herda o reset de momentum do SPAM para redefinir periodicamente os momentos de primeira e segunda ordem do Adam, mitigando o acúmulo de gradientes com picos. Experimentos extensivos mostram que o Stable-SPAM estabiliza efetivamente as normas dos gradientes no treinamento de LLMs em 4 bits, oferecendo desempenho superior em comparação com o Adam e o SPAM. Notavelmente, nosso modelo LLaMA-1B de 4 bits treinado com Stable-SPAM supera o LLaMA-1B BF16 treinado com Adam em até 2 pontos de perplexidade. Além disso, quando ambos os modelos são treinados em 4 bits, o Stable-SPAM alcança a mesma perda que o Adam, enquanto requer apenas cerca de metade das etapas de treinamento. O código está disponível em 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.