论文

带梯度共享边界的 Adam 优化器状态安全压缩方法

Guarded Compression of Adam Optimizer State with Shared-Gradient Bounds

精选理由

训练时优化器状态比权重还占显存,这篇论文给出一种带数学保证的压缩法,能把 payload 压到 FP32 的三成左右,恢复更新也不跑偏。

Adam 优化器的两个 FP32 动量占用的字节量是 FP32 权重的两倍,直接压缩容易因除以二阶矩平方根的运算被微小重构误差放大。这篇 arXiv 论文提出 guarded moment compression:只接受在容差 τ 内能复现原生 FP32 更新的坐标,其余坐标精确存储。单纯按梯度区间逐一界定的证书在 OPT-125M 的 100 万坐标上逃逸率达 98.9%,payload 反而涨到 126.6%;改为直接约束两个更新的差值并配合混合线性-对数梯度分区后,τ=0.01 时逃逸率降到 0.53%,payload 降到 28.3%。在 OPT-125M 和 ResNet-18 共 120 次完整 checkpoint 恢复实验中,guarded 8-bit log 编码相对 FP32 把动量 payload 缩减 69.5% 和 71.9%,且每次恢复后的首次更新都在容差内。

原文 · arXiv cs.LG

Guarded Compression of Adam Optimizer State with Shared-Gradient Bounds

Adam keeps two FP32 moments per parameter, twice the bytes of the FP32 weights, yet compressing them is risky: the update divides by the square root of the second moment, so small reconstruction errors can change the restored computation substantially. We study guarded moment compression, in which a coordinate's decoded moments are accepted only if they reproduce the native FP32 Adam update to within a tolerance $τ$, in units of the learning rate, for every gradient in a clipped domain; all other coordinates are stored exactly. The natural certificate, which ranges each update separately over gradient intervals, fails: on one million trained OPT-125M coordinates it escapes 98.9% of coordinates and the payload grows to 126.6% of raw FP32. Our key observation is that both updates consume the same gradient and parameter. Bounding their difference directly, with FP32 rounding-error allowances and a fixed mixed linear-logarithmic gradient partition, cuts escapes to 0.53% and the payload to 28.3% at $τ=0.01$. Across 120 complete-checkpoint restores of OPT-125M and ResNet-18, with four trained seeds and three training fractions per model, guarded 8-bit log codes shrink moment payloads by 69.5% and 71.9% relative to FP32, 40.8% and 45.4% relative to equally guarded BF16, and 29.3% and 37.6% relative to equally guarded blockwise 8-bit, and every restored first update stays within tolerance. Encoding and restoration together take 2.3-2.6$\times$ as long as guarded BF16 but less time than guarded blockwise 8-bit. Searching both moment codes jointly saves under 0.001% more bytes at 6% more codec time, so per-moment rounding suffices. The guarantee covers the first restored update under a pinned Torch operation model; we measure subsequent training empirically.