基于共享梯度界的 Adam 优化器状态受保护压缩

Adam 为每个参数保存两个 FP32 矩，其字节数是 FP32 权重的两倍，但压缩它们存在风险：更新需要除以二阶矩的平方根，因此微小的重构误差也可能显著改变恢复后的计算结果。我们研究受保护的矩压缩：只有当某个坐标解码后的矩在裁剪域内的每个梯度下，都能以学习率为单位、在容差 $τ$ 内复现原生 FP32 Adam 更新时，才接受这些矩；其余坐标均精确存储。直观的认证方法会分别计算两种更新在各梯度区间内的取值范围，但该方法失效了：在 OPT-125M 训练后的一百万个坐标上，98.9% 的坐标回退为精确存储，有效载荷增至原始 FP32 的 126.6%。我们的关键观察是，两种更新使用相同的梯度和参数。通过直接对二者之差求界，并计入 FP32 舍入误差余量、采用固定的线性与对数混合梯度分区，在 $τ=0.01$ 时，回退比例降至 0.53%，有效载荷降至 28.3%。在对 OPT-125M 和 ResNet-18 进行的 120 次完整检查点恢复中，每个模型涵盖四个训练随机种子和三个训练进度比例；受保护的 8 位对数编码使两个模型的矩有效载荷相较于 FP32 分别减少 69.5% 和 71.9%，相较于采用相同保护机制的 BF16 分别减少 40.8% 和 45.4%，相较于采用相同保护机制的分块 8 位编码分别减少 29.3% 和 37.6%，且每次恢复后的首次更新均处于容差范围内。编码与恢复的总耗时为受保护 BF16 的 2.3–2.6$\times$，但低于受保护的分块 8 位编码。联合搜索两个矩的编码所额外节省的字节数不足 0.001%，却使编解码时间增加 6%，因此对每个矩分别舍入已足够。该保证涵盖固定 Torch 运算模型下恢复后的首次更新；对于后续训练，我们通过实验进行评估。