Preconditioner-space rounding: closing the gap for 4-bit AdamW optimizer states

Li et al. identify preconditioner distortion in 4-bit AdamW quantization, proposing ZIP-SR and ZE-EDEN to cut validation loss gaps to FP32 by up to 70%.

Primary source: arXiv:2610.12444: https://arxiv.org/abs/2610.12444

What’s new: In large language model training, the optimizer states of AdamW consume twice as much accelerator memory as the model weights themselves. While 4-bit state quantization reduces this overhead, quantized second moments suffer from persistent approximation errors that perturb adaptive gradient updates. In Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization (arXiv:2610.12444, 8 Oct 2026), Li et al. analyze the coordinate space in which quantizers map continuous values to discrete levels.

The authors show that minimizing quantization error in state space does not guarantee minimal error in preconditioner space, namely the inverse square root factor $1 / (\sqrt{v} + \epsilon)$ that scales updates, especially near zero. To address this mismatch, they propose two methods. First, Zero-Inclusive Preconditioner-space Stochastic Rounding (ZIP-SR) retains zero in the second-moment codebook while computing stochastic rounding probabilities directly in preconditioner coordinates. Second, Zero-Excluding EDEN calibration (ZE-EDEN) uses a zero-excluding codebook with block-level rescaling to neutralize floor distortion. Paired with 4-bit NormalFloat (NF4) for the first moment, both recipes cut the validation-loss gap between TorchAO 4-bit AdamW and 32-bit AdamW by up to 70% across pretraining runs from 130M to 2.7B parameters, matching FP32 performance in full-parameter fine-tuning.

Why it matters: Accelerator high-bandwidth memory is the primary hardware bottleneck during long-context training and fine-tuning. Formulating quantization rounding in the functional space of the preconditioner rather than naive state storage unlocks major memory savings without paying the typical penalty in pretraining perplexity or downstream accuracy.

Caveats: Evaluated on autoregressive models up to 2.7 billion parameters. Whether preconditioner-space stochastic rounding preserves numerical stability across multi-trillion-token runs on hundred-billion-parameter clusters remains to be validated. Targeted stochastic rounding of the language model head also adds engineering overhead to distributed training loops.