arXiv 2402.115922024-05-28cs.LGcs.CL

重新审视用于内存高效LLM微调的零阶优化：一个基准

Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark

Michigan State University（密歇根州立大学）；The University of North Carolina at Chapel Hill（北卡罗来纳大学教堂山分校）；UT Austin（德克萨斯大学奥斯汀分校）；University of Minnesota Twin Cities（明尼苏达大学双城分校）；IBM Research（IBM研究院）；Princeton University（普林斯顿大学）；DAMO Academy, Alibaba Group（阿里巴巴达摩院）；MIT（麻省理工学院）；Harvard University（哈佛大学）

Yihua Zhang, Pingzhi Li, Junyuan Hong, Jiaxiang Li, Yimeng Zhang, Wenqing Zheng, Pin-Yu Chen, Jason D. Lee, Wotao Yin, Mingyi Hong, Zhangyang Wang, Sijia Liu, Tianlong Chen

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

本文提出转向无反向传播的零阶优化以降低LLM微调内存成本，通过首创性基准研究揭示优化原则并引入分块下降等增强技术，实现更高效的内存微调。

摘要：

在不断发展的自然语言处理（NLP）领域，使用SGD和Adam等一阶（FO）优化器对预训练大型语言模型（LLMs）进行微调已成为标准。然而，随着LLMs规模的增长，用于FO梯度计算的反向传播（BP）带来的巨大内存开销构成了重大挑战。解决这个问题至关重要，特别是对于设备端训练等内存效率至关重要的应用。本文提出转向无BP的零阶（ZO）优化，作为降低LLM微调内存成本的解决方案，基于MeZO引入的初步概念。与传统的ZO-SGD方法不同，我们的工作通过一项全面的首创性基准研究，将探索扩展到更广泛的ZO优化技术，涵盖五个LLM系列（Roberta、OPT、LLaMA、Vicuna、Mistral）、三种任务复杂度和五种微调方案。我们的研究揭示了先前被忽视的优化原则，强调了任务对齐的重要性、前向梯度方法的作用，以及算法复杂度与微调性能之间的平衡。我们进一步引入了对ZO优化的新颖增强，包括分块下降、混合训练和梯度稀疏性。我们的研究为实现进一步内存高效的LLM微调提供了有前景的方向。复现所有实验的代码位于https://github.com/ZO-Bench/ZO-LLM。

英文摘要：

In the evolving landscape of natural language processing (NLP), fine-tuning pre-trained Large Language Models (LLMs) with first-order (FO) optimizers like SGD and Adam has become standard. Yet, as LLMs grow {in size}, the substantial memory overhead from back-propagation (BP) for FO gradient computation presents a significant challenge. Addressing this issue is crucial, especially for applications like on-device training where memory efficiency is paramount. This paper proposes a shift towards BP-free, zeroth-order (ZO) optimization as a solution for reducing memory costs during LLM fine-tuning, building on the initial concept introduced by MeZO. Unlike traditional ZO-SGD methods, our work expands the exploration to a wider array of ZO optimization techniques, through a comprehensive, first-of-its-kind benchmarking study across five LLM families (Roberta, OPT, LLaMA, Vicuna, Mistral), three task complexities, and five fine-tuning schemes. Our study unveils previously overlooked optimization principles, highlighting the importance of task alignment, the role of the forward gradient method, and the balance between algorithm complexity and fine-tuning performance. We further introduce novel enhancements to ZO optimization, including block-wise descent, hybrid training, and gradient sparsity. Our study offers a promising direction for achieving further memory-efficient LLM fine-tuning. Codes to reproduce all our experiments are at https://github.com/ZO-Bench/ZO-LLM .

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。