Search papers, labs, and topics across Lattice.
This paper investigates the impact of quantization on optimizer states, specifically Exponential Moving Averages (EMAs), during LLM pre-training. It reveals that quantization leads to "state staleness" where updates round to the same value, hindering adaptation. The authors develop a predictive model for stalling probabilities and demonstrate that optimizer-state resets can restore responsiveness by temporarily alleviating staleness, leading to a theory-guided method for choosing reset periods.
Quantizing optimizer states in LLM pre-training introduces "staleness," but strategically timed resets can recover lost performance and reduce memory footprint.
Quantizing optimizer states is becoming an important ingredient of memory-efficient large-scale pre-training, but the resulting optimizer dynamics remain only partially understood. We study low-precision exponential moving average (EMA) optimizer states and show how quantization can cause many nominal updates to round back to the same stored value, making the state effectively stale and slowing adaptation beyond what the nominal decay would suggest. We then develop a simple predictive model of stalling that estimates one-step stalling probabilities and characterizes how stalling builds up over time after the initialization. This perspective provides a mechanistic explanation for why optimizer-state resets help in low precision: once a quantized EMA becomes effectively stale, resetting it can temporarily restore responsiveness. Motivated by this picture, we derive a simple theory-guided method for choosing useful reset periods, showing that in low precision the key question is not only whether resets help, but when they should be applied. Experiments in controlled simulations and LLM pre-training show that suitable reset schedules recover the performance lost to low-precision state storage while substantially reducing optimizer-state memory.
