Quantization error triggers loss spikes in language model training
The Role of Quantization in Loss Spikes
Machine LearningArtificial Intelligence
Summary
Training large language models can sometimes hit sudden problems called loss spikes that waste time and computing power. These spikes happen because parts of the model's output become unstable and grow too large during training. The authors found that errors caused by limited numerical precision in the model's calculations, called quantization errors, cause these late-stage loss spikes. Using lower precision makes the problem worse, while higher precision helps prevent it. They also show that switching the training method for the model’s final layer helps keep output sizes in check and makes training more stable for longer.
What this means in practice
- •For machine learning engineers: Improve pretraining stability in large language models by adjusting optimizer and precision settings to avoid costly training interruptions.
- •For hardware designers: Design hardware and numerical formats that reduce quantization errors to enhance the reliability of long-running neural network training tasks.
Authors
Zaiqian Chen, Aditya Somasundaram, Sophie Li, Sean McLeish, John Kirchenbauer, Tom Goldstein, Micah Goldblum
Abstract
Loss spikes in language model pretraining can waste substantial compute and derail otherwise healthy training runs. A primary driver of these instabilities is the growth and divergence of output logits in the final language modeling (LM) head. In this work, we study a particularly costly form of this instability associated with output logit growth: delayed, terminal divergences that emerge only after prolonged training, which we call time-bomb spikes. We find that these instabilities are preceded by rapidly growing quantization error in the LM head, which eventually culminates in a loss spike. Crucially, reducing numerical precision accelerates time-bomb spikes, while increasing precision improves stability. We then propose a simple yet effective mitigation that replaces AdamW on the LM head with SGD with momentum, suppressing logit growth and improving stability over long training horizons. Finally, we find that existing mitigation techniques from the literature (along with our own) all reduce quantization error, suggesting a shared numerical mechanism behind their effectiveness. Together, our results identify finite precision error as a hidden driver of time-bomb spikes and show that suppressing this error is key to stable long-horizon pretraining.