【 ICLR 2025 】Adam 有了 mini 版：内存减半，吞吐量可提升 50%

101 0 0

文章摘要

在大语言模型训练中，内存开销是一个关键挑战，尤其是Adam优化器的高内存需求。为了解决这一问题，研究团队提出了一种轻量化优化器Adam-mini，通过分析Transformer模型的Hessian结构，大幅减少了Adam优化器中学习率的数量。Adam-mini在降低50%优化器内存开销的同时，保持了与Adam优化器相当甚至更优的性能，提升吞吐量可达50%。这一成果已发表于ICLR 2025，并在多个平台获得广泛关注。

Adam优化器在训练过程中需要额外存储两个状态变量：一阶动量m和二阶动量v，这显著增加了内存负担。例如，训练一个Llama 2-7B模型的内存占用中，Adam状态变量占58GB，远超权重参数的28GB。如此高的内存需求对硬件造成了巨大压力，即便使用A100-80GB显卡，仍需依赖内存分片或CPU卸载技术，增加了训练复杂性和通信开销。

研究团队通过分析Transformer模型的Hessian结构，发现了Hessian块异质性现象，即不同参数块的Hessian矩阵特征谱存在显著差异。这一现象解释了为什么在Transformer架构中，Adam优化器优于SGD。实验表明，块异质性程度越高，SGD的性能越差，而Adam能够适应不同参数块的优化需求。

基于这一发现，研究团队提出了Adam-mini优化器。Adam-mini通过简化学习率，取子块梯度的均方值为唯一学习率，并在块内共享，从而无损地移除了超过99.9%的二阶动量v，节省了接近50%的内存。实验验证表明，Adam-mini在Llama系列模型的预训练中，不仅大幅降低了内存开销，还在多项指标上与AdamW持平，甚至在某些任务中表现更优。此外，Adam-mini在预训练Llama 2-7B模型时，吞吐量提升了49.6%，训练时间缩短了33%。

Adam-mini无需额外调参，使用AdamW原版超参数即可达到相同性能，且可扩展到更大规模的模型。研究团队还进一步结合低秩方法，提出了GaLore-mini优化器，通过利用梯度的低秩特性，将梯度投影到低维空间来更新优化器状态，同时结合Adam-mini的策略，显著减少了优化器的存储需求，实现了81%的内存节省。

Adam-mini与现有框架高度兼容，用户仅需一行代码即可轻松安装，使用方法与AdamW类似。研究团队还提供了GitHub项目主页，供用户了解更多代码使用介绍。此外，课题组正在招收博士后、MPhil/博士研究生以及实习生/科研助理，欢迎对LLM/优化算法有浓厚兴趣、具备扎实数学与编程基础的研究者加入。