Gefen优化器将AdamW内存占用降低8倍,无需修改超参数即可提升训练效率。
Gefen: Optimized Stochastic Optimizer
- 通过共享参数块的二阶矩并量化一阶矩,大幅减少内存占用。
- 每十亿参数节省6.5 GiB内存,性能与AdamW相当。
- 适合大模型训练,可显著提升单机或分布式场景下的吞吐量。
AdamW是现代深度学习的默认优化器,但其一阶和二阶矩状态会增加约两倍于参数大小的内存开销,加剧大规模预训练的成本。我们提出Gefen,一种内存高效的优化器,自动在参数块间共享二阶矩估计,并使用学习得到的码本对一阶矩进行量化,使AdamW的内存占用降低约8倍,每十亿参数减少6.5 GiB。该方法基于理论发现:大混合海森矩阵元素限制了平方梯度比值趋近于1,表明海森对齐参数天然适合共享二阶统计量。由于直接计算海森矩阵不可行,Gefen从初始平方梯度推断块结构,无需架构特定元数据或额外超参数,仅需AdamW默认设置。Gefen学习基于直方图的动态规划量化码本,并复用同一块进行一阶矩缩放。在多种预训练实验中,Gefen在对比的AdamW类方法中达到最低峰值优化器内存,同时保持与AdamW相当的性能。无论单机还是分布式训练,更小的内存开销支持更大微批次,显著提升吞吐量,提供一个无需改动即可替换AdamW、降低内存使用并提升效率的实用方案。完整Python实现(含融合CUDA内核)见https://github.com/ndvbd/Gefen。
原文摘要 · Abstract (English)
AdamW is a default optimizer for modern deep learning, but its first and second moment states add roughly two parameter-sized buffers to training memory, increasing the already substantial cost of large-scale pretraining. We propose Gefen, a memory-efficient optimizer that automatically shares second-moment estimates across parameter blocks and quantizes the first moment using a learned codebook, thereby reducing AdamW's memory footprint by ~8x while maintaining the same performance, corresponding to a reduction of 6.5 GiB per billion parameters. The method is motivated by a theoretical result showing that large mixed Hessian entries constrain the ratio of squared gradients toward one, suggesting that Hessian-aligned parameters are natural candidates for sharing second-moment statistics. Since computing Hessians is impractical at scale, Gefen infers block structure from the initial squared gradients, requiring no architecture-specific metadata or hyperparameters beyond AdamW defaults. Gefen learns an exact histogram-based dynamic-programming quantization codebook and reuses the same blocks for first-moment scaling. Across diverse pretraining experiments, Gefen achieves the lowest peak optimizer memory among the compared AdamW-like methods while maintaining AdamW-level performance. In single-machine or distributed training, the reduced memory footprint enables larger microbatches and improves throughput significantly over AdamW, providing a practical drop-in replacement with lower memory usage that can increase throughput and enable training larger models or using larger global batch sizes. We provide the complete Python implementation, including fused CUDA kernels at https://github.com/ndvbd/Gefen
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。