Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets

Kristi Topollai; Anna Choromanska

cs.LGMar 17, 2026

Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets

Authors: Kristi Topollai, Anna Choromanska

Organizations: New York University

Abstract

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.

Figures & tables

Figure 1: Mechanism of state-update stalling and reset recovery in a quantized EMA under nearest rounding. Top: The high-precision proposal $v_{t}$ changes the stored state $v_{t}^{q}$ only when $|\Delta_{t}|\geq u_{t}/2$ , escaping its local rounding cell. As $v_{t}^{q}$ grows, the spacing $u_{t}=\operatorname{ulp}(v_{t-1}^{q})$ also grows, causing more proposals to round back to the stored value and preventing the state from tracking changes in the local scale. Bottom: The resulting stalling probability increases with the state age $j$ toward $P_{\mathrm{stall}}^{\mathrm{ss}}$ . Resetting $v_{t_{r}}^{q}\leftarrow 0$ and $j\leftarrow 0$ restores responsiveness, allowing the state to track the new, lower scale.

Format	$\varepsilon$	$\hat{\rho}$	$P^{\mathrm{NR}}_{\mathrm{stall}}(\hat{\rho})$
BF16	$2^{-7}$	$2.71$	$\approx 0.946$
FP8 (E4M3)	$2^{-3}$	$43.3$	$\approx 1.000$
FP4 (E2M2)	$2^{-2}$	$86.6$	$\approx 1.000$

Table 1: Theoretical steady-state stalling probabilities under nearest rounding (NR) from ( 11 ).

Figure 2: Simulated EMA stalling during LLaMA-60M pre-training on WikiText. After warmup, each moment update is skipped independently with probability

p

Figure 3: Measured stalling fraction during LLM pre-training for quantized EMA states. As the second moment grows away from initialization, the stalled fraction rises and then plateaus near its steady-state level. Lower-precision formats enter the stalled regime earlier and remain there longer. 100M-parameter LLaMa model trained on C4.

Format	BF16	FP8	FP4
$P^{\mathrm{SR}}_{\mathrm{stall}}$	$0.825$	$0.989$	$0.994$

Table 2: Theoretical steady-state stalling probabilities under SR for

\beta_{2}=0.999

Figure 4: Agreement between predicted and measured steady-state stalled fractions. Bars show empirical measurements, while dashed lines show the corresponding theoretical predictions.

Figure 5: Validation loss curves for the 3 model scales under 2 different precisions

	BF16	FP8	FP4
$K^{\star}$	$\approx 1116$	$\approx 320$	$\approx 224$

Table 3: Derived reset periods from ( 29 ) with

\beta_{2}=0.999

and

s_{0}=0.6

Configuration	60M	100M	350M
BF16 AdamW	3.52	3.27	2.99
BF16 AdamW + SR + Resets	3.45	3.21	2.92
FP8 AdamW	3.53	3.33	3.02
FP8 AdamW + SR + Resets	3.44	3.22	2.92
FP4 AdamW	3.52	3.32	3.03
FP4 AdamW + SR + Resets	3.48	3.26	2.97

Table 4: Validation loss across model scales.

Figure 6: Sensitivity to the reset period.

Configuration	FP32	BF16	FP8	FP4
NR only	3.155	3.181	3.193	3.199
SR only	–	3.156	3.160	3.182
Resets only	3.161	3.158	3.159	3.166
SR + Resets	–	3.155	3.155	3.160
Memory savings	0%	50%	75%	87.5%

Table 5: Component ablation with FP32 training and quantized optimizer states only. In the FP32 column, we apply resets only, with period 1000 for both moments.

Reset Strategy	Validation Loss
Symmetric Resets	3.223
Asymmetric Resets	3.212
Adaptive Resets	3.225

Table 6: Different resetting policies under FP8.

Configuration	BF16	FP8
SOAP	3.22	4.15
SOAP + SR + Resets	3.19	3.23

Table 7: SOAP with and without SR + resets.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Model	Size	Format	Round.	Meas. $P_{\mathrm{ss}}$	Pred. $P_{\mathrm{ss}}$
LLaMA	60M	BF16	NR	.969	.946
LLaMA	60M	BF16	SR	.812	.825
LLaMA	60M	FP8	NR	.999	1.000
LLaMA	60M	FP8	SR	.986	.989
LLaMA	100M	BF16	NR	.967	.946
LLaMA	100M	BF16	SR	.814	.825

Appendix

Table 8: Complete LLaMA steady-state stalling measurements at

\beta_{2}=0.999

. The same prediction is used at every model scale. The 1B and 7B results are obtained from targeted diagnostic runs rather than complete pretraining runs.

Model	Size	Format	Round.	Meas. $P_{\mathrm{ss}}$	Pred. $P_{\mathrm{ss}}$
GPT-NeoX	160M	BF16	NR	.971	.946
GPT-NeoX	160M	BF16	SR	.817	.825
GPT-NeoX	160M	FP8	NR	.999	1.000
GPT-NeoX	160M	FP8	SR	.989	.989
GPT-NeoX	410M	BF16	NR	.971	.946
GPT-NeoX	410M	BF16	SR	.820	.825

Appendix

Table 9: Complete steady-state stalling measurements for the alternative architectures.

Format	$\chi^{2}_{1}$ KS	Stall-mass error	Mean mantissa	Relative ULP error
BF16	0.016	0.0001	1.318	0.193
FP8	0.011	0.0006	1.413	0.213
FP4	0.018	0.0031	1.427	0.307

Appendix

Table 10: Assumption audit after warmup on LLaMA-100M. The KS, stall-mass-error, and relative-ULP-error entries are medians across tensors. The theoretical mean mantissa under the log-uniform approximation is

1/\ln 2\simeq 1.443

	BF16	FP8	FP4
$P_{\mathrm{init}}$	$0.17$	$0.53$	$0.97$
$j^{\star}(P_{0}{=}0.5)$	$76$ / $\mathbf{50}$	$\mathbf{0}$ / $\mathbf{0}$	$\mathbf{0}$ / $\mathbf{0}$
$j^{\star}(P_{0}{=}0.8)$	$464$ / $\mathbf{360}$	$15$ / $\mathbf{15}$	$\mathbf{0}$ / $\mathbf{0}$
$j^{\star}(P_{0}{=}0.9)$	$1051$ / $\mathbf{700}$	$36$ / $\mathbf{43}$	$\mathbf{0}$ / $\mathbf{0}$
$j^{\star}(P_{0}{=}0.95)$	$3044$ / –	$61$ / $\mathbf{79}$	$\mathbf{0}$ / $\mathbf{0}$

Appendix

Table 11: Startup window under nearest rounding for

\beta_{2}=0.999

after incorporating the empirical floor

P_{\mathrm{init}}

. Theory uses ( 20 )–( 22 ); empirical values are measured in LLM pretraining. “

\mathbf{0}

” means the target is already met immediately because

P_{0}\leq P_{\mathrm{init}}

. “–” means the threshold was not reached empirically within the observation window.

Figure 7: Measured stalling fraction for the first moment during pretraining of the 100M-parameter LLaMA model on C4.

Format	$\hat{\rho}$	$0.5$	$0.6$	$0.7$	window
		$K^{\star}$ for tolerance level $s_{0}$			Empirical
BF16	$2.71$	$\approx 1004$	$\approx 1116$	$\approx 1262$	$500$ – $1500$
FP8	$43.3$	$\approx 295$	$\approx 320$	$\approx 351$	$300$ – $600$
FP4	$86.6$	$\approx 206$	$\approx 224$	$\approx 246$	$50$ – $250$

Appendix

Table 12: Reset periods from the cycle-averaged thresholded rule ( 58 ) under NR with

\beta_{2}=0.999

. We report several tolerance levels

s_{0}

to show the sensitivity of the rule to the tolerated amount of reset-sensitive staleness.

Params	Hidden	Intermediate	Heads	Layers	Steps	Tokens	LR
60M	512	1376	8	8	10K	1.3B	0.001
100M	640	1708	10	12	20K	2.6B	0.001
350M	1024	2736	16	24	60K	7.8B	0.0004

Appendix

Table 13: Model configurations used in our experiments. Data amounts are given in tokens.

Figure 8: Validation loss curves for the three model scales with FP4 optimizer-state storage.

States	AdamW	AdamW+SR+resets
BF16	2.91	2.77
FP8	2.89	2.78

Appendix

Table 14: End-to-end pretraining of LLaMA-1B for 120K steps (approximately 16B tokens). Lower loss is better.

States	AdamW	AdamW+SR+resets
BF16	3.162	3.044
FP8	3.223	3.098

Appendix

Table 15: Long-horizon LLaMA-100M pretraining for 60K steps. Lower loss is better.

States	Method	HellaSwag	PIQA	ARC-E	ARC-C	WinoGrande	Avg.
BF16	AdamW	31.4	61.6	30.9	25.3	49.0	$39.6\pm 0.0$
BF16	AdamW+SR+resets	33.1	63.6	32.4	25.5	50.0	$\mathbf{40.9\pm 0.1}$
FP8	AdamW	30.7	62.2	30.0	25.0	48.8	$39.3\pm 0.1$
FP8	AdamW+SR+resets	32.5	63.9	31.9	25.6	50.7	$\mathbf{40.9\pm 0.1}$
FP4	AdamW	31.4	62.8	29.9	24.0	48.2	$39.3\pm 0.2$
FP4	AdamW+SR+resets	32.8	63.4	31.7	25.1	49.8	$\mathbf{40.6\pm 0.1}$

Appendix

Table 16: Zero-shot evaluation of the 350M checkpoints. Scores are percentages. Avg. is the mean over the listed tasks, with standard deviation over five evaluation seeds.

Budget	Method	State configuration	Loss
BF16	AdamW	BF16	3.27
	AdamW+SR+resets	BF16	3.21
	Adam-mini	FP32	3.24
	Adam-mini+resets	FP32	3.24
FP8	8-bit AdamW	Dettmers blockwise INT8	3.24
	8-bit AdamW+SR+resets	Dettmers blockwise INT8	3.21

Appendix

Table 17: Comparison of low-memory optimizers on LLaMA-100M, grouped by optimizer-state-memory budget. The budget labels refer to the memory group, not necessarily to the precision of every stored entry. Lower loss is better.

Implementation	Resets	Sec./step	Tokens/sec	State memory (GiB)
FP32 native	No	3.888	25,718	0.746
BF16 native	No	3.862	25,895	0.373
BF16 native	$K=100$	3.859	25,915	0.373
BF16 simulated Quant/Dequant	No	4.120	24,268	0.373
BF16 simulated Quant/Dequant + simulated SR	No	4.130	24,209	0.373
BF16 simulated Quant/Dequant + simulated SR	$K=100$	4.068	24,581	0.373

Appendix

Table 18: BF16 runtime and optimizer-state-memory measurements on 100K token batches. K=100 is a deliberately frequent reset overhead.

Figure 9: Validation loss curves for LLaMA-60M with FP4 optimizer-state storage under per-tensor scaling.

P_{\mathrm{stall}}(k)\leftarrow\frac{1}{d}\sum_{i=1}^{d}\mathbf{1}\!\left[(\bar{s}_{t})_{i}=(\bar{s}_{t-1})_{i}\right]

Appendix

Algorithm 2 Stateful Optimizer with Quantized State Storage and Adaptive Resets

Low-Precision Training

See all 24 papers

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers
Aug 23, 2026Yan WangDeep Learning Optimization Memory-Efficient Optimization
Independent Researcher
STQuant: Spatio-Temporal Adaptive Framework for Optimizer Quantization in Large Multimodal Model Training
Apr 8, 2026Minglu Liu, Cunchen Hu, Liangliang Xu +3Neural Network Quantization Memory-Efficient Training
Xidian University Xian, Shaanxi, China · China Telecom Cloud Computing Research Institute Beijing, China
Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization
Oct 8, 2026Hanyang Li, Shao Tang, Daniel Thomas Braithwaite +7Stochastic Rounding 4-Bit Quantization
University of California, Berkeley · Nubank

Aug 23, 2026J/K move · Enter open · S save

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

Yan Wang

Independent Researcher

Deep Learning Optimization Memory-Efficient Optimization Low-Bit Quantization Neural Network Quantization

Open paper AskPDF arXiv

Optimizer-state quantization is commonly designed for Adam's dense, parameter-aligned first- and second-moment arrays. This abstraction breaks for memory-efficient optimizers, whose states may be factored, confidence-modulated, or maintained in a projected space, so similar reconstruction error can produce different update error. We formulate optimizer-state quantization as a joint problem over representation, topology, and update semantics. We then introduce Adaptive Log-Space (AL) quantization for non-negative states. AL fits each block's observed nonzero logarithmic interval and reserves a separate code for exact zero, enforcing

q = 0 \Leftrightarrow x = 0

; signed momentum and state precision remain independently selectable. Controlled probes show that adaptive ranges reduce update error and temporal drift, exact-zero reservation preserves dormant states, and state topology constrains useful block granularity. End-to-end language-model training evaluates the resulting policy across dense, factored, confidence, and projected optimizer states. On TinyLlama-1.1B, AL8 with uniform 8-bit momentum reaches 72.90 perplexity versus 73.54 for bitsandbytes 8-bit AdamW, with comparable optimizer-state storage and higher throughput. CAME matches reference-level final perplexity across three seeds when its non-negative states use AL16, while a semantic grouping-and-protection policy closes most of quantized Adafactor's 100K-step late-loss gap. These results make state topology and update semantics first-class design constraints for optimizer quantization.

Abstract

Low-Precision Training

Explore similar work

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

STQuant: Spatio-Temporal Adaptive Framework for Optimizer Quantization in Large Multimodal Model Training

Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization