A da Z eta: Adaptive Zeroth-Order Tensor-Train Adaption for Memory-Efficient Large Language Models Fine-Tuning (2024)

Abstract: Fine-tuning large language models (LLMs) has achieved remarkable performance across various natural language processing tasks, yet it demands more and more memory as model sizes keep growing. To address this issue, the recently proposed Memory-efficient Zeroth-order (MeZO) methods attempt to fine-tune LLMs using only forward passes, thereby avoiding the need for a backpropagation graph. However, significant performance drops and a high risk of divergence have limited their widespread adoption. In this paper, we propose the Adaptive Zeroth-order Tensor-Train Adaption (AdaZeta) framework, specifically designed to improve the performance and convergence of the ZO methods. To enhance dimension-dependent ZO estimation accuracy, we introduce a fast-forward, low-parameter tensorized adapter. To tackle the frequently observed divergence issue in large-scale ZO fine-tuning tasks, we propose an adaptive query number schedule that guarantees convergence. Detailed theoretical analysis and extensive experimental results on Roberta-Large and Llama-2-7B models substantiate the efficacy of our AdaZeta framework in terms of accuracy, memory efficiency, and convergence speed.

## introduction fine-tuning large language models (llms) has demonstrated outstanding performance in addressing numerous natural language processing applications, such as natural language understanding (kenton and toutanova, 2019), question-answering @xcite , and summarization @xcite . however, as the size of llms increases, the training process consumes progressively more gpu memory. in recent years, approaches such as quantization @xcite and parameter-efficient fine-tuning (peft) @xcite have been proposed to reduce memory costs during training by storing data with lower bit-depth or updating only a portion of the parameters. despite these strategies effectively reducing memory costs, overall memory usage remains high due to the continuous reliance on a backpropagation graph. to further reduce the memory overhead, @xcite proposed the memory-efficient zeroth-order (mezo) method for llm fine-tuning, which shows over 8× memory reduction compared with the first-order (fo) fine-tuning methods like sgd @xcite and adamw @xcite . unlike fo methods, which calculate gradients via backpropagation, the mezo method estimates gradients based on the difference between loss values obtained from two forward passes, thereby eliminating the need for a backpropagation graph. however, two main challenges persist in the zeroth-order (zo) fine-tuning of llms: 1) a significant performance gap between fo and zo approaches, and 2) increased risk of divergence, particularly in the zo fine-tuning of large-scale llms, as observed in recent studies @xcite . to improve the performance, various fo optimization techniques have been adapted for zo fine-tuning scenarios, like the zo-adamu method @xcite . however, these approaches fail to accommodate the specific needs of zo methods, and add significant memory overhead from the optimizer state. given the dimensionality-related nature of zo convergence rates, @xcite propose the sparse-mezo method that generates pruning masks based on the value of the weight elements. nevertheless, the sparse-mezo method yields inconsistent performance across various tasks and hyperparameter configurations. in contrast to this approach, we consider using the peft method to reduce the number of trainable parameters. although the zo peft method like mezo-lora has been considered in @xcite , the improvements are limited as the lora adapter fails to offer high representational ability with an ultra-low rank. to solve this problem, we involve tensorized adapters, which offer high performance with even lower trainable parameters than lora adapters. to address the variance-related divergence issue in large-scale zo fine-tuning, previous studies @xcite have primarily focused on adjusting the batch size, as increasing the batch size can reduce the noise in zo gradient estimation. however, these approaches introduce significant runtime overhead and fail to improve performance significantly. to further reduce variance, @xcite introduced the mezo-svrg method, adapting the first-order svrg technique to the zo context. despite its success, mezo-svrg suffers from a slow and memory-inefficient fine-tuning process due to the additional parameter copies and computation process that even doubles the memory cost of the mezo methods. in contrast to these works, we consider reducing the zo gradient variance with a sublinearly increasing query foot_1 schedule that achieves not only better accuracy but also faster convergence in terms of both steps and time. this paper explores task-specific peft training for zo fine-tuning scenarios. we introduce the adaptive zeroth-order tensor-train adaption (adazeta) framework, which incorporates fastforward tensorized adapters and an adaptive query schedule. this combination can significantly enhance the accuracy and convergence of zo fine-tuning, as demonstrated in fig. 1 . our contributions are summarized as follows: ## methods in this section, we first introduce some basic knowledge of the zo gradient estimator. then, we present our adazeta method, a powerful framework designed to improve the performance of zo llm fine-tuning with two main components: 1) the fast-forward tensorized adapters, and 2) an adaptive query number schedule. finally, we provide a theoretical analysis of the convergence rate of the adazeta method, demonstrating the improved convergence rate theoretically. ## experiments in this section, we conduct comprehensive experiments to evaluate the performance of our proposed adazeta framework across several llms with different scales on a variety of natural language understanding and generation tasks @xcite @xcite . we demonstrate that our methods surpass a comprehensive array of memory-efficient baselines, including inference-only methods such as zero-shot @xcite , in-context learning (icl), and linear probing (lp) @xcite , as well as zo fine-tuning methods like mezo, mezo-lora @xcite , and sparse-mezo @xcite . also, the first-order fine-tuning (ft) baseline is also provided as a reference. initially, we present experimental evidence using roberta-large models @xcite , illustrating that the integration of tensorized adapters can significantly enhance the efficiency of zo fine-tuning by reducing the number of trainable parameters. subsequently, we enabled our proposed adaptive query schedule method to show the effectiveness of the adazeta framework on large-scale llama-2-7b models @xcite , which not only enhances performance but also ensures robust convergence. all experiments are conducted on nvidia tesla a100-40gb gpus, with further details about the experimental setup available in appendix a. ## 992 finally, we need to bound the variance λ 2 , related to the zo gradient estimation noise. taking expectation with respect to the i.i.d. random perturbation vector u, we have: finally, we need to bound the term e u [∥ĝ(w k ; b i , u i,1 ) -∇ℓ(w k )∥], which gives: putting it all together we can obtain the upper bound for the gradient noise ∥h k ∥. now we begin to present the proof of theorem 1: we start from the gradient updating rule in the adazeta algorithm, which gives w t+1 = w t -η∇ l(w k ). by using taylor's theorem on the exact smoothed loss ℓ(w k ), we have: taking expectations on both sides gives: taking the norms on both sides and using the lipschitz condition, we have: finally, since this must hold for any y and x, and since the norm of the hessian matrix is the supremum of v 2 f (c)(yx) /∥y -x∥ for non-zero yx, it follows that: rearrange and we obtain: taking summation over steps @xmath0, • • • , k gives: to bound the last term, we first represent the noise gradient ∇ lk (w k ) 2 as a combination of the true gradient and the gradient noise introduced in eq. ( 5 ), which gives: taking eq. ( 29 ) back into eq. ( 26 ), using the results from lemma 1, taking the expectation over all randomness and average over the maximum steps k, we obtain: