FZOO让大模型微调速度逼近Adam,同时大幅降低显存占用。
FZOO: Fast Zeroth-Order Optimizer for Fine-Tuning Large Language Models towards Adam-Scale Speed
- 用批量单边估计和自适应步长减少梯度计算次数
- 相比MeZO准确率提升3%,前向传播次数减少3倍
- 适合单卡高精度微调,尤其适配参数高效训练技术
大模型微调常因Adam等一阶优化器的反向传播导致显存暴涨(如OPT-30B达633GB)。零阶优化器通过仅用前向传播估算梯度可缓解此问题,但传统方法如MeZO收敛慢。本文提出FZOO,一种接近Adam速度的快速零阶优化器。FZOO采用批量单边梯度估计并根据批次损失标准差自适应步长,结合Rademacher随机扰动与CUDA并行加速每批次计算。在RoBERTa-large、OPT(350M-66B)、Phi-2、Llama3等11个任务上验证,平均比MeZO准确率高3%,前向传播次数减少3倍;对RoBERTa-large,准确率提升5.6%,前向传播次数减少18倍,收敛速度媲美Adam。理论分析证明FZOO等价于归一化SGD并具备收敛性保障。其可无缝集成到PEFT中,进一步节省显存。结果表明,单卡全参数微调已可行,并为后续内存高效预训练指明方向。
原文摘要 · Abstract (English)
Fine-tuning large language models (LLMs) often faces GPU memory bottlenecks: the backward pass of first-order optimizers like Adam increases memory usage to more than 10 times the inference level (e.g., 633 GB for OPT-30B). Zeroth-order (ZO) optimizers avoid this cost by estimating gradients only from forward passes, yet existing methods like MeZO usually require many more steps to converge. Can this trade-off between speed and memory in ZO be fundamentally improved? Normalized-SGD demonstrates strong empirical performance with greater memory efficiency than Adam. In light of this, we introduce FZOO, a Fast Zeroth-Order Optimizer toward Adam-Scale Speed. FZOO reduces the total forward passes needed for convergence by employing batched one-sided estimates that adapt step sizes based on the standard deviation of batch losses. It also accelerates per-batch computation through the use of Rademacher random vector perturbations coupled with CUDA's parallel processing. Extensive experiments on diverse models, including RoBERTa-large, OPT (350M-66B), Phi-2, and Llama3, across 11 tasks validate FZOO's effectiveness. On average, FZOO outperforms MeZO by 3 percent in accuracy while requiring 3 times fewer forward passes. For RoBERTa-large, FZOO achieves average improvements of 5.6 percent in accuracy and an 18 times reduction in forward passes compared to MeZO, achieving convergence speeds comparable to Adam. We also provide theoretical analysis proving FZOO's formal equivalence to a normalized-SGD update rule and its convergence guarantees. FZOO integrates smoothly into PEFT techniques, enabling even larger memory savings. Overall, our results make single-GPU, high-speed, full-parameter fine-tuning practical and point toward future work on memory-efficient pre-training.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。