Sparse Gradient Compression for Fine-Tuning Large Language Models

We sincerely appreciate the time and effort you have devoted to reviewing our work and offering such valuable feedback. Below, we have addressed each of your points in detail. Note that we have split our response into two parts due to the word limit.

Weaknesses

Dataset Limitation: The authors only use a single dataset (Commonsense) in the experimental sections. I strongly recommend adding at least one more dataset to demonstrate the generalizability of the algorithm across different data domains.

We have added results from a second dataset, alpaca_gpt4, fine-tuned on LLaMA2-7B, and evaluated on the MMLU benchmark. Our results show that under similar memory consumption, our approach is competitive with LoRA and GaLore.

Comparison to LoRA in Speed: While SGC effectively reduces optimizer memory costs similar to LoRA, LoRA offers additional advantages by significantly speeding up the fine-tuning process. Through low-rank adapters and fewer trainable parameters, LoRA can accelerate fine-tuning by 12–16 times compared to full fine-tuning. Since SGC performs full forward and backward propagation, it does not offer the same speed benefit and is likely to be significantly slower than LoRA. I suggest the authors include a training time profile for SGC to clarify this difference.

The forward propagation speed is the same for all approaches since the same calculations are made (assuming merged LoRA weights, otherwise there is a marginal increase in computation time for LoRA). We would like to emphasize that SGC does not perform full backward propagation, and only updates the target weight matrices as with LoRA and GaLore. SGC does require computing full gradients which is slower than LoRA, but GaLore suffers from the same issue (as seen in table below). The remaining gap in runtime between CESGC and LoRA can be attributed to the OMP calculation, which we will discuss in more detail in your question at the end.

Comparison to LoRA in Activation Memory: LoRA also has the advantage of substantially reducing activation memory costs. For example, in LLaMA-2-7B full fine-tuning, the activation memory cost for a batch size of 128 can reach up to 40GB. With low-rank adapters, this can be reduced to under 1GB. Since SGC performs full forward and backward propagation, it does not reduce activation memory cost and is expected to be comparable to full fine-tuning. I recommend the authors discuss SGC's activation memory usage in more detail.

We would like to clarify that the activation memory usage of LoRA and SGC is identical when applied to the same target layers (e.g., the Q and V weight matrices). If activation checkpointing is not used, the activations for these layers are stored during forward propagation regardless of the method employed. During backpropagation, the computation of gradients with respect to the Q and V matrices requires the same set of activations for both LoRA and SGC. Therefore, the activation memory cost of SGC is not comparable to full fine-tuning, as it is restricted to the activations of the target layers, just like in LoRA.

Base Model Limitations: Although the authors utilize LLaMA-2-7B, LLaMA-2-13B, and LLaMA-3-8B, I recommend including more state-of-the-art models, such as LLaMA-3.1. Additionally, incorporating models outside the LLaMA family, like Phi-2/Phi-3 or Mistral, could further demonstrate SGC's generalizability across different model architectures.

To extend our experimental results, we fine-tuned LLaMA-3.1-8B on the commonsense dataset. In addition, we also fine-tuned Mistral-7B on a subset of the alpaca gpt4 dataset evaluated on MMLU. With these results, we show that SGC does generalize across different model architectures.

LLaMA3-8.1B

Mistral-7B

We sincerely appreciate the time and effort you have devoted to reviewing our work and offering such valuable feedback. Below, we have addressed each of your points in detail.

Weaknesses

The author highlights limitations in the flexibility and granularity of LoRA due to the dependency on model dimensions. However, in practical applications, these constraints may not significantly impact performance. Many real-world tasks do not require extreme reductions in trainable parameters, and the existing flexibility of LoRA is often sufficient. For instance, as shown in Table 2, LoRA fine-tunes only 0.2% of the parameters, meaning the LoRA weights and optimizer states are not the bottleneck—the base model weights and activations are. Reducing this further to 0.08% would likely not yield significant benefits.

The importance of flexibility and granularity becomes apparent in limited-data fine-tuning. SGC’s flexibility allows it to adapt to diverse resource constraints such as extreme memory limitations or specific downstream task performance requirements. Here we present a scenario that offers practical performance and efficiency benefits as a result of the enhanced flexibility of SGC. In this example, we fine-tune LLaMA2-7B on a 2k subset of the BoolQ dataset and show how under extreme memory requirements, SGC outperforms both GaLore and LoRA. With LoRA and GaLore, r=1 is the minimum number of optimizer parameters (equivalent to total sparisty = 512 in table below). However, for our approach, we can use a smaller number of optimizer parameters by varying the total sparsity level. By going below r=1, we can select values for total sparsity such that SGC outperforms LoRA and GaLore while using a significantly smaller number of optimizer states.

Questions

Can you justify in what scenarios is fine-grained control over the training parameters necessary?

Please see the answer to the weakness above.

If fine-grained control of training parameters is required, are there simpler methods to achieve similar results, such as using different ranks in different transformer layers?

Alternative methods such as layer-wise rank adjustments are tightly coupled to either the model architecture or downstream task, requiring careful manual tuning of which layers should receive higher or lower ranks. In contrast, SGC provides model-agnostic flexibility by decoupling optimizer state size from parameter dimensions, enabling granular memory-performance tradeoffs without requiring architectural modifications. SGC’s low-dimensional projection operates uniformly across all layers, simplifying implementation and making it more scalable for large-scale models with many layers.

We sincerely appreciate the time and effort you have devoted to reviewing our work and offering such valuable feedback. Below, we have addressed each of your points in detail. Note that we have split our response into two parts due to the word limit.

Weaknesses

The novelty of the proposed approach is limited. The concept of projecting optimizer states into a subspace with a dimension independent of the original model size has been previously discussed in the top-k compressor as shown in [1] and [2].

Thank you for bringing these papers to our attention and we will be sure to cite them as part of the related work section. Papers [1] and [2] are focused on the distributed setting where communication between workers is the bottleneck and there is an emphasis on error compensation to mitigate the lossy compression. However, the idea of using top-k sparsification along with dimensionality reduction has not been readily explored in the context of fine-tuning LLMs. The gradient compressors in these papers are functions from $\mathbb{R}^d$ to $\mathbb{R}^d$, and the dimensionality of the optimizer states remains unchanged. On the other hand, in our approach, we use a random matrix to project the gradients to a lower dimensional subspace in $\mathbb{R}^k$ (with an arbitrary size $k << d$) while doing sparsification to enable memory efficient storage of optimizer states.

The paper lacks a theoretical analysis of the relationship between the choice of k and the model’s convergence, a detail that has been explored in [1] and [2].

Following [1], it is possible to show that top-$k$ sparsification leads to convergence at the same rate as vanilla SGD. The key difference in our algorithm is the use of chunking and sparsification applied to every chunk. Thus, the proof of convergence boils down to bounding the distance between the sparse form of gradient vector $G$ and the sparse form of every sub-vector after chunking the gradient vector $G$. Let us denote $G=[G_1, \dots, G_c]$, where $G_i \in \mathbb{R}^{d/c}$ is the subvector corresponding to each chunk. We then represent the $s$-sparse form of $G$ with $G’$, i.e.,

$$G' = \text{Sparsify}_s(G) = \text{Sparsify}_s([G_1, \dots, G_c]),$$

and the vector of sparsified chunks with

$$\tilde{G} = [\tilde{G}_1 , \dots, \tilde{G}_c], \quad \tilde{G}_i = \text{Sparsify}s_c(G_i).$$

In the following, we will provide a bound for $E$ || G' - \tilde{G} ||_2^2 $$ where the randomness is from the gradient $G$.

If we assume a uniform distribution for non-zero indices in $G'$, it is easy to verify that

$$E \left[ \|| G' - \tilde{G} \||_2^2 \right] = 0.$$

On the other hand, it is straightforward to show that the worst case for $E$ || G' - \tilde{G} ||_2^2 $$ occurs when all the non-zero entries are located contiguously in $G'$. This is due to the way we perform the chunking, i.e., having consecutive indices in a chunk. Without loss of generality and for ease of presentation, we assume that entries $1$ to $s$ in $G'$ are non-zero.

Let us denote the number of chunks these $s$ entries span is equal to

$$l = \lceil s / (d / c) \rceil.$$

Therefore, the error can be computed in two parts: $D_1$, due to the contribution from chunks $i \leq l$, where there are missing entries not selected in $\tilde{G}$, and secondly, $D_2$, in chunks $i > l$, where extra entries are selected in $\tilde{G}$. Adding these two together, we get:

$E\left[ \|| G' - \tilde{G} \||^2_2 \right] = D_1 + D_2, \text{where}$

$D_1 = E \left[ \sum_{i=1}^l || G_i - \tilde{G}_i ||^2_2 \right] \le (s - l s_c) E \left[ \frac{\|| G' ||^2_2}{s} \right],$

$D_2 = E \left[ \sum_{i=l+1}^c || G_i - \tilde{G}_i ||^2_2 \right] \le (c - l)s_c E\left[ \frac{\|| G' ||^2_2}{s} \right],$

Putting everything together and simplifying, we get a worst-case bound:

$$E\left[ \|| G' - \tilde{G} \||_2^2 \right] \leq 2 \left( 1 - \frac{s}{d} \right) E\left[ \|| G' \||_2^2 \right] \le 2 \left ( 1 - \frac{s}{d} \right) G_m,$$

where $G_m$ is the bound over the expectation of the l-2 norm of the gradients.

Using this bound, we can then formulate the theoretical analysis for convergence, and we leave this as future work.

We sincerely appreciate the time and effort you have devoted to reviewing our work and offering such valuable feedback. Below, we have addressed each of your points in detail.

Weaknesses

Although SGC is more flexible, this advantage is somewhat marginal, as LoRA and GaLore are already quite flexible.

SGC’s flexibility allows it to adapt to diverse resource constraints such as extreme memory limitations or specific downstream task performance requirements. Here we present a scenario that offers practical performance and efficiency benefits as a result of the enhanced flexibility of SGC. In this example, we fine-tune LLaMA2-7B on a 2k subset of the BoolQ dataset and show how under extreme memory requirements, SGC outperforms both GaLore and LoRA. With LoRA and GaLore, r=1 is the minimum number of optimizer parameters (equivalent to total sparisty = 512 in table below). However, SGC can use a smaller number of optimizer parameters by varying the total sparsity level. By going below r=1, we can select values for total sparsity such that SGC outperforms LoRA and GaLore while using a significantly smaller number of optimizer states.

The paper lacks throughput experiments and runtime analysis of OMP.

We conducted throughput experiment for CESGC fine-tuning LLaMA2-7B on commonsense dataset = 627 tokens/s.

The additional time complexity of our approach mainly comes from the OMP calculation which requires an extra O(k*s) time complexity at each timestep with the efficient batched GPU implementation that we developed. Time profiling results are shown below where we performed a runtime comparison when fine-tuning LLaMA2-7B on the commonsense 170k dataset for one epoch, on a single H100 GPU.

It does not include empirical experiments comparing memory usage of SGC and baseline methods to validate the theoretical analysis.

Empirical experiments for fine-tuning LLaMA2-7b on a subset of the commonsense dataset shows approximately equal max GPU memory usage during the training process. Here memory usage is almost equivalent between LoRA, GaLore, and SGC because parameters $k$ (SGC) and $r$ (LoRA, GaLore) were deliberately selected to be equal. Even in this setting, we can see that SGC is more memory efficient than LoRA, and this gap would increase if we increase the model size and/or the targets to apply SGC to (Only Q, V weight matrices were targeted here).

There is no information on the error magnitude after gradient compression.

We conducted additional experiments to obtain the error magnitude from before gradient compression, and after OMP recovery. We used the average MSE across all targeted weights to compute the difference between MESGC and a sparsified gradient approach without compression.

Mischaracterization in lines 350-352 and 368-369, where it mentions GaLore as a type of PEFT method; GaLore is not a PEFT method.

Thank you for pointing out this important clarification regarding the characterization of GaLore. We agree that GaLore is primarily a gradient compression method and not a Parameter-Efficient Fine-Tuning (PEFT) method. We will make sure to fix this in the final version of the paper.

Proposed revisions: Lines 350-352: For this purpose, we leverage the compatibility of SGC with various optimization techniques. Specifically, we utilize GaLore, a gradient compression method, to obtain $\mathbf{B}_t$​, which reduces the dimensionality of the vector input to the SGC algorithm.

Lines 368-369: Here, we analyze the memory requirements of our efficient SGC implementations and compare them with established methods for reducing memory usage during fine-tuning, specifically gradient compression techniques like GaLore and PEFT methods like LoRA.

Questions

What is the practical benefit of increased flexibility and granular control, given that memory usage is primarily dominated by parameters rather than optimizer states in methods like PEFT or GaLore?

Please see the answer to the first weakness above.

We sincerely appreciate the time and effort you have devoted to reviewing our work and offering such valuable feedback. Below, we have addressed each of your points in detail. Note that we have split our response into two parts due to the word limit.

Weaknesses

Limited Applicability: While the paper claims that SGC offers a more flexible, fine-grained tradeoff, PEFT methods typically target compute-constrained scenarios, where such granular control may require extra tuning that reduces practicality. It would be beneficial to include a plot with sparsity on the x-axis and performance on the y-axis to directly compare the flexibility of SGC with LoRA. This visualization could more intuitively demonstrate whether SGC’s fine-grained control offers practical performance benefits at different sparsity levels.

Thank you for your observation regarding the tradeoff between flexibility and practicality in compute-constrained scenarios. Below, we clarify the advantages of this tradeoff and address concerns about practicality.

While PEFT methods focus on compute-constrained scenarios, SGC’s flexibility allows it to adapt to diverse resource constraints (such as extreme memory limitations or specific downstream task performance requirements). The additional hyperparameters introduced by SGC can be optimized through automated techniques such as grid search mitigating the overhead of manual tuning. Here we present a scenario that offers practical performance benefits at different sparsity levels. Since we are unable to include a plot in this response, we present the results in a table format. In this example, we fine-tune LLaMA2-7B on a 2k subset of the BoolQ dataset and show how under extreme memory requirements, SGC outperforms both GaLore and LoRA. With LoRA and GaLore, r=1 is the minimum number of optimizer parameters (equivalent to total sparisty = 512 in table below). However, for our approach, we can use a smaller number of optimizer parameters by varying the total sparsity level. By going below r=1, we can select values for total sparsity such that SGC outperforms LoRA and GaLore while using a significantly smaller number of optimizer states.

Questionable Memory Advantage: The memory usage for first-order optimization methods largely comes from the model parameters, gradients, activations, and optimizer states. Even with Adam’s two states, optimizer memory costs are typically less than half. SGC, based on Adam, can’t reduce memory below that of simple SGD without momentum, and since it still calculates full gradients, its GPU memory consumption may surpass LoRA, which doesn’t require full gradient computations.

Thank you for raising this point. Instead of storing all gradients in memory, we use per-layer weight updates [1], which compute the gradients only for the current layer being processed. This approach significantly reduces memory requirements by avoiding the need to retain gradients for the entire model at once. By using a memory profiler, we can show that our approach does not require more GPU memory than LoRA given similar optimizer state sizes (see table below).

Subpar Performance: As seen in Table 2, SGC shows no clear performance advantage over methods like LoRA and GaLore, raising questions about its efficacy as a fine-tuning method.

We acknowledge that Table 2 does not show a consistent performance advantage for SGC across all metrics, and we appreciate the opportunity to clarify and expand on its contributions. Firstly, Table 2 demonstrates that SGC achieves comparable accuracy to LoRA and GaLore while using significantly fewer optimizer state parameters, showing that it is still an effective method for fine-tuning. The performance of SGC also becomes evident in memory-limited or data-limited settings, as demonstrated in Sections 5.2 and 5.3. In small datasets, CESGC consistently outperforms LoRA and GaLore (e.g., on BoolQ in Figure 3). In addition, MESGC also achieves higher accuracy while being more memory efficient in commonsense reasoning tasks (see Table 3).

[1] Lv, Kai, et al. "Adalomo: Low-memory optimization with adaptive learning rate." arXiv preprint arXiv:2310.10195 (2023).