Zeroth-Order Fine-Tuning of LLMs with Extreme Sparsity

Authors: wentao  guo
jikai  long
yimeng  zeng
zirui  liu
xinyu  yang
yide  ran
jacob r gardner
osbert  bastani
christopher  de sa
xiaodong  yu
beidi  chen
zhaozhuo  xu
Venue: ICML_cc_2024_Workshop_WANT
Published: 2024-05-24 06:16:34

Abstract: Zeroth-order optimization (ZO) is a memory-efficient strategy for fine-tuning Large Language Models using only forward passes. However, the application of ZO fine-tuning in memory-constrained settings such as mobile phones and laptops is still challenging since full precision forward passes are infeasible. In this study, we address this limitation by integrating sparsity and quantization into ZO fine-tuning of LLMs. Specifically, we investigate the feasibility of fine-tuning an extremely small subset of LLM parameters using ZO. This approach allows the majority of un-tuned parameters to be quantized to accommodate the constraints of limited device memory. Our findings reveal that the pre-training process can identify a set of "sensitive parameters" that can guide the ZO fine-tuning of LLMs on downstream tasks. Our results demonstrate that fine-tuning 0.1% sensitive parameters in the LLM with ZO can outperform the full ZO fine-tuning performance, while offering wall-clock time speedup. Additionally, we show that ZO fine-tuning targeting these 0.1% sensitive parameters, combined with 4 bit quantization, enables efficient ZO fine-tuning of an Llama2-7B model on a GPU device with less than 8GiB of memory and notably reduced latency.

Main paper: 
 introduction: 
 Large language models (LLMs) have demonstrated superior performance in general-purpose language generation (Brown et al., 2020; Radford et al., 2019; Liu et al., 2019) . Despite their success, it remains necessary to fine-tune LLMs for specific tasks to achieve optimal results. However, 1 Princeton University 2 Cornell University 3 Stevens Institute of Technology 4 University of Pennsylvania 5 Rice University 6 Carnegie Mellon University. Correspondence to: Wentao Guo <wg0420@princeton.edu>, Zhaozhuo Xu <zxu79@stevens.edu>.Accepted to the Workshop on Advancing Neural Network Training at International Conference on Machine Learning (WANT@ICML 2024). fine-tuning LLMs often requires much more memory compared to the inference process. Specifically, there are mainly four parts that occupy the memory during fine-tuning LLMs:(1) the weight parameter itself; (2) the optimizer state, which contains the information about the past gradient (Kingma & Ba, 2015) ; (3) the weight gradient used to update the parameters; (4) the activation cached to calculate the weight gradient (Liu et al., 2024b) ; In previous work like QLoRA (Dettmers et al., 2023) , it can reduce both (1) and (2) by combining weight quantization and low-rank adaption (Hu et al., 2021) , which enables fine-tuning huge LLMs under data-center level GPUs. However, on memory-constrained hardware like cell phones, the memory of caching (3) weight gradient and (4) activation required by backpropagation still cannot be overlooked. The disparity between the demand of LLM fine-tuning and hardware capacity limits the adaptability of LLMs, especially when personalizing them for edge devices.Exploring Zeroth-Order Optimization in LLM Fine-Tuning. Recently, there has been a resurging interest in zeroth-order (ZO) optimization methods for LLM finetuning (Malladi et al., 2023a; Liu et al., 2024a; Chen et al., 2024) . ZO optimization method perturbs model parameters in random directions and utilize the loss value difference to compute the gradient direction for parameter update. One advantage of ZO methods in LLM fine-tuning is that they do not require backpropagation procedures, which significantly saves the computation and memory. In this way, ZO is backpropagation-free and does not need to cache (3) weight gradients and (4) activations during fine-tuning. In practice, ZO methods have demonstrated the potential to achieve performance comparable to first-order methods in LLM fine-tuning, which opens the doors for various efficient LLM adaptation strategies.Efficient ZO LLM Fine-Tuning with Sparsity. Although ZO methods remove the need for backpropagation, a significant drawback of these methods is the slow convergence rate (Zhao et al., 2024; Liu et al., 2024a) . A recent approach addresses this by fine-tuning with a sparse mask (Liu et al., 2024a; Zhang et al., 2024) , achieving approximately ∼ 75% sparsity. Nonetheless, this sparsity level barely reduces computational overhead, as the latency during the forward pass with even ∼ 90% sparsity is still comparable to that of As the sensitive sparse fine-tuning method achieves great performance via optimizing only 0.1% parameters (performance comparable to ZO full fine-tuning and 10% random subsets), during inference we achieve an end-to-end 1.49× speedup, with 2.15× speedup at sparse operations.dense matrix operations. This latency increase can greatly impact user experience on applications such as personal assistants, where even a twofold increase in latency is perceptible. In addition, merging the sparse weights back into the base model is impractical on these devices due to memory constraints prohibiting dequantization and quantization. Empirical evidence suggests that higher sparsity levels can significantly decrease the time required for sparse matrix operations, as shown in Figure 1 . This raises the question:Is it possible to leverage the benefits of higher sparsity levels in reducing inference latency while preserving performance on downstream tasks? If so, how far can sparsity be pushed in this context?Our Proposal: ZO LLM Fine-Tuning with Fisher-Informed, Transferable Sparsity. In this paper, we answer the raised research question by proposing an efficient sparse ZO LLM fine-tuning strategy. We observe an extreme sparsity pattern in LLM parameters: a subset, determined by selecting the top k magnitude entries from the diagonal of empirical Fisher information matrix, is effective for ZO fine-tuning. Moreover, we find this sparsity pattern can be obtained through LLM's continuous pre-training process and be transferred to various downstream tasks without modification.Building on these insights, our work proposes a comprehensive framework for ZO finetuning, making the following contributions:• We identify that only an extremely small portion (0.1%) of LLM parameters should be updated during ZO LLM fine-tuning. Moreover, we utilize this insight to guide the memory-efficient on-device personalization of LLMs by low-bit quantization of model parameters.  
background and related works: 
 In this section, we present the formulation for ZO optimization. We also discuss related works about sparsity in LLMs.ZO surrogate gradient estimator. ZO optimizers have been studied widely in the machine learning community. Given a dataset D = {(x 1 , y 1 ), . . . , (x n , y n )} and a loss function f with model parameters w ∈ R d , ZO optimizer will estimate the gradient at w via ZO surrogate gradient estimator. Simultaneous Perturbation Stochastic Approximation (SPSA) (Spall, 1992) is such an estimator that would first sample a random vector z ∈ R d and uses the loss value difference to scale the update direction. z is usually sampled from an Gaussian distribution N (0, I d ).Definition 2.1 (Simultaneous Perturbation Stochastic Approximation (SPSA) (Spall, 1992) ). SPSA estimates the gradient w.r.t. w with a data example (x, y), a small constant ϵ ∈ R, and a sampled random vector z ∈ R d as follows:EQUATIONThere are other ZO surrogate gradient estimators available (Liu et al., 2020; Ohta et al., 2020) , but in practice SPSA achieves good performance in ZO optimization, particularly when fine-tuning LLMs. Some ZO algorithms such as DeepZero (Chen et al., 2024) would utilize the parameterwise finite difference of loss values to derive parameter-wise update directions. This would yield O(d) query costs per training step even when combining with certain sparse masking methods and not practical for LLM fine-tuning scenarios.We therefore select SPSA with random Gaussian perturbation as our ZO gradient estimator.ZO-SGD algorithm. ZO-SGD is an optimizer similar to SGD but replaces the FO gradient with ZO surrogate gradient estimate per training step, as defined below: Definition 2.2 (ZO-SGD update rule). ZO-SGD is an optimizer that uses ZO surrogate gradient to update parameters w t with learning rate η t and a data example (x t , y t ) sampled at timestep t:EQUATIONMeZO (Malladi et al., 2023a ) is a ZO-SGD algorithm that uses the "random seed trick" to save the need of caching ZO surrogate gradient. The choice of optimizer (SGD) is orthogonal to ZO optimization techniques, but in our preliminary experiments we find adaptive optimizers such as Adam (Kingma & Ba, 2015) would not necessarily accelerate ZO convergence in LLM fine-tuning scenarios. There are other ZO optimizers aware of the parameter-wise heterogeneity of loss curvatures to accelerate the optimization convergence (Zhao et al., 2024) , and we leave how to combine our method with theirs as future works.Sparsity-driven techniques are widely adopted in improving ML model's efficiency (Tan et al., 2024a; Xia et al., 2023; Liu et al., 2023; Peng et al., 2013; Frankle & Carbin, 2019) and robustness (Zhong et al., 2024; 2021) . Frankle & Carbin (2019) showed that within large feed-forward networks, there exists a subnetwork that, when trained in isolation, can achieve test accuracy comparable to that of the original network. In the foundation models era, Liu et al. (2023) demonstrated that transformer-based models, such as OPT (Zhang et al., 2022) , exhibit great sparsity (≥ 95%) in activations. Moreover, Panigrahi et al. (2023) discovered that for RoBERTa (Liu et al., 2019) , fine-tuning a very small subset of parameters (∼ 0.01%) can yield performance exceeding 95% of that achieved by full fine-tuning.In the context of ZO optimization, Liu et al. (2024a) and Zhang et al. (2024) also suggest that sparsity would potentially accelerate ZO optimization convergence. We believe that ZO has an intrinsic need for sparse training, as the procedure of ZO gradient estimator usually requires uniform coordinate-wise scale (in expectation) perturbation which grows with d. In tradition, people usually resolve this with knowledge from parameter-wise loss curvature heterogeneity (replace z with Σ1/2 z where Σ 1/2 serves as a Hessian-informed preconditioner) (Ye et al., 2018; Zhao et al., 2024) . However, they do not provide a comprehensive investigation on massive parameter models like LLMs. In particular, we also observe that during first-order (FO) finetuning of LLMs, the FO gradient can be quite sparse. We will elaborate more on this insight in the following section (see Figure 2 and Figure 7 ). We would like to explore how sparsity can benefit the ZO LLM fine-tuning. 
chasing extreme sparsity in zo llm fine-tuning: 
 In this section, we describe the extreme sparsity pattern we observed in LLMs and how we utilize it for efficient ZO fine-tuning including on-device personalization of LLMs. m k = argmax m ∥m ⊙ ∇f (w; (x, y))∥ 2 2 .(3)In the context of ZO optimization, we will update sensitive parameters only. Denote that z = z ⊙ m k . We will modify the SPSA gradient estimator from ĝ(w, (x, y), z) to ĝ(w, (x, y), z), and accordingly:Definition 3.2 (Sensitive sparse ZO-SGD update rule).EQUATIONThe theoretical support of sensitive parameters can be derived from the lens of SPSA gradient estimator and Fisher information matrix as follows:• Maximum zeroth-order loss value changes, from the lens of SPSA estimator.The square (account for negativity) of loss value difference for ĝw (w t , (x t , y t ), zt ) is as follows:E z{f (w + ϵz; (x, y)) -f (w -ϵz; (x, y))} 2 ≈ E z{2ϵz ⊤ ∇ w f (w; (x, y))} 2 = 4ϵ 2 ∥m k ⊙ ∇ w f (w; (x, y))∥ 2 2Since by Definition 3.1 our sensitive mask would maximize ∥m k ⊙ ∇ w f (w; (x, y))∥ 2 for a given sparsity ratio, we would expect our sensitive mask to maximize the magnitude of the loss value difference for any given sparsity ratio.• Maximum coverage of Hessian diagonal, from the lens of Fisher matrix.LLMs are often pre-trained on large text corpus 1 to reach low perplexity before entering the fine-tuning stage. In this case, we would assume p LLM (y|x) ∼ p D (y|x), which implies the empirical Fisher F should be close to the (true) Fisher matrix F as follows:F = E x∼p D ,ŷ∼pLLM(•|x) ∇ w log p LLM (ŷ|x)(∇ w log p LLM (ŷ|x)) ⊤ ≈ F = E (x,y)∼p D ∇ w log p LLM (y|x)(∇ w log p LLM (y|x)) ⊤As we assume the empirical Fisher matrix approximates Fisher, which also approximates the Hessian, and empirical i of linear layers during Llama2-7B (Touvron et al., 2023) fine-tuning. For each linear layer, we first sort parameters by the decreasing order of their gradient square value [∇F(w)] 2 i , i ∈ [d layer ], and we take the cumulative sum and normalize it to draw a blue curve, and the red-shaded region is the mean ± std of all blue curves. More similar figures are in Figure 7 . We observe that roughly 0.1% parameters in all linear layers contribute about 50% gradient norm square ∥∇F(w)∥ 2 2 . Fisher's diagonal is equal to the coordinate-wise gradient square vector when computing with downstream taskspecific loss, our sensitive parameters would cover a large fraction of the largest Hessian diagonal entries. This idea of sensitive parameters has been studied in the quantization community (Kim et al., 2023; Guo et al., 2023) and FO optimization (Sung et al., 2021) . However, we are the first one to leverage the extremely sparse sensitive parameters in LLM fine-tuning to accelerate ZO fine-tuning with LLMs. When we have perturbation and updating in the scale of billion parameters, finding which parameters to fine-tune would be important for improving ZO performance. Notice that here we use sensitive masks m k for understanding purposes. In Section 3.4, we will discuss how to transform Definition 3.2 to a parameter-efficient optimization pipeline by optimizing fixed sensitive parameters.We would investigate the theoretical convergence of sensitive sparse ZO-SGD on sensitive parameters under the non-convex optimization settings. Our assumptions are included in Appendix B.2.Theorem 3.3 (Convergence rate of sensitive sparse ZO-SGD (Definition 3.2)). If we pick η t = 1/(L(k + 2)),under Assumptions B.1 (bounded gradient error), B.2 (Lipschitz smoothness), and B.4 (sparse sensitive parameters), we would haveEQUATION)Moreover, if we still pick η t = 1/(L(k + 2)), with an extra Assumption B.3 (P.L. condition), we would haveEQUATIONThe proof for Inequality 5 is in Appendix B. We want to emphasize that our contributions are more on empirical LLM fine-tuning instead of general machine learning tasks, and in Section 4.1 we extensively compare our sparse ZO methods with other sparse ZO methods and we demonstrate its superiority during LLM fine-tuning. We do not use the strict "local r-effective rank" assumption that Malladi et al. (2023a) uses, and our Assumption B.4 can be easily observed empirically in Figure 2 . Liu et al. (2024a) and Ohta et al. (2020) also provide similar analysis on the convergence. However, they do not include our sensitive sparse mask in their studies.Sparse fine-tuning with fixed sensitive parameters. Our Theorem 3.3 focuses on dynamic sparse fine-tuning. However, Panigrahi et al. (2023) notice that in real LLM fine-tuning scenario, the fine-tuning performance could be attributed to a sparse subset of weights (∼ 0.01%). Malladi et al. (2023b) also find certain finetuning tasks would demonstrate kernel behaviors, which include "fixed (gradient) features": ∇ w f (w after FT ; (x, y)) ∼ ∇ w f (w before FT ; (x, y)).The similarity of gradient features during fine-tuning would imply that we do not need to re-select our sensitive parameters during fine-tuning i.e. select once before fine-tuning should be sufficient. This hypothesis can be validated by Figure 3 and Figure 5b . In Figure 3 , the fact that "task grad, static" does not vanish and still has a large ratio over "task grad, dyn." at the end of training demonstrate that we can select parameters before fine-tuning. We also include similar figures for Mistral-7B and OPT-6.7B in Figure 8 in Appendix C.3. We will describe Figure 5b in Section 4.3.Surrogate sensitive sparse mask from pre-training datasets. Another observation from Figure 3 is that the sensitive parameters derived from pre-training datasets (C4) would still cover a large fraction of model sensitivity. Therefore, we could use it as a surrogate sensitive sparse mask when gradients on downstream tasks are unavailable, particularly in scenario of on-device personalization.2 3.4. Our Proposal: ZO LLM Fine-Tuning with Fisher-Informed, Transferable SparsityThe sparse optimization on fixed parameters can be implemented as a parameter-efficient optimization workflow, which will reduce the perturbation and updating time during ZO optimization. Suppose we have derived a sensitive sparse mask m k , and we know it is fixed during fine-tuning.Instead of applying m k to z, we would apply it directly to w and extract the nonzero parts as below:EQUATIONDenote z k,t ∼ N (0 k , I k ) as the Gaussian perturbation sampled in timestep t. We will determine w sparse before finetuning and optimize on w sparse only and leave w dense frozen during fine-tuning. In this case, our sensitive sparse ZO-SGD update rule will become:EQUATIONIn Section 3.5, we will describe how this decomposition would seamlessly combine with existing post-training quantization (PTQ) methods, which creates an opportunity for on-device personalization. In Appendix C.6, we will discuss efficient implementations of linear layers after our decomposition.As LLMs are often pre-trained with user-agnostic public datasets, personalizing LLMs with individual user's preferences and meet user's specific needs before real-world deployment are vital. (Tan et al., 2024b; Mairittha et al., 2020) However, transferring the user-specific data to upstream cloud before fine-tuning LLMs would raise privacy concerns. (Xu et al., 2018) On the other hand, personal devices usually have less computational budget and are more memory-constrained than the cloud (Zhu et al., 2023) , and performing full fine-tuning would easily exceed the device memory budget.If we want to fine-tune a 7B model (e.g., Llama2-7B) on memory-constrained devices, we need to reduce the memory consumption on model weights, gradients, forward activations, and optimizer states:• Model weights. We would quantize the w dense to 4 bits, which reduces the model size of a Llama-2 7B model from 13.5 to 3.4 GiB.• Forward activations. ZO optimization already saves the need of caching activations.• Gradients. We would use the "random seed trick" same as MeZO (Malladi et al., 2023a) to reproduce layer-wise gradients instead of caching them.• Optimizer states. In this paper we use SGD. Our method can also be implemented as a parameter-efficient optimization method which would also work with other optimizers (even with Adam).As a result, our memory consumption is nearly minimum: we can fine-tune a Llama2-7B model under 8 GiB GPU memory without any offloading. This would satisfy the memory constraint by a wide range of edge or mobile devices as illustrated in Table 3 .Integration with quantization. In Section 3.4, we know that we can obtain surrogate sensitive sparse masks before fine-tuning. We would first decompose sensitive w to w sparse and w dense . We will then quantize w dense . During this process, we will use surrogate gradient information that many PTQ algorithms already have: they need gradients to calibrate their quantization errors.In addition, our method does not put strict constraints on specific choices of quantization algorithms since any algorithm that aims to minimize the least-square quantization error term (as follows) or its variant would suffice: (Chee et al., 2024; Nagel et al., 2020; Frantar et al., 2022; Lin et al., 2023; Kim et al., 2023 )Q(w) = argmin Q(w) E x ∥(w -Q(w))x∥ 2 2 (9)On-device personalization workflow. The workflow is il- For each line, the colors represent the fraction of parameters and the line style represents the category. "task grad, dyn." refers to the sensitive parameters selected at the given timestep (x-axis), and "task grad, static" refers to the sensitive parameters selected before fine-tuning. "C4 grad, static" refers to the sensitive parameters selected with gradients taken from causal language modeling on C4 datasets (Raffel et al., 2019) , and we keep it unchanged during fine-tuning. More similar figures are in Figure 8 .  
experiments: 
 In this section, we want to validate the effectiveness of our sensitive sparse ZO optimization method. We also investigate the effectiveness of our on-device personalization recipe in Figure 4 . There are a few research questions we want to answer:• RQ1: Is optimizing sensitive parameters more effective than optimizing other subset of parameters during ZO finetuning? Can we optimize surrogate sensitive sparse parameters when downstream gradient information is unavailable?• RQ2: Can optimizing extremely sparse and fixed parameters (Equation 8) lead to iteration-wise and total wall-clock time speedup?• RQ3: Can we match the full performance of ZO full finetuning by employing our on-device personalization recipe (Figure 4 )?We focus on 7B-level LLM models (Llama2-7B (Touvron et al., 2023) , Mistral-7B (Jiang et al., 2023) , OPT-6.7B We first investigate the performance of optimizing our sensitive parameters versus other subsets of parameters. Our baseline sparsity methods are random subsets and weight outliers. As illustrated in Figure 5a , we can find that ZO fine-tuning would benefit from sparse optimization, as all methods would achieve higher than ZO full fine-tuning at 90% sparsity. However, only sensitive parameters would maintain its performance as we move to the extreme sparsity region (> 99%). In fact, the performance curve of sensitive parameters w.r.t. different sparsity levels is near a flat curve, which indicates the performance loss by moving from 90% to 99.9% is minimal. Therefore, we can optimize 100 × less parameters compared with random and weight outliers and still get same performance. We also validate whether optimizing fixed and surrogate sensitive parameters should still yield satisfactory performance.In Figure 5b , we compare the performance of optimizing sensitive parameters with gradients on C4 dataset with its theoretical upper bound: fixed sensitive parameters derived from gradients on each fine-tuning task as the solid line and its dynamic version as the dash-dotted line. We also include the fixed and dynamic random subset parameters as a baseline. We can find that the gap of sensitive parameters between deriving from gradients on C4 dataset and gradients on each fine-tuning task at sparsity level 99.9% is small and blue line is still far above the random and full fine-tuning baseline. We also present a summary of our approaches with 99.9% sparsity on various datasets and models in Table 1 . By employing parameter-efficient ZO fine-tuning with extreme sparsity, we also achieve 1.2 -2.5× wall-clock time convergence speedup compared with ZO full fine-tuning as we nearly eliminate the ZO perturbation and optimizer update time, as Figure 6 shows. This also boosts the GPU utilization rate as large-batched ZO forward is often computebounded while the perturbation and optimization steps are often memory-bounded. Furthermore, the reduced memory footprint of parameter-efficient ZO fine-tuning allows for training larger models on the same hardware, potentially leading to even better performance. As a result, we answer this question that optimizing extremely sparse and fixed parameters leads to substantial iteration-wise and total wall-clock time improvements.We validate whether our sensitive sparse ZO optimization method would fit with on-device personalization pipeline described in Section 3.5 with Table 1 . We follow the exact recipe as described Figure 4 to report a number as "Sensitive (C4, static)", where we only optimize 0.1% sensitive parameters on top of a 4-bit quantized model. As ZO fine-tuning happens after model is quantized, ablating on extracting 0.1% random subsets of parameters would produce a different quantized model according to Figure 4 . So we choose to report the result for optimizing a fixed random subset on top of the 16-bit model as the "Random (static)" as a Table 1 : Performance of difference methods on Llama2-7B fine-tuning tasks. In the first column, "Q" means the full model is quantized with 4-bit quantization method (SqueezeLLM (Kim et al., 2023 )), and "ZO" means the model is finetuned with ZO-SGD optimizer. For each cell, we use the same hyperparameters and repeat it with 3 random seeds. We report the average and standard deviation of test set accuracy in the format of mean std . In last 2 columns, "Acc" means the average test set accuracy and "Rank" means the average rank of test set accuracy among all methods across datasets.(a) Llama2-7B theoretically performance upper bound.We also compare with optimizing with LoRA (Hu et al., 2021) and Prefix Tuning (Li & Liang, 2021) with ZO-SGD optimizer on top of the same quantized model as "Sensitive (C4, static)". The zero-shot inference and in-context learning (ICL) baselines are also included as the last two rows in each subtable. We follow the LoRA r and α and prefix length shown in Malladi et al. (2023a) , and for LoRA, we add it to all linear layers where our sensitive parameters are extracted. We find that integrating sensitive sparse ZO optimization with on-device personalization pipelines would still yield good performance exceeding all baselines across models and tasks. Particularly, the performance is higher than ICL, and ZO full fine-tuning in 16 bits. In addition, we have surpassed other common ZO-PEFT methods and random sparse ZO fine-tuning methods. This demonstrates the superiority of optimizing sensitive parameters only in ZO fine-tuning recipes. We also notice that optimizing sensitive parameters derived from C4 gradients still produce close re-sults as from task-specific gradients (in average less than 1% accuracy difference). This indicates optimizing surrogate sensitive parameters is still empirically successful. 
conclusion: 
 We have shown that the sensitive parameters provided by the pre-training process can effectively assist in ZO LLMs finetuning. Our experiments suggest that the ZO fine-tuning guided by 0.1% sensitive parameters in the LLM can even perform better than the full parameter ZO fine-tuning. The experiment results also demonstrate that the quantization of parameters other than sensitive parameters allows us to perform ZO fine-tuning of an LLM on limited memory devices. 
 
 Appendix: We present the notations used in this work as follows. training loss of w evaluated at a single data example (x, y) We start with listing standard assumptions in nonconvex optimization literature:F(w) = E (x,y) f (w; (x, y)) full-Assumption B.1 (Bounded stochastic gradient errors). For any data example (x, y) ∈ D and for any w ∈ R d , denote the full-batched loss function F(w) = E (x,y)∈D f (w; (x, y)), we haveEQUATIONEQUATIONAssumption B.3 (PL inequality). We assume that F(w) fulfills the Polyak-Lojasiewicz (PL) condition: there exists some µ > 0, for anyEQUATIONInspired by Figure 7 , we would assume the sensitive parameters of w are sparse.Assumption B.4 (Sensitive parameters are sparse). We assume at timestep t ∃m t ∈ {0, 1} d with the number of nonzero entries as k, ∃c ∈ [0, 1] such that∥m t ⊙ ∇ w f (w t ; (x t , y t ))∥ 2 = c∥∇ w f (w t ; (x t , y t ))∥ 2 .Here we assume c ≫ k/d.3 B.2. Proof for Equation 5, Theorem 3.3We will start with formulating the expectation of sensitive sparse ZO surrogate gradient norm square in terms of its corresponding stochastic gradient norm square.Lemma B.5 (Sensitive sparse ZO surrogate gradient norm square).E z[∥ĝ(w t , (x t , y t ), zt )∥ 2 ] = (2 + k)c∥∇ w f (w, (x t , y t ))∥ 2Proof for Lemma B.5. We know that our z can be considered as being sampled from N (0, Ĩd,m k ) where Ĩd,m k is the identity matrix I d with the main diagonal masked by m k .We expand the sensitive sparse ZO surrogate gradient covariance matrix E z ĝ(w, (x, y), z)ĝ(w, (x, y), y), z) ⊤ as follows:E z ĝ(w, (x, y), z)ĝ(w, (x, y), z) ⊤ = E zi [z i z⊤ i (m k ⊙ ∇ w f (w; (x, y)))(m k ⊙ ∇ w f (w; (x, y))) ⊤ zi z⊤ i ] = 2 (m k ⊙ ∇ w f (w; (x, y)))(m k ⊙ ∇ w f (w; (x, y))) ⊤ + ∥m k ⊙ ∇ w f (w; (x, y))∥ 2 Ĩd,m kThen the sensitive sparse ZO surrogate gradient norm square is the square of the diagonal of its corresponding covariance matrix:E z[∥ĝ(w t , x t , zt )∥ 2 ] = diag E z ĝ(w, (x, y), z)ĝ(w, (x, y), y), z) ⊤ 2 = 2c∥∇ w f (w, (x t , y t ))∥ 2 + kc∥∇ w f (w, (x t , y t ))∥ 2 = (2 + k)c∥∇ w f (w, (x t , y t ))∥ 2Then we are in good shape of deriving the convergence rate under the Lipschitz smoothness condition:Proof for Equation 5, Theorem 3.3. F(w t+1 ) ≤ E z,(x,y) {F(w t ) -cη t ∥∇ w F(w t )∥ 2 + cσ 2 η t + Lη 2 t 2 c(k + 2)∥∇ w F(w t )∥ 2 + Lη 2 t 2 c(k + 2)σ 2 } E z,(x,y) F(w t+1 ) ≤ E z,(x,y) {F(w t ) -cη t - Lη 2 t 2 c(k + 2) ∥∇ w F(w t )∥ 2 + cσ 2 η t + Lη 2 t 2 c(k + 2)σ 2 } Denote α = Lc(k + 2), we will have E z,(x,y) F(w t+1 ) ≤ E z,(x,y) {F(w t ) -η t c - α 2 η t ∥∇ w F(w t )∥ 2 + cσ 2 η t + α 2 σ 2 η 2 t } Set η t < c α = 1 L(k + 2), we haveE z,(x,y) F(w t+1 ) ≤ E z,(x,y) {F(w t ) - cη t 2 ∥∇F(w t )∥ 2 + cσ 2 η t + α 2 σ 2 η 2 t }If we apply our sparse ZO update rule recursively for T steps,1 T T -1 t=0 E z,(x,y) ∥∇ w F(w t )∥ 2 ≤ 2α T c 2 (F(w 0 ) -F * ) + 1 T T -1 t=0 cσ 2 η t + α 2 σ 2 η 2 t cη t 2 ≤ 2α T c 2 (F(w 0 ) -F * ) + (2σ 2 + σ 2 ) ≤ 2L(k + 2) c 1 T (F(w 0 ) -F * ) + 3σ 2 ≤ O k c • L T (F(w 0 ) -F * ) + 3σ 2 B.3. Proof for Equation 6, Theorem 3.3We can derive a convergence rate of sensitive sparse ZO-SGD optimization method under PL inequality and Lipschitzsmoothness as follows (this proof resumes from our prior proof with the Lipschitz-smoothness condition alone):Proof for Equation 6, Theorem 3.3. Denote κ as the condition number κ = µ L .E z,(x,y) F(w t+1 ) ≤ E z,(x,y) {F(w t ) - cη t 2 ∥∇F(w t )∥ 2 + cσ 2 η t + α 2 σ 2 η 2 t } ≤ E z,(x,y) {F(w t ) -cµη t (F(w t ) -F * ) + cσ 2 η t + α 2 σ 2 η 2 t } E z,(x,y) {F(w t+1 ) -F * } ≤ E z,(x,y) {(F(w t ) -F * ) -cµη t (F(w t ) -F * ) + cσ 2 η t + α 2 σ 2 η 2 t } E z,(x,y) {F(w t+1 ) -F * } ≤ E z,(x,y) {(F(w t ) -F * ) -cµη t (F(w t ) -F * ) + cσ 2 η t + α 2 σ 2 η 2 t }Plugging in η t ≤ c α and applying recursively for T iterations. i of linear layers for Llama2-7B (subfigure 7a), Mistral-7B (subfigure 7b), and OPT-6.7B (subfigure 7c) across RTE, WiC, and COPA tasks during FO-SGD fine-tuning. For each linear layer, we first sort parameters by the decreasing order of their gradient square value [∇F(w)] 2 i , i ∈ [d layer ], and we take the cumulative sum and normalize it to draw a blue curve, and the red-shaded region is the mean ± std of all blue curves.E z,(x,y) {F(w T ) -F * } ≤ (1 - cκ (k + 2) ) T (F(w 0 ) -F * ) + 3σ 2 c 2 2α ≤ (1 - cκ (k + 2) ) T (F(w 0 ) -F * ) + 3σ 2 c 2L(k + 2) ≤ 1 -O µ L • c k T (F(w 0 ) -F * ) + 3σ 2 c 2L(k + 2)In Figure 3 , we explore the transferability of gradient features from pre-training datasets (C4) to downstream tasks, and here we will also validate this phenomenon across models, as shown in Figure 8 . As there are no solid lines (top-(1e-2,1e-3,1e-4)) parameters with C4 gradient entries prior to fine-tuning) vanish to 0, we know the transferability of gradient features from C4 datasets to downstream datasets hold across models and downstream tasks. In this case, sensitive parameters determined from C4 gradients would still be similar to sensitive parameters determined from downstream task-specific gradients across models. Figure 8 : Cumulative normalized gradient square values of Llama2-7B (subfigure 8a), Mistral-7B (subfigure 8b), and OPT-6.7B (subfigure 8c)'s linear layers during FO fine-tuning. For a given model and training checkpoint, we report the average value across all linear layers as a line in each subfigure. For each line, the colors represent the fraction of parameters (1e-2,1e-3,1e-4) and the line style represents the category. "task grad, dyn." refers to the sensitive parameters selected at the given timestep (x-axis), and "task grad, static" refers to the sensitive parameters selected before fine-tuning. "C4 grad, static" refers to the sensitive parameters selected with gradients taken from causal language modeling on C4 datasets, and we keep it unchanged during fine-tuning.For all experiments, we use 20,000 training steps with ZO-SGD optimizer (Definition 2.2). We will save a model checkpoint every 500 steps, and load the checkpoint with the lowest loss on the validation set at the end of the training, and report its test set accuracy as result. Usually, the training/validation set will be sampled from the original dataset with size 1000/500 respectively and the test set is of size min(1000, |original test set|), except for CB and COPA that we use 100 for the validation set size. For all ZO experiments (Table 4 and Table 5 ), we use batch size of 16. This experiment setting is identical to Malladi et al. (2023a) . Our hyperparameters (learning rate η, perturbation scaling constant ϵ, and the number of ICL examples) for Table 1 is reported in Table 4 for reproducibility. We use constant η and ϵ throughout our experiments. We also report the chosen hyperparameter for Figure 5a and Figure 5b in Table 5 . For LoRA, we always add to all linear layers with r = 8 and α = 16, and for Prefix Tuning, we always add to W K and W V with length as 5, as what Malladi et al. (2023a) uses.In light of this observation, in our Figure 1 , we implement both "SparseAdd" and "SparseAddMM" methods for "Sensitive (0.1%)" and "Random (10%)". For each method we report the lowest time out of these 2 implementations: for "Sensitive (0.1%)" training and "Random (10%)" training and inference, we use "SparseAdd" approach. For "Sensitive (0.1%)" inference, we use the "SparseAddMM" approach.Integration with weight quantization method. Weight quantization algorithms can be categorized into 2 categories: uniform quantization method and non-uniform quantization method. For uniform quantization method, (Xi et al., 2023) indicates that we could use integer matrix multiplication to compute Q(w dense )x efficiently without first dequantizing Q(w dense ) to 16 bits. However, this creates difficulty on our "SparseAdd" approach as we will violate the constraint of uniformly-spaced quantization bins by computing SparseAdd(Q(w dense ) + w sparse ). In this case, we also have 3 different implementations:EQUATION= SparseAddMM Dequantize IntMM(Q(w dense ), x) , w sparse , x fits with integer matmul (19)= SparseAddMM Dequantize(Q(w dense )), x, w sparse similar to Equation 15(20)= (Dequantize(Q(w dense ))+w sparse )x (21) = DenseMM(SparseAdd (w sparse , Dequantize(Q(w dense )), x) similar to Equation 17(22)Equation 19 would compute IntMM(Q(w dense ), x) before dequantizing it to 16 bits. This would make "SparseAdd" approach infeasible and we can only employ "SparseAddMM" approach in this case. Notice that both Equation 20 and Equation 22 would still dequantize Q(w dense ) first and the choice of implementation would follow into our discussion of input vector size x in last paragraph. We leave a practical implementation and thorough benchmarking into a future work.We recommend using Equation 19 when we use efficient integer matmul to compute Q(w dense )x and in other cases, using Equation 20 or Equation 22 follows our previous recommendation based on the size of input vectors.C.7. Hardware, platform, libraries, and other details for fine-tuning and benchmarking In subfigure 2 and 4 in Figure 9 , we use Nvidia A100-SXM4 (40 GiB) and AMD EPYC 7543P 32-Core CPU with PyTorch version 2.1, HuggingFace version 4.38.2, and CUDA 12.2. We use Flash Attention 2 (Dao, 2023) in HuggingFace Transformers library throughout our experiments, and the base model for ZO full fine-tuning and benchmarking is always Llama2-7B with Float16 datatype (torch.float16). We also use the Float16 datatype (torch.float16) for all of our sparse parameters (sensitive sparse, random subsets, etc.) in ZO fine-tuning experiments. Notice that for all of the FO fine-tuning demonstrations (Figure 7 and Figure 8 ) we use the BrainFloat16 datatype (torch.bfloat16) to avoid the NaN issue from the Float16 datatype.In Figure 1 and Figure 9 , we use sequence length of 512 and batch size 16 sampled from WikiText-2 dataset (Merity et al., 2016) as a representative computational intensity for ZO training, and for inference we generate 128 tokens with top-p (p = 0.9) sampling from the prompt "Please describe the effect of sparse zeroth-order optimization methods on memory-efficient LLM fine-tuning: ". We still use the Float16 datatype (torch.float16) for both benchmarks.Here we assume data examples (x, y) ∼ pD in fine-tuning datasets after verbalization would also appear in the large text corpus during pre-training.Obtaining gradients of LLMs on edge devices is expensive, and we usually cannot transfer data from edge devices to the cloud to compute the gradient on downstream tasks on cloud. In this case we would need some surrogate gradient information to derive sensitive sparse masks on cloud. We will discuss this in Section3.5.From Figure7, we know that for c ∼ 0.5, we only need k/d ∼ 0.001. In this case k/c ∼ 0.002d. 


Decision: The reviewers are positive about the paper and its merits. I agree with their assessment and recommend acceptance.
I would also suggest and encourage the following:
1. It would be very helpful to explain the GPU memory demand for ZO fine-tuning in detail. (As reviewers suggested, to improve memory management in future work) 
2. For 4-bit quantization on edge devices, since some of them support INT8, it would be worth exploring how much the proposed methods could benefit from INT8 in future work.

The author should clarify the meaning of "with Float16 datatype" and "16 bit". As (for throughput) official Flash Attention 2 does not support fp16, and (for fine-tuning) not all (edge) devices support bf16.