Gefen：优化随机优化器

AdamW 是现代深度学习的默认优化器，但其一阶矩和二阶矩状态向训练内存添加了大约两个参数大小的缓冲区，从而增加了大规模预训练本已相当大的成本。我们提出了 Gefen，一种内存高效的优化器，它自动在参数块之间共享二阶矩估计，并使用学习的码本量化第一矩，从而在保持相同性能的同时将 AdamW 的内存占用减少约 8 倍，相当于每十亿个参数减少 6.5 GiB。该方法的动机是理论结果表明，大型混合 Hessian 条目将平方梯度的比率限制为 1，这表明 Hessian 对齐参数是共享二阶矩统计数据的自然候选者。由于大规模计算 Hessians 是不切实际的，Gefen 从初始平方梯度推断块结构，不需要超出 AdamW 默认值的特定于架构的元数据或超参数。 Gefen 学习精确的基于直方图的动态编程量化码本，并重用相同的块进行一阶矩缩放。在各种预训练实验中，Gefen 在类似 AdamW 的方法中实现了最低的峰值优化器内存，同时保持了 AdamW 级别的性能。在单机或分布式训练中，与 AdamW 相比，减少的内存占用可实现更大的微批次，并显着提高吞吐量，从而提供了一种实用的直接替代方案，内存使用量更低，可以提高吞吐量并支持训练更大的模型或使用更大的全局批次大小。我们提供完整的 Python 实现，包括融合的 CUDA 内核，网址为 https://github.com/ndvbd/Gefen