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

TL;DR

Memory constraints hinder fine-tuning large language models, motivating low-rank gradient projection (LoRP) as a memory-efficient alternative. The paper introduces VLoRP, which adds a Projection Granularity parameter $c$ alongside rank $r$ with memory budget $\mathcal{M}=cr$, and proposes ProjFactor, a memory-efficient optimizer for this framework. Theoretical results establish $G^o$ as an unbiased gradient estimate with variance bounded by $\mathcal{O}((m+c)/\mathcal{M})$, and an $O(1/T)$ convergence rate under SGD, with Hamiltonian-descent guarantees for ProjFactor. Empirically, finer granularity under a fixed budget improves performance and stability across multiple tasks (commonsense, MMLU, GSM8K) with competitive memory usage and throughput, demonstrating practical gains for memory-constrained LLM fine-tuning.

Abstract

Building upon the success of low-rank adapter (LoRA), low-rank gradient projection (LoRP) has emerged as a promising solution for memory-efficient fine-tuning. However, existing LoRP methods typically treat each row of the gradient matrix as the default projection unit, leaving the role of projection granularity underexplored. In this work, we propose a novel framework, VLoRP, that extends low-rank gradient projection by introducing an additional degree of freedom for controlling the trade-off between memory efficiency and performance, beyond the rank hyper-parameter. Through this framework, we systematically explore the impact of projection granularity, demonstrating that finer-grained projections lead to enhanced stability and efficiency even under a fixed memory budget. Regarding the optimization for VLoRP, we present ProjFactor, an adaptive memory-efficient optimizer, that significantly reduces memory requirement while ensuring competitive performance, even in the presence of gradient accumulation. Additionally, we provide a theoretical analysis of VLoRP, demonstrating the descent and convergence of its optimization trajectory under both SGD and ProjFactor. Extensive experiments are conducted to validate our findings, covering tasks such as commonsense reasoning, MMLU, and GSM8K.

Figures (12)

• Figure 1: Overview of VLoRP versus standard LoRP. Bottom (gray) : In ordinary LoRP, the gradient matrix $G$ is directly projected row‐by‐row from $\mathbb{R}^{n\times m}$ into $\mathbb{R}^{n\times r}$ and stored as $G^{s}$. Top (pale beige): In contrast, VLoRP reshapes the original gradient matrix $G$ first to adjust the granularity of projection (from $\mathbb{R}^{n\times m}$ to $\mathbb{R}^{nc \times\frac{m}{c}}$), and during the update, the project-backed gradient $\tilde{G}^o$ would be reshaped back into $\mathbb{R}^{n\times m}$ to update the parameter $W\in \mathbb{R}^{n\times m}$.
• Figure 2: Left: Schematic illustration of the Subspace Scheme (SS, green) operating in a learned subspace, and the Original Scheme (OS, blue) operating in the original space. Right (top row, panels a--c): Fine-tuning loss curves of SS (green), OS (blue), and Adam (orange) on three tasks, showing that OS outperforms SS by a large margin while has a comparable performance with Adam. Right (bottom row, panels d--f): Comparison of OS (blue) and its approximate algorithm ProjFactor (purple), indicating that ProjFactor closely approximates the dynamic of OS. Specifically, columns (a) and (d) test on the Commonsense Reasoning task, columns (b) and (e) test on the MMLU, and columns (c) and (f) test on the GSM8K.
• Figure 3: Left: Performance comparison of different methods on GSM8K. Right: Performance comparison among the configurations of VLoRP with $\mathcal{M}=256$. The x-axis indicates configurations from fine to coarse (left to right).
• Figure 4: Performance Evaluation for Different Projection Granularities under Varying Memory Budgets on Commonsense. $\mathcal{M}$ denotes the memory budget. All subcolumns inside the same category $\mathcal{M}/r$ share the same rank $r$.
• Figure 5: GPU Memory Usage Comparison on LLaMA2-7B Model with Batch Size 16 and Max Length 1024.

Theorems & Definitions (12)

• Proposition 3.2
• Theorem 3.3
• Theorem 4.1
• Lemma 3.1
• proof
• Proposition 3.2
• proof
• Theorem 3.3
• proof
• Theorem 3.5