Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam
Authors: name: Tianjin Huang; name: Haotian Hu; name: Zhenyu Zhang; name: Gaojie Jin; name: Xiang Li; name: Li Shen; name: Tianlong Chen; name: Lu Liu; name: Qingsong Wen; name: Zhangyang Wang; name: Shiwei Liu
Published: 2025-02-24 11:09:15

Abstract: 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.