Memory-Efficient LLM Training by Various-Grained Low-Rank Projection of Gradients

The rapid development of large language models (LLMs) such as GPT-3, LLaMA, and PaLM has demonstrated remarkable capabilities in NLP tasks. However, their enormous parameter counts—often hundreds of billions—bring substantial computational and memory costs, hindering widespread fine-tuning and deployment. Parameter-efficient fine-tuning (PeFT) methods like LoRA, which decomposes weight matrices into low-rank factors, have been effective in reducing storage and computation. Similarly, low-rank gradient projection (LoRP) exploits the low-rank structure within gradients to further compress memory. Nonetheless, existing LoRP approaches typically fix the projection unit at the row level, limiting the potential for finer control. As models grow larger, the need for more flexible, scalable, and memory-efficient methods becomes urgent, prompting research into multi-scale gradient compression techniques.

The core challenge lies in balancing memory constraints with training stability and model performance. Traditional low-rank methods fix the projection granularity at the row level, which may not be optimal across different tasks or model sizes. Coarser granularity reduces memory but can impair gradient estimation accuracy, leading to unstable training. Conversely, finer granularity improves stability but increases memory overhead. Finding an optimal configuration that adapts to resource limits while maintaining high performance remains an open problem, especially as model sizes continue to expand beyond current hardware capabilities.

VLoRP introduces the concept of Projection Granularity, allowing the reshaping of the gradient matrix into various granularities beyond the fixed row-wise unit. This flexibility enables a systematic exploration of the trade-off between memory and accuracy. The framework combines this with low-rank projection, stochastic approximation, and random Gaussian projections to improve gradient estimation. Additionally, ProJFator, an adaptive optimizer, is designed to work efficiently with these configurations, reducing memory footprint during training. Theoretical analysis guarantees convergence rates comparable to standard SGD, even with varying granularity and rank, providing a solid foundation for practical deployment.

• �� Reshape the gradient matrix G from Rn×m to Rnc×(m/c), where c is the granularity factor. • Generate a random projection matrix of shape (m/c)×r with entries sampled from N(0, 1/r). • Project the reshaped gradient into a low-dimensional subspace, store it, and during updates, project back to the original space. • Adjust the granularity c to control the size of the fundamental projection unit, balancing memory and estimation accuracy. • Design ProJFator to adaptively optimize the projected gradients, incorporating low-rank approximations and momentum. • Theoretically, prove that the convergence rate remains O(1/T) under fixed memory budgets, regardless of granularity settings.

Experiments conducted on LLaMA2-7B fine-tuning across tasks like MMLU, GSM8K, and commonsense reasoning. Different configurations of granularity c and rank r were tested under fixed memory budgets (e.g., M=256). Baseline comparisons included LoRA, Galore, and Adam. The impact of granularity on performance, stability, and memory consumption was analyzed. ProJFator was evaluated in gradient accumulation scenarios, demonstrating significant memory savings while maintaining accuracy. Convergence behavior was validated through loss curves and theoretical bounds, confirming the robustness of the approach.

Finer granularity configurations consistently outperformed coarser ones, with performance gains of approximately 3% on MMLU. ProJFator reduced memory usage by over 50%, enabling larger batch sizes and faster training. Theoretical proofs confirmed the O(1/T) convergence rate, matching standard SGD. Ablation studies showed that granularity has a more significant impact than rank within fixed memory budgets, highlighting the importance of multi-scale gradient compression.

This framework is suited for resource-constrained environments requiring large model fine-tuning, such as edge devices and cloud services. It enables efficient multi-task learning, rapid deployment, and adaptation of models in NLP, computer vision, and multimodal applications. The approach can be integrated into existing training pipelines, offering scalable solutions for industry and academia, especially as models continue to grow in size.

The method's effectiveness depends on the choice of granularity and rank, which may require empirical tuning. Extreme granularity settings could lead to instability or performance degradation. The theoretical assumptions of smoothness and randomness may not hold in all practical scenarios. Managing random projection matrices at very large scales adds computational overhead, necessitating further optimization.