Stable-SPAM: Hoe stabieler trainen in 4-bit dan met 16-bit Adam

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

Samenvatting

Dit artikel evalueert uitgebreid verschillende recent voorgestelde optimalisatiemethoden voor 4-bit training, waarbij wordt aangetoond dat lage-bitprecisie de gevoeligheid voor leerpercentages versterkt en vaak instabiele gradiëntnormen veroorzaakt, wat leidt tot divergentie bij hogere leerpercentages. Onder deze methoden presteert SPAM, een recente optimalisatiemethode met momentumreset en spike-aware gradiëntclipping, het beste over verschillende bitniveaus, maar slaagt er niet in om gradiëntnormen te stabiliseren, wat zorgvuldige afstemming van het leerpercentage vereist. Om deze beperkingen aan te pakken, stellen wij Stable-SPAM voor, dat verbeterde gradiëntnormalisatie en clippingtechnieken omvat. In het bijzonder past Stable-SPAM (1) adaptief de clippingdrempel voor gespikete gradiënten aan door hun historische maxima bij te houden; (2) normaliseert de gehele gradiëntmatrix op basis van zijn historische l_2-normstatistieken; en (3) erft de momentumreset van SPAM om periodiek de eerste en tweede momenten van Adam te resetten, waardoor de accumulatie van gespikete gradiënten wordt verminderd. Uitgebreide experimenten tonen aan dat Stable-SPAM effectief gradiëntnormen stabiliseert in 4-bit LLM-training en superieure prestaties levert in vergelijking met Adam en SPAM. Opmerkelijk is dat ons 4-bit LLaMA-1B-model, getraind met Stable-SPAM, de BF16 LLaMA-1B die met Adam is getraind, overtreft met een verbetering van maximaal 2 perplexiteit. Bovendien bereikt Stable-SPAM, wanneer beide modellen in 4-bit worden getraind, hetzelfde verlies als Adam terwijl slechts ongeveer de helft van de trainingsstappen nodig is. De code is beschikbaar op 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.