Fine-Tuning Language Models with Just Forward Passes

Jason Lee, Sanjeev Arora, Sadhika Malladi, Danqi Chen, Alex Damian, Tianyu Gao, Eshaan Nichani

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 17 functions out of this paper's own repositories and ran 9 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

Repositories linked to this paper

• princeton-nlp/mezo — marked official in Papers with Code; named in the paper
• liangyuwang/zo2
• mathisall/zo-adamu

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a prohibitively large amount of memory. Zeroth-order (ZO) methods can in principle estimate gradients using only two forward passes but are theorized to be catastrophically slow for optimizing large models. In this work, we propose a memory-efficient zerothorder optimizer (MeZO), adapting the classical ZO-SGD method to operate inplace, thereby fine-tuning LMs with the same memory footprint as inference. For example, with a single A100 80GB GPU, MeZO can train a 30-billion parameter model, whereas fine-tuning with backpropagation can train only a 2.7B LM with the same budget. We conduct comprehensive experiments across model types (masked and autoregressive LMs), model scales (up to 66B), and downstream tasks (classification, multiple-choice, and generation). Our results demonstrate that (1) MeZO significantly outperforms in-context learning and linear probing; (2) MeZO achieves comparable performance to fine-tuning with backpropagation across multiple tasks, with up to 12× memory reduction and up to 2× GPU-hour reduction in our implementation; (3) MeZO is compatible with both full-parameter and parameter-efficient tuning techniques such as LoRA and prefix tuning; (4) MeZO can effectively optimize non-differentiable objectives (e.g., maximizing accuracy or F1). We support our empirical findings with theoretical insights, highlighting how adequate pre-training and task prompts enable MeZO to fine-tune huge models, despite classical ZO analyses suggesting otherwise. 1

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2305.17333")
get_code_for_paper("2305.17333")
have("2305.17333")

Connect an agent — have() is free.