PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2402.11592v3cs.LGcs.CL

TL;DR

Large-model fine-tuning faces substantial memory overhead from back-propagation. The paper benchmarks a broad range of BP-free zeroth-order methods and introduces further enhancements, identifying task alignment and other optimization principles while advancing memory-efficient fine-tuning.

  • Problem

    Back-propagation for first-order LLM fine-tuning creates substantial memory overhead, motivating BP-free methods and broader evidence beyond MeZO’s ZO-SGD.

  • Method

    The paper benchmarks six BP-free or ZO methods across five LLM families, three task complexities, and five fine-tuning schemes, then proposes block-wise, hybrid ZO–FO, and sparsity-induced enhancements.

  • Results

    The benchmark reveals task alignment, forward gradient, and trade-offs among algorithm complexity, accuracy, query efficiency, and memory efficiency as important ZO fine-tuning principles.

  • Takeaways & Limitations

    The study offers a direction for improving the memory efficiency of LLM fine-tuning through broader ZO optimization and targeted enhancements.

  • Takeaways & Limitations

    The paper’s potential impacts depend on how the resulting fine-tuned LLMs are used, including possible misuse for misinformation, phishing, or releasing private and copyrighted information.

Abstract

from arXiv · show

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 .

1. Introduction

The paper addresses the memory overhead of back-propagation in LLM fine-tuning by benchmarking broader zeroth-order methods and proposing enhancements for accuracy and efficiency.

  • 12× more memory is required to compute OPT-13B gradients than for model inference, motivating memory-efficient LLM fine-tuning.
  • The study benchmarks six BP-free or ZO methods across five LLM families, three task complexities, and five full-parameter or PEFT schemes.
  • The benchmark examines task alignment, forward gradient, and trade-offs among algorithm complexity, accuracy, query efficiency, and memory efficiency.
  • Beyond assessment, the paper introduces block-wise, hybrid ZO–FO, and sparsity-induced ZO techniques to improve accuracy while maintaining memory efficiency.

2. Related Work

Related work frames zeroth-order optimization as a BP-free alternative for large-model training, while prior LLM studies demonstrate scalability but leave broader method comparisons open.

  • ZO optimization estimates gradients from function values and has been applied across machine-learning problems, but its use has primarily remained at small model scales.
  • High variance and slow convergence worsen with model dimension, motivating acceleration through historical data, gradient estimators, and random perturbations.
  • Forward-gradient learning provides another BP-free alternative based on forward-mode automatic differentiation for memory-intensive large-model training.
  • MeZO extended ZO-SGD to LLMs with up to 60 billion parameters and was reported as competitive with first-order and structured fine-tuning methods.

3. Reviewing ZO Optimization and Beyond

This section reviews randomized zeroth-order optimization and related BP-free methods, then motivates a broader benchmark using distinct update rules and task-aligned LLM fine-tuning.

  • Basics of ZO optimization: ZO optimization replaces first-order gradients with function-value estimates while retaining the corresponding optimizer’s algorithmic structure.
  • Basics of ZO optimization: Randomized gradient estimation uses finite differences along random directions; its variance is roughly O(d/q), trading estimation variance against query complexity.
  • Forward gradient: Forward-Grad uses forward-mode automatic differentiation without backward evaluation and can serve as a BP-free performance reference for ZO optimization.
  • A focused spectrum of ZO optimization methods: The benchmark covers ZO-SGD, sign-based, momentum, conservative-update, and ZO-Adam variants selected for minimal implementation changes and distinct update characteristics.
  • Task alignment in ZO optimization for LLM fine-tuning: Task alignment converts downstream classification into pre-training-like next-token prediction, and removing prompts caused about 10% SST2 and 8% RTE accuracy degradation for ZO methods.

4. LLM Fine-Tuning Benchmarking

The benchmark evaluates BP-free and FO optimizers across tasks, models, and fine-tuning schemes, showing that ZO performance depends strongly on task difficulty, tuning strategy, and algorithm choice. ZO methods offer substantial memory advantages, especially for long sequences, but their accuracy and scalability vary across settings.

  • Results: ZO-Adam achieves the best performance in 4 of 8 fine-tuning settings, but its higher algorithmic complexity increases memory consumption.The benchmark compares multiple ZO methods across fine-tuning settings and notes the memory cost of ZO-Adam.
  • Results: Forward-Grad is competitive with ZO methods, particularly for full fine-tuning, although its advantage diminishes as the fine-tuning scheme becomes simpler.The benchmark identifies Forward-Grad as a compelling BP-free baseline, especially for larger-scale problems.
  • Results: ZO performance varies substantially across settings: ZO-Adam can degrade dramatically, ZO-SGD is not consistently best, and ZO-SGD-Sign is usually weakest except for Prompt tuning.ZO-SGD-Cons and ZO-SGD-MMT also show strong performance, highlighting unstable relative rankings across scenarios.
  • Results: On LoRA, COPA performance can match or exceed FO methods for some models, while WinoGrande retains a 5% ∼6% FO–ZO performance gap across models.Forward-Grad and ZO-Adam outperform the best FO method on OPT-13B and Vicuna-7B for COPA, whereas WinoGrande remains harder for ZO methods.
  • Ablation study on query budget: Forward-Grad surpasses ZO-SGD by 1% ∼2% when the query number exceeds 500 and approaches FO-SGD, while computation cost scales linearly with the query budget.Both methods improve with more queries, but Forward-Grad improves more substantially.
  • Memory efficiency: ZO methods are generally more memory-efficient than FO methods, with larger advantages at longer sequence lengths because they avoid storing intermediate activations and full-model gradients.Empirical results align with the theoretical analysis, and the memory advantage strengthens as sequence length increases.

5. Extended Study to Improve ZO Fine-Tuning

The study introduces block-wise, hybrid ZO-FO, and sparsity-induced enhancements to improve ZO fine-tuning while preserving memory efficiency. Results show gains from block-wise updates and a controllable performance–memory trade-off in hybrid training.

  • Block-wise ZO fine-tuning: ZO-SGD-Block outperforms ZO-SGD across query budgets and fine-tuning tasks, showing the benefit of block-wise ZO tuning.The comparison uses MeZO as the q = 1 baseline and matches query complexity with a full-model variant.
  • Hybrid ZO-FO fine-tuning: Hybrid ZO-FO training confines ZO optimization to shallow layers and FO optimization to deep layers to trade memory cost against fine-tuning accuracy.The split is controlled by the number of layers assigned to each optimization method.
  • Hybrid ZO-FO fine-tuning: Using ZO on only the first third of model layers yields performance comparable to full FO optimization while reducing memory usage by approximately 10%.Using ZO on at least half the layers produces performance similar to full ZO fine-tuning.
  • Gradient pruning: Gradient pruning induces sparse parameter perturbations to reduce variance in random gradient estimation and improve ZO fine-tuning performance.The study evaluates sparsity ratios on OPT-1.3B using COPA and SST2, with 0% sparsity representing vanilla ZO-SGD.

6. Conclusion

The paper broadens ZO optimization for memory-efficient LLM fine-tuning through benchmarking and new algorithmic enhancements. It identifies task alignment and forward gradients as important principles and reports improved accuracy while maintaining memory efficiency.

  • 6. Conclusion: The study benchmarks diverse ZO methods across tasks and metrics, revealing task alignment and forward gradient as overlooked optimization principles.It evaluates accuracy and efficiency across varied ZO optimization settings.
  • 6. Conclusion: Block-wise descent, hybrid ZO and FO training, and gradient sparsity are proposed to improve ZO fine-tuning accuracy while maintaining memory efficiency.These techniques extend the benchmark findings into algorithmic enhancements.

Impact Statement

The paper studies memory-efficient LLM fine-tuning and notes potential benefits for energy consumption and resource-constrained environments. It also acknowledges possible misuse of resulting fine-tuned models, while identifying no specific societal consequences directly stemming from the technical work.

  • Impact Statement: Memory-efficient LLM fine-tuning could reduce energy consumption and support improved performance in resource-constrained environments.The statement presents these as potential positive impacts of the paper’s technical direction.
  • Impact Statement: Fine-tuned models could potentially be misused for misinformation, phishing attacks, or releasing copyrighted and private information.The impact statement frames these as potential negative consequences of model use.
  • Impact Statement: The paper identifies no specific societal consequences directly stemming from its technical focus.The impact discussion limits its claims to the paper’s technical scope.

A. Zeroth-Order Optimization Algorithms

Zeroth-order methods estimate gradients from function evaluations, avoiding derivative computation and back-propagation. The section summarizes several ZO optimizers, forward gradients, PEFT schemes, and the memory components governing FO/ZO training.

  • ZO gradient estimation: Random gradient estimation approximates gradients using function evaluations, with Gaussian direction u and small scalar µ defining the estimator.This surrogate enables gradient-based optimization without directly computing derivatives.
  • ZO algorithms: ZO-SGD directly updates parameters with the estimated gradient, while ZO-Sign-SGD applies a sign operation to make updates more robust to coordinate-wise gradient noise.The sign operation is intended to mitigate the effects of large-variance noise.
  • ZO algorithms: Momentum-based ZO-SGD uses moving averages to estimate the global gradient, whereas ZO-Adam additionally adapts the learning rate through a diagonal normalization matrix.These mechanisms target variance and noise in stochastic gradient estimates.
  • Forward Gradient: Forward Gradient estimates directional derivatives with JVPs during the forward pass, reducing computation and memory relative to back-propagation-based approaches.Although it uses first-order gradients, its forward-pass implementation keeps memory complexity comparatively low.
  • Parameter-efficient fine-tuning: LoRA, prompt tuning, and prefix tuning train compact parameter subsets while keeping the pretrained model parameters frozen.LoRA updates low-rank matrices, prompt tuning learns prompt embeddings, and prefix tuning updates attention keys and values.
  • Memory efficiency: Optimizer memory consists of constant model and optimizer states plus dynamic forward, backward, and optimization allocations during training.The peak memory depends on the maximum number of variables stored at a given point in the update pipeline.

C.1. Theoretical Memory Efficiency Analysis of Different Optimizers

The analysis contrasts memory states for FO, ZO, and forward-gradient optimizers, showing why ZO avoids storing back-propagation activations while identifying implementation limits for stateless forward gradients.

  • FO-SGD versus ZO-SGD: FO-SGD stores model, activation, gradient, and optimizer states, whereas ZO methods estimate gradients without the backward activation state.The FO analysis defines stored model and optimizer states, while vanilla ZO-SGD uses random vectors and projected gradients from forward passes.
  • ZO-SGD: ZO-SGD replaces stored random vectors with their random seed, regenerating each vector on demand to reduce memory consumption.The paper uses this memory-reduction trick by default and refers to the resulting implementation as ZO-SGD.
  • ZO-SGD-Momentum: ZO-SGD with momentum requires additional momentum storage matching the model-parameter size, increasing memory relative to ordinary ZO-SGD.Momentum is treated as an optimizer state in the memory accounting.
  • Forward Gradient: Forward Gradient uses forward-mode AD and Jacobian-vector products, requiring random vectors, projected gradients, and intermediate forward results.The resulting forward-state cost includes intermediate activations, while the backward state contains gradients estimated from the random direction.
  • Forward Gradient: PyTorch does not support stateless Forward Gradient through its built-in forward-mode AD, although stateless computation remains possible by regenerating layerwise random vectors on demand.The paper leaves more memory-efficient practical Forward Gradient computation as future work.

C.2. Other Implementation Details

The implementation details distinguish F16, FP16, and optimizer-specific memory behavior, including temporary mixed-precision model copies and extra Adam update storage.

  • Precision settings: F16 loads models at 16-bit precision for memory efficiency, but it cannot be used with Forward Gradient or FO methods in the stated automatic-differentiation setup.The paper applies F16 to ZO methods that do not require differentiation.
  • Precision settings: FP16 mixed precision accelerates FO gradient computation and reduces intermediate-result memory, but does not reduce stored model or gradient memory.The paper distinguishes this from F16 model loading.
  • Adam implementation: Adam's foreach implementation merges layer weights into one tensor during updates, requiring extra memory despite faster computation.This overhead is separate from the model and activation terms in the optimizer comparison.
  • Memory comparison setup: The full-finetuning comparison uses OPT-13B on MultiRC with examples averaging 400 tokens and reports theoretical and empirical memory values together.The table caption states that the theoretical and empirical values mutually corroborate one another.

C.3. Additional Experiments

Additional experiments compare empirical memory costs for full and LoRA fine-tuning, finding that the observed pattern agrees with the theoretical analysis and favors ZO methods.

  • Empirical memory results: ZO methods are generally much more memory-efficient than their FO counterparts, and loading the model in F16 further reduces overhead.This conclusion is reported for both full and LoRA fine-tuning comparisons.
  • LoRA fine-tuning: The LoRA comparison evaluates OPT-13B on MultiRC with examples averaging 400 tokens, omitting adapter-parameter memory because it is assumed negligible relative to model and activation memory.The table compares memory consumption under the same general settings as the full-finetuning comparison.
Loading 2402.11592v3…