量化优化器状态会引发更新停滞,重启机制可有效缓解。
Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets
- 提出量化导致优化器状态滞后的机制分析
- 低精度下重启能恢复性能并减少内存占用
- 给出基于理论的重启时机选择方法
量化优化器状态是实现大规模预训练内存效率的关键,但其动态特性尚不清晰。本文研究低精度指数移动平均(EMA)优化器状态,发现量化会使大量名义更新回绕至相同存储值,导致状态实际停滞,影响自适应速度。我们建立简单预测模型,估算单步停滞概率,并揭示停滞随时间累积的过程。该视角解释了为何在低精度下重置优化器状态有效:当量化后的EMA趋于停滞时,重置可临时恢复响应性。据此,我们推导出一种基于理论的重置周期选择方法,表明在低精度下关键不仅是是否重置,更是何时重置。受控模拟与大语言模型预训练实验均表明,合理重置策略可在显著降低优化器状态内存的同时,恢复因低精度存储损失的性能。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。