提出无需存储梯度矩的零阶优化器,高效微调大模型且显存更低。
AdaMeZO: Adam-style Zeroth-Order Optimizer for LLM Fine-tuning Without Maintaining the Moments

- 用前向传播估计动量与二阶矩,不存储中间结果
- 比MeZO少70%前向传播次数,收敛更快
- 适合显存受限的大模型微调场景
微调大语言模型需应对下游任务,但传统反向传播方法消耗大量显存。近期的MeZO方法仅依赖前向传播,显著降低显存需求,但因忽略损失曲面结构导致收敛慢。标准优化器如Adam通过估计一阶和二阶矩引导优化,但直接使用会增加三倍显存开销,抵消了MeZO的优势。为此,我们提出AdaMeZO,一种零阶优化器,在不维护动量和二阶矩的情况下,仍实现类似Adam的曲率自适应能力。理论分析与大量实验验证表明,AdaMeZO性能优于MeZO,同时最多减少70%的前向传播次数。轨迹可视化进一步证实其能有效适应多样化的损失曲面。
原文摘要 · Abstract (English)
Fine-tuning LLMs is necessary for various dedicated downstream tasks, but classic backpropagation-based fine-tuning methods require substantial GPU memory. To this end, a recent work, MeZO, which relies solely on forward passes to fine-tune LLMs, significantly reduces GPU requirements at the cost of slower convergence due to its indifference to loss landscapes. Standard solutions, such as Adam, explore loss landscapes by estimating the first- and second-order moments and storing them in memory to guide the model's movement through dimensions with lower curvature and vice versa. However, directly applying Adam negates MeZO's advantage as it will triple the memory requirement. In light of this, we propose AdaMeZO, a zeroth-order optimizer that leverages Adam-style first- and second-moment estimates without maintaining them in memory. We present a theoretical analysis of AdaMeZO, corroborated by extensive experiments demonstrating AdaMeZO's performance, showing that AdaMeZO can outperform MeZO while requiring up to $70\%$ fewer forward passes. Trajectory visualizations affirm AdaMeZO's ability to adapt to diverse loss landscapes.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。