AdaMeZO: Adam-style LLM fine-tuning without storing gradient moments in GPU memory

Researchers Zhijie Cai, Haolong Chen, and Guangxu Zhu have introduced AdaMeZO, a zeroth-order optimizer that brings the benefits of the popular Adam algorithm to large language model fine-tuning — without requiring gradient moments to be stored in GPU memory.

Why is GPU memory a bottleneck for LLM fine-tuning?

The standard Adam optimizer, routinely used for neural network training, tracks two statistics for each model parameter: the first moment (running mean of gradients) and the second moment (running mean of squared gradients). For models with billions of parameters, this triples memory requirements. MeZO, a prior approach that uses only forward passes without computing true gradients, solves the memory problem — but converges more slowly because it lacks per-parameter adaptive learning rate adjustment.

How does AdaMeZO combine both approaches?

AdaMeZO estimates Adam’s moments without storing them persistently: it uses random weight perturbations and measures loss changes to reconstruct adaptive moment behavior on-the-fly, separately for each optimization step. The result is an optimizer that behaves like Adam — adjusting the learning rate according to the estimated loss surface geometry — while keeping the same memory footprint as MeZO.

What do the trajectory visualizations show?

The authors present optimization trajectory visualizations on various loss surfaces demonstrating how AdaMeZO adaptively navigates flat and curved regions of parameter space, in contrast to MeZO’s more uniform behavior. Quantitatively, AdaMeZO achieves up to 70% fewer forward passes to reach the same performance level, with improved convergence compared to the original MeZO.