Q-GaLore: Quantized GaLore with INT4 Projection and Layer-Adaptive Low-Rank Gradients

Thank you for the insightful questions and constructive suggestions, which have been immensely helpful in improving the quality of our work. Below, we provide detailed responses:

[Q1: Clarification of memory reduction] Thanks for the question. We’d like to clarify that the memory overhead of performing SVD on large matrices is minimal compared to the memory required for model weights and optimization states. This is because SVD is conducted in a layerwise manner. For example, consider one “mlp.up” layer of LLaMA-2-7B, which has a shape of 4096×11008. In 32-bit precision, this layer consumes approximately 0.17 GB of memory. During the SVD operation, the resulting matrices U, S, V have shapes 4096×4096, 4096, and 11008×11008, respectively, consuming a total of approximately 0.51 GB of memory. In comparison, the entire model weights require around 14 GB of memory. After completing the SVD for each layer, only the projection matrix is retained, and any intermediate memory usage is released. This approach ensures that the memory cost of the SVD operation is negligible. The memory savings in Q-GaLore come from reductions in weight and optimization states compared to GaLore, as illustrated in Figure 5.

[Q2: Explanation about lazy SVD update] Great suggestion. We’ve updated more results in Section D of the updated draft. As illustrated in Figure 8, the gradient subspace shows the most drastic changes in the q and k projections (with the smallest cosine similarity). Intuitively, these projections are responsible for generating the attention patterns and are more influenced by the token relationships within each training sample. Consequently, their gradient information is more diverse, leading to a more diverse gradient subspace. When comparing different layers, the middle layers tend to exhibit a more consistent gradient subspace compared to the initial and final layers. This behavior may be related to the oversmoothing issue in transformers[1-2], where the middle layers are not well-optimized, resulting in oversmoothed token representations.

[Q3: Comparison of the training time] Thanks for the question. We measure the latency of training a 1B model on two A6000 GPUs where we use the largest possible batch sizes for each method. For GaLore, a batch size of 32 can be used without OOM, while Q-GaLore can use a batch size of 64. In this case, the average throughput for Q-GaLore is 36.87 samples/s, higher than GaLore’s 35.56 samples/s. Additionally, we’d like to clarify the benefits of memory-efficient training across three key scenarios: (i) Single GPU Training: Reducing memory requirements for 7B models makes training more accessible on consumer-grade GPUs. For example, an RTX 4060 Ti with 16GB memory, costs approximately 499 dollars while RTX 4090 with 24GB memory requires around $1599. (ii) Distributed Training with Data Parallelism: When using the same number of GPUs, reducing memory overhead allows for training with larger batch sizes. This can reduce communication costs between GPUs, improving overall efficiency. (iii) Model Parallelism: For very large models that exceed the memory capacity of a single GPU for even one batch, model parallelism becomes necessary. By lowering memory overhead, it is possible to avoid the complexity of model parallelism, enabling training on larger models without additional overhead, and scaling up training larger models more effectively.

[1] PanGu-π: Enhancing Language Model Architectures via Nonlinearity Compensation.

[2] FFN-SkipLLM: A Hidden Gem for Autoregressive Decoding with Adaptive Feed Forward Skipping.

Thank you for clarification. I've adjusted my rating, and here are my follow-up questions.

• I am still a bit concerned with the training perplexity - memory & runtime tradeoff. It will be clearer if Tab.1 includes runtime results like training time or throughput like #batches/sec.
• Is the cosine similarity threshold model-dependent? Q-Galore will be a practical method if the recommended value of 0.4 works for other model architectures and sizes.

We sincerely appreciate Reviewer CuTC for the detailed and constructive suggestions. Below, we provide pointwise responses.

[Q1: Limited novelty with increased complexity.] We respectfully disagree that our novelty is modest: Firstly, training LLMs with weight quantization and low-rank gradients is non-trivial. The main challenge is the substantial reduction in gradient information. This introduces additional complexity in identifying optimal low-rank subspaces. For example, directly quantizing GaLore’s weights to 8-bit results in a significant performance degradation, with perplexity increasing by as much as 7.86 when pre-training a LLaMA-60M model. To address this, we investigated stochastic rounding and for the first time demonstrated the possibility of training large-scale quantized LLMs with low-rank gradients.

Then, Our investigation into the adaptive convergence properties of gradient subspaces across layers revealed novel insights. These findings not only provide an effective solution for reducing the computational overhead of expensive SVD operations but also deepen our understanding of transformer training dynamics. Specifically, we observed: (i) some layers’ gradient subspaces converge early in training. (ii) others exhibit stable subspaces within specific windows. (iii) certain layers only converge later in the training process. These novel observations connect to prior research on the diverse low-rank properties of weight spaces [1,2]. For instance, a stable gradient subspace over an extended training window may cause weights to settle into the same low-rank space, offering valuable insights for future studies on transformer training dynamics.

Additionally, we conducted an ablation study on the threshold for update intervals, as presented in Figure 7. The x-axis in the figure represents the relative SVD count. In this study, we performed a grid search on the cosine similarity threshold within the range [0, 1], reporting the corresponding SVD counts and perplexity values. The results indicate that a threshold of 0.4 (marked as the red star in Figure 7) is a turning point. Below this threshold, performance improves significantly as the threshold decreases. This experiment demonstrates the extra hyperparameter of cosine threshold can be set as 0.4 without performance degradation.

[1] JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention.

[2] The Truth is in There: Improving Reasoning in Language Models with Layer-Selective Rank Reduction.

[Q2: Latency breakdown of different operations.] Thanks for the question. We compare the end-to-end latency on two A6000 GPUs using distributed training. One key benefit of memory-efficient training is that by reducing memory requirements, we can use larger batch sizes during training, which enhances throughput since training is generally memory-bound.

In our experiments, we trained a 1B model using the largest possible batch sizes without causing an out-of-memory (OOM) error. GaLore supports a batch size of 32, while Q-GaLore enables a batch size of 64. The average throughput for Q-GaLore is 36.87 samples/s, higher than GaLore’s 35.56 samples/s. Under the same batch size, the additional computation required for quantization and dequantization in Q-GaLore results in a throughput of 26.60 samples/s. These additional computation costs, however, are compensated for by the ability to use larger batch sizes.

For cosine similarity comparisons, these operations are executed only during projection updates, which occur every 200 iterations. The latency caused by this operation is negligible compared to the SVD computation where its latency is smaller than the time fluctuation observed during different projection update stages.

Additionally, by lowering memory overhead, Q-GaLore makes training more accessible. For example, an RTX 4060 Ti with 16GB memory (priced at approximately 499 dollars) becomes sufficient, compared to an RTX 4090 with 24GB memory, which costs around $1599. Furthermore, for extremely large models that cannot fit a single batch on a single GPU, model parallelism is typically required. However, by reducing memory requirements, Q-GaLore can potentially avoid the complexity of model parallelism, enabling more efficient scaling for training larger models.

[Q3: Experiment details and threshold for cosine similarity.] Thanks for pointing it out. We’ve provided all the training and fine-tuning hyperparameters in Section C of the appendix in the updated draft. Additionally, we’ve conducted an ablation study on the threshold for cosine similarity. The results are presented in Figure 7, from which a threshold of 0.4 (marked as the red star in Figure 7) is a turning point. Below this threshold, performance improves significantly as the threshold decreases. Thus, we recommend using 0.4 as the threshold for cosine similarity.

We thank Reviewer 1Bwe for acknowledging our method “achieves substantial memory reduction” as well as the lazy SVD update strategy is “well-motivated”, To further address Reviewer 1Bwe’s concerns, we provide pointwise responses in the following:

[Q1: Motivation in the fine-tuning scenario] Thanks for the suggestion, we’ve strengthen the motivation in Section 1 in the updated draft, as “Even without factoring in any considerations for product efficiency, fine-tuning a LLaMA 7B model with 16-bit precision necessitates at least 56 GB memory for maintaining the model weight, Adam optimizer states and weight gradient, which is prohibitively expensive.”

[Q2: Incremental modification about quantization section] Thanks. We respectfully disagree that the quantization aspect in Q-GaLore is incremental compared to GaLore. In the original GaLore, quantization was limited to the use of an 8-bit Adam optimizer, which is a plug-and-play quantized optimizer from [1]. In contrast, Q-GaLore incorporates quantization from two distinct perspectives:

Weight Quantization: different from post-training quantization, one primary challenge of weight quantization during training lies in the significant reduction of gradient information. During each iteration, small gradient information is often lost due to the rounding operation in quantization, leading to ineffective training. We address this by stochastic rounding, which effectively preserves gradient information. And we demonstrate the efficacy of this approach across a wide range of experiment scales.

Projection Quantization: We show that the projection matrix is particularly friendly to quantization, offering a dual benefit. First, it provides an effective solution for further reducing memory costs. Second, it reveals insights into transformer training dynamics. Specifically, we observe that the gradient subspaces of certain layers exhibit an “early-bird” property, converging quickly during the initial stages of training. This empirical analysis provides valuable insights into the training dynamics of transformers.

Given these, we believe the quantization components in Q-GaLore represent a significant advancement rather than incremental improvements.

[Q3: Ablation study of the quantization of weight and projection matrices]: Great question. We conducted additional experiments to validate the individual effect of weight and projection quantization. As shown in Table R1, the final performance is slightly more sensitive to the weight quantization than projection matrix quantization.

Table R1: Comparison results of different precision of weight and projection quantization. Results are reported in perplexity. Where “W” represents weight and “P” represents the projection matrices.

[1] 8-bit Optimizers via Block-wise Quantization

We sincerely thank Reviewer PunU for the helpful comments and positive initial feedback on our work. Below, we provide detailed responses.

[Q1: Performance gap] Thanks for pointing it out. In Table 1, the perplexity difference for the 1B model is slightly smaller than that of the 350M model but slightly larger than the 130M model, indicating no consistent trend of an increasing gap as the model size scales up. Furthermore, compared to the baselines of Low-Rank, LoRA, and ReLoRA, the perplexity gap with GaLore is significantly larger than that of Q-GaLore. For instance, in the case of the Llama-1B model, the perplexity differences are 126.89, 3.57, and 2.69, respectively, which are considerably higher compared to the 0.61 of Q-GaLore. This suggests that Q-GaLore maintains comparable pre-training performance, as mentioned in Line 328.

[Q2: Threshold for cosine similarity] Thanks for pointing it out. We’ve conducted an ablation study on the threshold for cosine similarity, as shown in Figure 7, where the x-axis represents the relative SVD count. In this study, we performed a grid search for the cosine similarity threshold within the range [0, 1] and reported the corresponding SVD counts along with the perplexity. The results indicate that setting the threshold to 0.4 (marked as the red star in Figure 7) serves as a turning point, where performance begins to increase significantly when the threshold is below 0.4. And we have included the details of this experiment in the updated draft to avoid potential misunderstanding.

Thanks for your response. I still positively evaluate this work, but do not think I can further increase the score

[Q1. Performance Gap] I believe there are two issues here.

1. Noting that the performance degradation starts at 350M, I believe it's possible that the effectiveness of Q-GaLore decreases with larger models.

2. While the author noted that the absolute perplexity difference decreased at 1B (compared to 350M), I want to point out that the loss scale generally decreases as the model size increases following the scaling laws. I am not sure about the best way to measure the performance degradation across different model scales, but my concern 1 still remains valid.

[Q2. Threshold for Cosine Similarity] Thanks for providing additional details!