Quantizing AdamW Optimizer States and the Rounding Space Perspective
Adam and AdamW are the default optimizers for large-scale neural network training, yet their FP32 first- and second-moment buffers require 8 bytes per parameter, creating substantial memory overhead. Block-wise dynamic quantization has enabled practical 8-bit optimizers implemented in frameworks such as bitsandbytes and TorchAO, but extending quantization to 4 bits promises even greater storage savings. Li et al. (2023) extended optimizer-state quantization to 4 bits and identified the second-moment zero-point failure: mapping a small positive second-moment value to zero produces an excessively large update. Excluding zero from the second-moment codebook prevents this failure, and TorchAO's 4-bit AdamW implements a strictly positive codebook. Subsequent studies have likewise excluded zero, using linear grids or alternative positive codebooks. Our key insight is that the rounding rule is critical to whether the zero-point failure occurs with a zero-inclusive codebook. Conventional rounding chooses between neighboring levels by their values in second-moment state space, but AdamW transforms the second moment nonlinearly before scaling the update. Rounding that is accurate in state space can still distort the resulting update. By defining rounding decisions in preconditioner space, we retain the same zero-inclusive codebook while reducing the probability of harmful zero selections. More broadly, optimizer-state quantizers should account for how stored states affect optimization, not solely for reconstruction error in state space.
ZIP-SR and ZE-EDEN: Two Redesigned 4-bit Configurations
The paper introduces two 4-bit AdamW configurations that share the same persistent storage cost as TorchAO 4-bit AdamW but achieve lower validation loss. 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.
Zero-Inclusive Preconditioner-space Stochastic Rounding (ZIP-SR) retains zero in the second-moment codebook and computes stochastic-rounding probabilities in preconditioner space. At each token, the quantizer chooses between adjacent reconstruction levels based on the preconditioner induced by each candidate level rather than the raw second-moment value. This changes the rounding coordinate from state space to preconditioner space, preserving the zero-inclusive codebook while reducing the probability of storing zero when the working second moment is small and positive. ZIP-SR combines the zero-inclusive Dyn4 codebook with Update-SR, a rounding rule that preserves the mean of the current step preconditioner.
Zero-Excluding EDEN calibration (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. The zero-excluding Dyn4-NZ codebook prevents the zero-endpoint failure by construction, but imposes a positive floor that inflates small second moments and suppresses corresponding large preconditioner entries. ZE-EDEN applies guarded EDEN-inspired calibration to the decoded block, adjusting the block scale downward when the calibrated inner product exceeds the uncalibrated norm. This lowers the effective positive floor and permits larger preconditioner entries, partially restoring the tail suppressed by uncalibrated quantization. ZE-EDEN combines Dyn4-NZ with State-RTN and EDEN calibration.
Experiments: 70% Gap Reduction Across Model Sizes
Across GPT- and Llama-style pretraining experiments ranging from 130M to 2.7B parameters, both ZIP-SR and ZE-EDEN 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 70%. The experiments measure the difference in final validation loss between the 4-bit configuration and matched 32-bit AdamW, reported as positive gaps. Relative to TorchAO 4-bit AdamW, ZIP-SR and ZE-EDEN close a substantial portion of this gap at all scales.
Component attribution on GPT-style 834M pretraining reveals the contribution of each design choice. Replacing the Lin4-NZ second-moment codebook with Dyn4-NZ increases the validation-loss gap, confirming the importance of the dynamic codebook. Switching the first moment from SDyn4 to NF4 gives the largest stepwise gap reduction. From that baseline, both EDEN calibration (ZE-EDEN) and the Dyn4 Update-SR configuration (ZIP-SR) further reduce the gap, with ZIP-SR achieving a 32.5% reduction and ZE-EDEN achieving a 14.7% reduction relative to the NF4-first-moment baseline. The sample standard deviations across three seeds confirm that these improvements are consistent.
In full-parameter supervised fine-tuning, both recipes achieve lower validation loss than TorchAO while remaining close to 32-bit AdamW on downstream tasks. ZIP-SR and ZE-EDEN each match or exceed the 32-bit baseline on several checkpoint evaluations, demonstrating that the memory savings do not come at the expense of downstream task performance.
First-Moment Quantification and LM-head Stochastic Rounding
The paper examines first-moment codebook choice and rounding rule, finding that NF4 spreads its levels across the high-density region of block-normalized first-moment values, yielding lower mean squared reconstruction error than SDyn4 or FP4. Complete single-seed pretraining of the GPT-style 834M model confirms this ordering: NF4 achieves lower final validation-loss gaps than both SDyn4 and FP4 under both ZIP-SR and ZE-EDEN configurations. However, NF4-RTN trajectories can depart from 32-bit AdamW during the learning-rate cooldown in larger models. In GPT-style 2.7B pretraining, the NF4-RTN trajectory becomes unstable at cooldown, while NF4-SR remains stable.
Matched checkpoint interventions localize the instability to deterministic rounding of the LM-head first moment: applying NF4-RTN only to this state reproduces the instability, while replacing RTN with SR restores the stable trajectory. The radial bias of the LM-head first moment under NF4-RTN is predominantly negative, indicating systematic contraction along the direction of the moment vector. Under NF4-SR, the bias remains centered near zero. Because the quantized first moment is repeatedly fed back into its recurrence, the persistent inward bias under RTN can accumulate across iterations, explaining the subsequent loss instability. Guided by this diagnosis, the final methods use NF4-RTN for the first moment by default and switch only the LM-head first moment to NF4-SR during the final 10% of training.
From Persistent Storage to Downstream Task Performance
All three 4-bit configurations—TorchAO 4-bit, ZIP-SR 4-bit, and ZE-EDEN 4-bit—use packed 4-bit storage for eligible moment tensors and therefore have essentially the same persistent optimizer-state footprint. The codebook and rounding changes require no additional per-parameter state, and ZE-EDEN absorbs its calibration factor into the existing block scale. With block size 128 and one FP32 scale per block for each moment, storing both AdamW moments requires 1.0625 bytes per parameter on eligible tensors, compared with 8 bytes per parameter for FP32 moment storage, corresponding to a 7.53× compression. At the scale of the experiments, this corresponds to approximately 19.2 GB less persistent moment-state storage for the 2.7B-parameter pretraining model and roughly 1–2 TB less for 3B–8B full-parameter SFT models. To keep the storage comparison controlled, all 4-bit methods use the same tensor-eligibility criterion as the TorchAO baseline; only normalization-layer moment tensors fall outside this criterion and therefore remain in FP32.
Both ZIP-SR and ZE-EDEN reduce TorchAO 4-bit AdamW's mean validation-loss gap to 32-bit AdamW at every evaluated model size across GPT- and Llama-style pretraining from 130M to 2.7B parameters. The largest reported gap reduction reaches 70%, meaning that these configurations achieve validation loss within 30% of the 32-bit baseline even under 4-bit quantization. In full-parameter supervised fine-tuning, both recipes achieve lower validation loss than TorchAO while remaining close to 32-bit AdamW on downstream tasks, confirming that the memory savings do not compromise task performance.