用分块均值压缩优化器状态，训练大模型省90%内存
FOAM: Blocked State Folding for Memory-Efficient LLM Training
- 分块计算梯度均值，压缩优化器状态内存占用
- 实测节省90%优化器内存，收敛速度更快
- 兼容现有高效优化器，适合大规模模型训练
大型语言模型因参数量大和训练数据多而表现卓越，但其规模导致训练时面临严重内存瓶颈，尤其在使用Adam等内存密集型优化器时。现有高效方法常依赖奇异值分解、投影或权重冻结，会引入额外计算开销、增加内存消耗或降低模型性能。本文提出折叠优化器近似矩（FOAM），通过计算分块梯度均值压缩优化器状态，并引入残差校正恢复丢失信息。理论上，FOAM在标准非凸优化设定下收敛率与原始Adam相当。实验表明，FOAM可消除高达90%的优化器状态内存开销，同时加速收敛。此外，FOAM与其它内存高效优化器兼容，性能与吞吐量优于或媲美全秩及现有高效基线。代码已开源。
原文摘要 · Abstract (English)
Large language models (LLMs) have demonstrated remarkable performance due to their large parameter counts and extensive training data. However, their scale leads to significant memory bottlenecks during training, especially when using memory-intensive optimizers like Adam. Existing memory-efficient approaches often rely on techniques such as singular value decomposition (SVD), projections, or weight freezing, which can introduce substantial computational overhead, require additional memory for projections, or degrade model performance. In this paper, we propose Folded Optimizer with Approximate Moment (FOAM), a method that compresses optimizer states by computing block-wise gradient means and incorporates a residual correction to recover lost information. Theoretically, FOAM achieves convergence rates equivalent to vanilla Adam under standard non-convex optimization settings. Empirically, FOAM eliminates up to 90\% of the memory overhead of optimizer states and accelerates convergence. Furthermore, FOAM is compatible with other memory-efficient optimizers, delivering performance and throughput that match or surpass both full-rank and existing memory-efficient baselines. Code is available at https://github.com/zqOuO/FOAM.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。