Revisiting Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning: A Benchmark (2402.11592v3)

Abstract: In the evolving landscape of NLP, fine-tuning pre-trained 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 .

Enhancing Memory Efficiency in Fine-Tuning LLMs through Zeroth-Order Optimization

Overview

Fine-tuning pre-trained LLMs is a pervasive practice in natural language processing tasks. However, the substantial memory overhead associated with gradient computation through back-propagation remains a significant barrier, particularly for computational platforms with limited memory. This challenge has motivated a shift towards memory-efficient approaches, such as Zeroth-Order (ZO) optimization which eliminates the necessity for explicitly computing gradients. Building on the concept introduced by Malladi et al. (2023), this paper proposes a comprehensive analysis of Zeroth-Order Optimization for memory-efficient LLM fine-tuning, unveiling previously overlooked optimization principles and introducing novel enhancements.

Related Work and Theoretical Background

Previous efforts in Parameter-Efficient Fine-Tuning (PEFT) strategies and zeroth-order optimization have laid the groundwork for memory-efficient model training. Traditional approaches like Adapter-based methods, Low-Rank Adaptation (LoRA), and prompt tuning significantly reduce the number of parameters required for fine-tuning but still require considerable memory for gradient computation. In contrast, Zeroth-Order (ZO) optimization utilizes function value-based gradient estimation, thereby circumventing the need for back-propagation and subsequently reducing memory usage. Despite its promise, the exploration of ZO optimization techniques beyond basic ZO-Stochastic Gradient Descent (ZO-SGD) is scant, prompting this paper.

Methodology and Key Contributions

1. Benchmark Creation: The paper creates the first benchmark for ZO optimization in LLM fine-tuning, evaluating six BP-free ZO optimization methods across five LLM families, three task complexities, and five fine-tuning schemes.
2. Insights on Optimization Principles: The benchmark paper reveals critical insights including the importance of task alignment, the utility of the forward gradient method as a baseline for ZO optimization, and the balance between algorithm complexity and fine-tuning performance.
3. Enhancements to ZO Optimization: Drawing from the underline insights, the paper proposes techniques of block-wise descent, hybrid ZO and FO (First-Order) training, and gradient sparsity to improve ZO optimization-based LLM fine-tuning.

Theoretical and Practical Implications

From a theoretical standpoint, this work advances understanding of the optimization landscape for LLM fine-tuning, particularly under resource constraints. Practically, the introduced benchmark and ensuing insights offer a structured foundation for future research and development in memory-efficient fine-tuning methods. The proposed enhancements to ZO optimization—block-wise descent, hybrid training, and gradient sparsity—not only improve fine-tuning accuracy but also maintain memory efficiency. These advancements possess the potential to facilitate on-device training and deployment of sophisticated LLMs in memory-constrained environments.

Future Directions

Looking ahead, the exploration of further ZO optimization methods and their combinations with established PEFT strategies presents a promising avenue for research. Additionally, investigating the applicability of these memory-efficient fine-tuning techniques beyond LLMs to other domains of deep learning could broaden their utility.

Concluding Thoughts

This paper's comprehensive benchmarking and innovative enhancements to ZO optimization mark significant steps towards overcoming the memory limitations in fine-tuning LLMs. By elucidating the trade-offs between algorithm complexity, accuracy, and memory efficiency, it lays the groundwork for more sustainable and accessible AI models, pushing the boundaries of what's possible within constrained computational environments.