AdaMeZO: Adam-style Zeroth-Order Optimizer for LLM Fine-tuning Without Maintaining the Moments

TL;DR

AdaMeZO is a novel zeroth-order optimizer for fine-tuning large language models that mimics Adam's moment estimation without extra memory, achieving faster convergence with fewer forward passes.

Contribution

It introduces AdaMeZO, an Adam-style zeroth-order optimizer that estimates moments without storing them, reducing memory use and improving convergence over prior zeroth-order methods.

Findings

AdaMeZO outperforms MeZO in convergence speed.

Requires up to 70% fewer forward passes than MeZO.

Effectively adapts to diverse loss landscapes.

Abstract

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.