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 in-place, 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 12x memory reduction and up to 2x 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.
翻译:微调语言模型(LM)已在多种下游任务中取得成功,但随着模型规模的增长,反向传播需要极大的内存开销。零阶方法原则上只需两次前向传播即可估计梯度,但理论上对大规模模型的优化速度极慢。本文提出一种内存高效的零阶优化器(MeZO),通过改进经典ZO-SGD方法实现原位优化,使语言模型微调的内存消耗与推理阶段相当。例如,单张A100 80GB GPU上,MeZO可训练300亿参数模型,而反向传播微调在相同预算下仅能训练27亿参数模型。我们针对模型类型(掩码与自回归语言模型)、模型规模(高达660亿参数)及下游任务(分类、多项选择与生成)开展全面实验。结果表明:(1)MeZO显著优于上下文学习和线性探测;(2)MeZO在多数任务中性能与反向传播微调相当,内存占用减少12倍,GPU耗时减少2倍;(3)MeZO兼容全参数微调及LoRA、前缀微调等参数高效微调技术;(4)MeZO可有效优化不可微分目标(如准确率或F1值)。我们通过理论分析支持实证发现,揭示充分的预训练和任务提示如何使MeZO成功微调大型模型,尽管经典ZO理论分析给出相反结论。