Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization
This study addresses the issue of rounding errors accumulating recursively through moment estimation during optimizer state quantization in AdamW, which leads to preconditioner distortion and perturbed adaptive updates. To mitigate this, the work reformulates 4-bit AdamW quantization from a rounding-space perspective, proposing ZIP-SR (preconditioner-space stochastic rounding) and ZE-EDEN (zero-exclusion calibration) strategies, integrated with the NF4 format to optimize quantization dynamics. Evaluated on models ranging from 130M to 2.7B parameters, the proposed approach reduces the validation loss gap by 70%, achieving training performance closely approximating full 32-bit precision.