在现代深度学习中，AdamW是默认的优化器，但其一阶和二阶矩状态大约增加了与参数数量相同的内存缓冲区。我们提出了Gefen，一种内存高效的优化器，它自动在参数块之间共享二阶矩估计，并使用学习到的码本对一阶矩进行量化，从而将AdamW的内存占用降低约8倍，同时保持相同的性能，相当于每十亿个参数减少6.5 GiB的内存。该方法的理论基础表明，大的混合Hessian项限制了平方梯度的比率接近于1，这表明与Hessian对齐的参数自然适合共享二阶矩统计信息。由于在大规模计算中计算Hessian不切实际，Gefen从初始平方梯度推断块结构，无需超出AdamW默认设置的架构特定元数据或超参数。Gefen学习了一个基于直方图的动态规划量化码本，并重用相同的块用于一阶矩缩放。在各种实验中，Gefen在比较的类似AdamW的方法中实现了最低的峰值优化器内存，同时保持了AdamW级别的性能。在FSDP和DDP训练中，减少的内存占用使得更大的微批量成为可能，并显著提高了吞吐量，提供了一种内存使用更低的实用替代方案，能够提高吞吐量并支持训练更大的模型或使用更大的批量大小。我们提供了完整的Python实现，包括融合的CUDA内核，链接在此：Gefen GitHub。
博主点评： Gefen优化器的提出是深度学习领域内存管理的一个重要进展。通过共享二阶矩和量化一阶矩，Gefen不仅显著降低了内存占用，还保持了与AdamW相同的性能，为大规模模型的训练提供了更高效的解决方案。其在FSDP和DDP训练中的表现尤其值得关注，展现了其在实际应用中的广泛潜力。