论文
预处理器空间中的舍入:重新设计 4 位 AdamW 优化器状态量化
Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization
摘要
量化 AdamW 的优化器状态会减少持久存储,但量化误差会通过时刻重复传播并扰乱后续的自适应更新。我们从 emph{舍入空间} 的角度重新设计 AdamW 的 4 位优化器状态量化:量化器在相邻重建级别之间进行选择的坐标。对于第二时刻,对邻近零的量化单元的局部分析表明,较小的平均状态误差不一定意味着下一步的较小的平均预处理器误差。一维二次构造进一步显示了状态空间和预处理器空间舍入下不同质的优化动态。这些结果激发了零包含预处理器空间随机舍入(\textbf{ZIP-SR}),它在二阶矩码本中保留零并计算预处理器空间中的随机舍入概率。作为补充途径,零排除 EDEN 校准 (\textbf{ZE-EDEN}) 使用零排除二阶矩码本并重新调整量化二阶矩块以减轻 由正量化基底引起的预处理器失真。两种配置在第一时刻都使用 4 位 NormalFloat (NF4),在训练的最后 10\% 期间对 LM 头第一时刻进行有针对性的随机舍入。在从 \textbf{130M} 到 \textbf{2.7B} 参数范围的 GPT 和 Llama 式预训练实验中,两种方法在每个评估的模型大小下都将 TorchAO 4 位 AdamW 的平均验证损失差距减少到 32 位 AdamW,报告的最大差距减少达到 \textbf{70\%}。在全参数监督微调中,两种方法都实现了比 TorchAO 更低的验证损失,同时在下游任务上保持接近 32 位 AdamW。