Q-GaLore: Quantized GaLore with INT4 Projection and Layer-Adaptive Low-Rank Gradients (2407.08296v1)

Abstract: Training LLMs is memory-intensive due to the large number of parameters and associated optimization states. GaLore, a recent method, reduces memory usage by projecting weight gradients into a low-rank subspace without compromising performance. However, GaLore relies on time-consuming Singular Value Decomposition (SVD) operations to identify the subspace, and the frequent subspace updates lead to significant training time overhead. Moreover, GaLore offers minimal improvements in accuracy and efficiency compared to LoRA in more accessible fine-tuning scenarios. To address these limitations, we introduce Q-Galore, a novel approach that substantially reduces memory usage by combining quantization and low-rank projection, surpassing the benefits of GaLore. Our method is based on two key observations: (i) the gradient subspace exhibits diverse properties, with some layers converging early in training while others are subject to frequent changes; (ii) the projection matrices are highly resilient to low-bit quantization. Leveraging these insights, Q-GaLore adaptively updates the gradient subspace based on its convergence statistics, achieving comparable performance while significantly reducing the number of SVD operations. We maintain the projection matrices in INT4 format and weights in INT8 format, incorporating stochastic rounding to capture accumulated gradient information. This approach enables a high-precision training trajectory using only low-precision weights. We demonstrate that Q-GaLore achieves highly competitive performance with exceptional memory efficiency. At pre-training, Q-GaLore facilitates training a LLaMA-7B model from scratch on a single NVIDIA RTX 4060 Ti with only 16 GB memory. At fine-tuning, it reduces memory consumption by up to 50% compared to LoRA and GaLore, while consistently outperforming QLoRA at the same memory cost.

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

Introduction

The increasing prominence and efficacy of LLMs have concurrently spotlighted the intensive memory requirements necessary to train such models effectively. LLMs, with billions of parameters, pose significant resource challenges, often necessitating extensive computing infrastructure. Traditional training methods are substantially memory-intensive, with significant amounts allocated for trainable parameters, optimizer states, and gradients. GaLore, a recent memory-optimization technique, introduced low-rank gradient representations using Singular Value Decomposition (SVD) to alleviate memory overhead. However, GaLore still suffers from notable computational complexity and memory requirements. Q-GaLore emerges as an enhancement over GaLore by synergizing quantization techniques and layer-adaptive low-rank projections to considerably reduce memory usage without compromising performance.

Key Contributions

Q-GaLore leverages two primary insights:

1. Layer-Specific Gradient Behavior: The gradient subspace exhibits varying behaviors across different network layers. Some layers stabilize early in the training process, while others continue evolving, necessitating frequent updates.
2. Quantization Robustness: Projection matrices exhibit high tolerance to low-bit quantization, performing efficiently even at 4-bit precision.

Methodological Framework

Preliminaries on Quantization

Q-GaLore employs Quantization-Aware Training (QAT), specifically utilizing INT8 for model weights and INT4 for projection matrices. This significantly reduces memory overhead without needing full-precision parameters, unlike traditional QAT methods.

Adaptive Layer-Wise Subspace Exploration

A novel adaptive update mechanism monitors the convergence of gradient subspaces across different layers. By adaptively tuning the frequency of SVD operations based on a monitored threshold of convergence stability, Q-GaLore reduces redundant computations, allowing for significant computational savings. This lazy update strategy is pivotal in reducing SVD operations while maintaining training performance consistency.

Stochastic Rounding for Training Stability

To mitigate information loss during low-precision updates, Q-GaLore adopts Stochastic Rounding (SR). SR provides an unbiased estimation of gradients, ensuring the training trajectory remains stable despite low-bit quantization. This mechanism is crucial in preserving training quality while reducing memory footprint.

Experimental Results

Pre-Training Efficiency

Q-GaLore demonstrates exceptional memory efficiency in pre-training various LLaMA-based models (ranging from 60M to 7B parameters) on the C4 dataset. Noteworthy results include:

• Memory Reduction: Q-GaLore achieves up to 29.68% memory savings compared to GaLore and significant reductions over full-rank training.
• Comparable Performance: Despite aggressive memory optimizations, Q-GaLore's performance remains close to traditional methods, with minimal increases in perplexity.

For instance, training a 7B LLaMA model, Q-GaLore facilitated training within a 16GB memory constraint on an NVIDIA RTX 4060 Ti GPU, showcasing its substantial practicality.

Fine-Tuning Applications

In fine-tuning scenarios across GLUE and MMLU tasks, Q-GaLore maintains competitive performance with significantly reduced memory requirements. The experiments span several architectures (RoBERTa, LLaMA-3-8B, Gemma-7B, and Mistral-7B), consistently outperforming other memory-efficient approaches like LoRA and QLoRA in terms of performance per memory overhead.

Implications and Future Directions

The implications of Q-GaLore's contributions are twofold:

1. Practical Deployment: By enabling effective training of large models on constrained hardware, Q-GaLore democratizes access to high-performance LLM training, making it viable for smaller research entities and applications in edge-computing environments.
2. Theoretical Advances: The adaptive low-rank gradient exploration introduces a novel paradigm in gradient approximation, offering insights into layer-specific behaviors and their potential exploitation for computational savings.

Future developments may focus on further optimizing quantization schemes and exploring the extension of Q-GaLore’s adaptive strategies to other forms of model compression and optimization techniques. Additionally, integrating these methods with advanced hardware architectures could further enhance training throughput and efficiency.

Conclusion

Q-GaLore represents a significant stride in memory-efficient LLM training. Through meticulous integration of low-bit quantization and layer-adaptive low-rank gradients, Q-GaLore achieves noteworthy reductions in memory usage while preserving training performance. This methodology sets a precedent for future research aiming to balance computational efficiency with training efficacy in large-scale neural networks.

## References - Zhao, et al. GaLore: Gradients via Low-Rank Projection for LLMs \cite{zhao2024galore} - Brown, et al. LLMs are Few-Shot Learners. 2020. - Touvron, et al. LLaMA: Open and Efficient Foundation LLMs. 2023. - Hendrycks, et al. Measuring Massive Multitask Language Understanding. 2020. - von Neumann, J. Various techniques used in connection with random digits. 1947.