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

Hanyang Li; Shao Tang; Daniel Thomas Braithwaite; Gregory Dexter; Leonardo Neves; Aman Gupta; Hiroto Udagawa; Abhishek Shivanna

cs.LGOct 8, 2026

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

Authors: Hanyang Li, Shao Tang, Daniel Thomas Braithwaite, Gregory Dexter, Leonardo Neves, Aman Gupta, Hiroto Udagawa, Abhishek Shivanna, +2 more

Organizations: University of California, Berkeley · Nubank

Abstract

Quantizing AdamW's optimizer states reduces persistent storage, but quantization errors propagate through the moment recurrences and perturb subsequent adaptive updates. We redesign 4-bit optimizer-state quantization for AdamW from the perspective of \emph{rounding space}: the coordinate in which a quantizer chooses between adjacent reconstruction levels. For the second moment, a local analysis of the quantization cell adjacent to zero shows that small mean state error need not imply small mean preconditioner error at the next step. A one-dimensional quadratic construction further shows qualitatively different optimization dynamics under state-space and preconditioner-space rounding. These results motivate Zero-Inclusive Preconditioner-space Stochastic Rounding (\textbf{ZIP-SR}), which retains zero in the second-moment codebook and computes stochastic-rounding probabilities in preconditioner space. As a complementary route, Zero-Excluding EDEN calibration (\textbf{ZE-EDEN}) uses a zero-excluding second-moment codebook and rescales the quantized second-moment block to mitigate the preconditioner distortion caused by the positive quantization floor. Both configurations use 4-bit NormalFloat (NF4) for the first moment, with targeted stochastic rounding of the LM-head first moment during the final 10% of training. Across GPT- and Llama-style pretraining experiments ranging from \textbf{130M} to \textbf{2.7B} parameters, both methods reduce TorchAO 4-bit AdamW's mean validation-loss gap to 32-bit AdamW at every evaluated model size, with the largest reported gap reduction reaching \textbf{70%}. In full-parameter supervised fine-tuning, both recipes achieve lower validation loss than TorchAO while remaining close to 32-bit AdamW on downstream tasks.

Figures & tables

Rule	$\mathbb{P}(Q(\tilde{v}_{t})=0\mid\tilde{v}_{t})$	State error $\Delta_{v}$	Precond. error $\Delta_{r}$
State-RTN	$1$	$-\tilde{v}_{t}$	$+\Theta(\epsilon^{-1})$
State-SR	$1-\tilde{v}_{t}/b$	$0$	$+\Theta(\epsilon^{-1})$
Update-RTN	$0$	$b-\tilde{v}_{t}$	$-\Theta(1)$
Update-SR	$\Theta(\epsilon)$	$b-\tilde{v}_{t}-\Theta(\epsilon)$	$-\Theta(1)$

Table 1: Conditional mean errors for fixed

0<\tilde{v}_{t}<b/2

\epsilon\downarrow 0

, with

|g_{t+1}|=\mathcal{O}(\epsilon)

Figure 1: Rounding space controls training stability. Training loss over the first

4

B tokens of GPT-style 834M pretraining, with FP32-

\mathbf{m}

and zero-inclusive Dyn4-

\mathbf{v}

. State-RTN/SR become unstable and exceed the plotted range, while Update-RTN/SR closely track 32-bit AdamW.

Figure 2: Next-step preconditioner distributions at 50% of GPT-style 834M pretraining. EDEN partially restores the large-preconditioner tail suppressed by Dyn4-NZ in the embedding and LM-head states. Log-RMSE measures coordinate-wise error from FP32.

Figure 3: Empirical distribution of block-normalized FP32 first-moment values from the layer-12 MLP up-projection at step 15,908 of GPT-style 834M pretraining, together with the probability mass assigned to each reconstruction level by FP4, SDyn4, and NF4 under RTN.

Figure 4: Late-training instability and LM-head directional bias in GPT-style 2.7B pretraining. (a) Cooldown training loss averaged in 100-update bins; NF4-SR uses SR on the LM-head first moment. (b) Radial bias

b_{\mathrm{rad}}(\tilde{\mathbf{m}}_{t})

logged every 100 steps, with FP32 non-head first moments and Dyn4/Update-SR second moments. The vertical dashed line marks the first global gradient-clipping event of the NF4-RTN continuation.

Configuration	First moment		Second moment
Configuration	Non-LM-head	LM-head	Format	Rule	EDEN
32-bit	FP32	FP32	FP32	—	$\times$
TorchAO 4-bit	SDyn4-RTN	SDyn4-RTN	Lin4-NZ	State-RTN	$\times$
ZIP-SR 4-bit	NF4-RTN	NF4-RTN/SR ∗	Dyn4	Update-SR	$\times$
ZE-EDEN 4-bit	NF4-RTN	NF4-RTN/SR ∗	Dyn4-NZ	State-RTN	$\checkmark$

Table 2: Optimizer-state configurations evaluated in the end-to-end experiments. ZIP-SR and ZE-EDEN share the NF4 first-moment design established in Section 4 , but use the two different second-moment quantization routes developed in Section 3 . The superscript

\ast

indicates a phase-dependent rule: the LM-head first-moment uses NF4-RTN before the final 10% of training and switches to NF4-SR only during that phase.

Step	Setting	First moment	Second moment			Gap	vs. parent
Step	Setting	Format	Format	Rule	EDEN	Gap	$\Delta$	Gain
0	TorchAO 4-bit	SDyn4	Lin4-NZ	State-RTN	$\times$	$0.0246$	—	—
1	Dyn4-NZ control	SDyn4	Dyn4-NZ	State-RTN	$\times$	$0.0312$	$+0.0066$	$-26.9\%$
2	NF4 first moment	NF4	Dyn4-NZ	State-RTN	$\times$	$0.0207$	$-0.0105$	$33.7\%$
3a	ZE-EDEN	NF4	Dyn4-NZ	State-RTN	$\checkmark$	$0.0177$	$-0.0031$	$14.7\%$
3b	ZIP-SR	NF4	Dyn4	Update-SR	$\times$	$0.0140$	$-0.0067$	$32.5\%$

Table 3: Component attribution on GPT-style 834M. Gaps are relative to matched 32-bit AdamW. Steps 1, 2, and 3a are means over three random seeds; steps 0 and 3b are single runs. Steps 3a and 3b both compare with step 2. Positive Gain denotes a smaller gap; derived values use unrounded gaps.

Family	Size	32-bit loss	4-bit gap to 32-bit AdamW			Best reduction vs. TorchAO (%)
Family	Size	32-bit loss	TorchAO	ZE-EDEN	ZIP-SR	Best reduction vs. TorchAO (%)
GPT-style	162M	$3.1335$ $(0.0016)$	$+0.0290$ $(0.0028)$	$+0.0200$ $(0.0021)$	$\mathbf{+0.0195}$ $(0.0020)$	32.8
	405M	$2.8197$ $(0.0008)$	$+0.0254$ $(0.0015)$	$\mathbf{+0.0167}$ $(0.0007)$	$+0.0180$ $(0.0005)$	34.3
	1.4B	$2.5186$ $(0.0012)$	$+0.0521$ $(0.0048)$	$+0.0184$ $(0.0010)$	$\mathbf{+0.0156}$ $(0.0026)$	70.1
	2.7B	$2.4044$ $(0.0019)$	$+1.5435\smash{\makebox[0.0pt][r]{\raisebox{4.30554pt}{$ \dagger $}\hskip-3.00003pt}}$ $(0.1437)$	$+0.0535$ $(0.0024)$	$\mathbf{+0.0343}$ $(0.0027)$	—
Llama-style	130M	$2.6974$ $(0.0027)$	$+0.0205$ $(0.0015)$	$+0.0177$ $(0.0009)$	$\mathbf{+0.0171}$ $(0.0003)$	16.6

Table 4: Pretraining validation loss: absolute 32-bit loss and mean paired 4-bit gaps to matched 32-bit runs (sample standard deviation); bold marks the smallest gap per row. The final column reports the larger gap reduction achieved by ZE-EDEN or ZIP-SR. All rows use three seeds. The percentage reduction for GPT-style 2.7B is omitted because the

\dagger

-marked TorchAO endpoint diverges late in cooldown.

Model	Configuration	Tulu val. loss $\downarrow$	MMLU $\uparrow$	GSM8K $\uparrow$	HumanEval $\uparrow$	IFEval strict $\uparrow$
Qwen3- 8B-Base	32-bit	$0.55135$ $(0.00010)$	$78.5$ $(0.2)$	$82.8$ $(0.6)$	$72.6$ $(0.4)$	$49.5$ $(2.1)$
	TorchAO 4-bit	$+0.00060$ $(0.00009)$	$0.0$ $(0.1)$	$+0.2$ $(1.1)$	$-0.1$ $(1.2)$	$+0.6$ $(1.8)$
	ZE-EDEN 4-bit	$-0.00003$ $(0.00014)$	$+0.1$ $(0.1)$	$+0.2$ $(0.9)$	$+1.1$ $(1.3)$	$-0.1$ $(2.4)$
	ZIP-SR 4-bit	$+0.00042$ $(0.00016)$	$+0.1$ $(0.2)$	$+0.6$ $(0.6)$	$+0.5$ $(2.0)$	$-0.1$ $(1.9)$
Llama- 3.2-3B	32-bit	$0.73255$ $(0.00044)$	$52.0$ $(2.0)$	$34.6$ $(1.2)$	$33.4$ $(1.8)$	$37.4$ $(1.9)$
	TorchAO 4-bit	$+0.00175$ $(0.00012)$	$+1.4$ $(1.5)$	$+0.4$ $(1.6)$	$+0.6$ $(1.4)$	$+0.4$ $(1.2)$

Table 5: All SFT runs use AdamW. The 32-bit values are five-seed means, while the 4-bit values are mean paired gaps relative to matched 32-bit runs; parentheses report sample SDs. Downstream scores are percentages.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Family	Size	$\boldsymbol{L}$	$\boldsymbol{d_{\mathrm{model}}}$	Heads	$\boldsymbol{n_{kv}}$	FFN
GPT-style	162M	12	768	12	—	3072
	405M	24	1024	16	—	4096
	834M	24	1536	16	—	6144
	1.4B	24	2048	16	—	8192
	2.7B	32	2560	32	—	10240
Llama-style	130M	12	768	12	4	2304

Appendix

Table 6: Architectures used for pretraining and component studies. Size denotes the rounded parameter count;

L

d_{\mathrm{model}}

, heads,

n_{kv}

, and FFN denote the layer count, hidden width, attention heads, key/value heads, and feed-forward width. GPT-style models use RoPE/GELU with untied embeddings. Llama-style models use RoPE/SwiGLU/RMSNorm, including an output norm, with

n_{kv}=4

Family	Size	Candidate peak LRs	Selected LR	Val. loss
GPT-style	162M	$\{2,1,0.6,0.3\}\!\times\!10^{-3}$	$1\!\times\!10^{-3}$	3.135
	405M	$\{1,0.6,0.3,0.15\}\!\times\!10^{-3}$	$6\!\times\!10^{-4}$	2.818
	834M	$\{0.6,0.3,0.25,0.15\}\!\times\!10^{-3}$	$6\!\times\!10^{-4}$	2.641
	1.4B	$\{0.3,0.2,0.15,0.1\}\!\times\!10^{-3}$	$2\!\times\!10^{-4}$	2.517
	2.7B	Not swept	$1.5\!\times\!10^{-4}$	—
Llama-style	130M	$\{4,3,2,1,0.6,0.3\}\!\times\!10^{-3}$	$2\!\times\!10^{-3}$	2.698

Appendix

Table 7: Peak learning rates for pretraining. Each sweep uses

20

training tokens per parameter. The final column reports the minimum observed validation loss at the precision retained in the experiment logs. The GPT-style 2.7B rate is fixed without a sweep.

$\mathbf{m}$	Raw	+ EDEN	Update-SR
	Dyn4-NZ (State-RTN)		Dyn4
FP32	$0.0049$ $(0.0001)$	$0.0038$ $(0.0003)$	—
SDyn4	$0.0312$ $(0.0008)$	$0.0308$ $(0.0027)$	$0.0239$
NF4	$0.0207$ $(0.0010)$	$0.0177$ $(0.0015)$	$0.0140$

Appendix

Table 8: GPT-style 834M component grid. Validation-loss gaps to 32-bit AdamW. Rows vary the first moment (RTN throughout); columns vary the second moment. Entries with parentheses report mean (sample standard deviation) over three seeds; the others are single runs.

Configuration	Validation loss	Gap to 32-bit
32-bit AdamW	$2.6274$	—
TorchAO 4-bit	$2.6554$	$0.0280$
ZE-EDEN	$2.6421$	$0.0147$
ZIP-SR	$2.6385$	$0.0111$

Appendix

Table 9: GPT-style 834M pretraining with cosine learning-rate decay. Final validation loss at update 31,810 for seed 42. Gaps to the 32-bit AdamW run are shown directly.

Figure 5: Matched next-step preconditioner distributions at 25%, 50%, and 75% of training. Columns show checkpoints and rows show embedding and LM-head states. Curves compare FP32, Lin4-NZ, and Lin4-NZ with EDEN; open circles mark the Lin4-NZ floor atom.

LM-head $\mathbf{m}_{t}$	Non-head $\mathbf{m}_{t}$	Final train loss	Final val. loss
FP32	NF4-RTN	2.392	2.440
NF4-RTN	FP32	2.529	2.579
NF4-SR	FP32	2.391	2.440

Appendix

Table 10: Localization of the late-training instability in GPT-style 2.7B pretraining. Matched continuations vary the quantization of the LM-head and non-head first moments while holding the remaining training configuration fixed.

Figure 6: LM-head gradient growth in the matched continuations. The solid orange curve denotes NF4-RTN and the dashed blue curve denotes NF4-SR. The vertical dashed line marks step

102{,}500

(

53.74

B tokens), at which the NF4-RTN continuation first triggers global gradient clipping.

Diagnostic	NF4-RTN	NF4-SR
Relative RMS quantization error (%)	9.21	14.75
Nonzero-to-zero fraction (%)	9.16	12.08
Accumulated shadow-state error (%)	20.95	32.18

Appendix

Table 11: Magnitude-based diagnostics for the matched LM-head continuations. Values are means over

91

recorded evaluations; lower is better. The first two rows measure instantaneous storage-rounding error; the last measures accumulated deviation from the shadow recurrence in the working first moment before storage.

Figure 7: Validation loss during GPT-style 2.7B pretraining for a single run per configuration (seed 42). The dashed line marks the switch from NF4-RTN to NF4-SR for the LM-head first moment at

90\%

of training in ZE-EDEN and ZIP-SR. TorchAO exhibits a sharp regression, whereas ZE-EDEN and ZIP-SR remain close to 32-bit AdamW.

Explore similar work

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers
Aug 23, 2026Yan WangDeep Learning Optimization Memory-Efficient Optimization
Independent Researcher
MuonQ: Enhancing Low-Bit Muon Quantization via Directional Fidelity Optimization
May 12, 2026Yupeng Su, Ruijie Zhang, Ziyue Liu +24-Bit Quantization Muon Optimizer
University of California, Santa Barbara
Gefen: Optimized Stochastic Optimizer
Jun 11, 2026Nadav Benedek, Tomer Koren, Ohad FriedDeep Learning Optimization Memory-Efficient Optimization
Reichman University · Tel Aviv University, Google Research

Aug 23, 2026, cs.LGJ/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

Explore similar work

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

MuonQ: Enhancing Low-Bit Muon Quantization via Directional Fidelity Optimization

Gefen: Optimized Stochastic Optimizer