Unlocking More Granular Control of Memory-Efficient LLM Finetuning

1 Introduction

Large Language Models (LLMs) are powerful but costly to finetune due to their substantial memory requirements. For example, finetuning an LLM with 7 billion parameters in the bfloat16 format [Dean et al., 2012] requires approximately 14 GB of memory for the parameters alone. The gradients demand a similar amount of memory, 1 and, with Adam/AdamW [Kingma and Ba, 2014; Loshchilov and Hutter, 2019] serving as the de facto optimizer, an additional 28 GB is needed for optimizer states. Including the activations required for backpropagation, the total memory footprint exceeds 60 GB, rendering such training prohibitively costly.

1 Techniques like layer-wise updates can mitigate this overhead, but gradient accumulation still necessitates storing gradients in practice.

To mitigate this training overhead, various parameterefficient finetuning (PeFT) techniques [Razdaibiedina et al., 2023] have been developed. Among the most notable are Low-Rank Adapters (LoRA) [Hu et al., 2022] and its variants [Dettmers et al., 2023; Hayou et al., 2024; Wang et al., 2024], which decompose matrix-shaped parameters into two low-rank matrices to enable more memory-efficient training. Recently, a novel family of methods, Low-Rank Gradient Projection (LoRP) [Hao et al., 2024; Zhao et al., 2024; Zhu et al., 2024], has emerged. LoRP methods also leverage low-rank properties during LLM training, but rather than operating at the parameter level, they focus on exploiting the low-rank structure within the gradient. Concretely, LoRP methods achieve memory reduction by projecting the gradient matrix G ∈ Rn×m into a low-dimensional subspace spanned by the columns of a projection matrix P ∈ Rm×r with r ≪ m. When required for parameter updates, the gradient is approximated as (GP )P ⊤ by projecting the stored lowdimensional gradient GP back to the original space. In general, LoRP can offer better memory efficiency than LoRA, as it only requires the storage of gradient information in a single low-dimensional subspace, whereas LoRA necessitates two.

An interesting observation is that LoRP can also be interpreted from the perspective of stochastic approximation, specifically in the form of the forward gradient method [Baydin et al., 2022]: when P is a random matrix sampled from Gaussian distribution N (0, 1r ), LoRP can be viewed as an unbiased stochastic approximation of the original gradient, applied row-wisely (elaborated in Sec.2):

where i = 1, . . . , n, vj ∼ N (0, Im) is the j-th column of √rP, and Gi,: is the i-th row of G, representing a segment of the gradient being approximated. In equation 1, r serves a dual role: it defines the dimensionality of the projected space in LoRP and also determines the number of samples used for stochastic approximation, akin to the query count in Monte Carlo estimation. From the lens of stochastic approximation, the effectiveness of estimation is highly influenced simultaneously by the target dimensionality (i.e., the size of Gi,:) and the number of samples r [Berahas et al., 2022; Shen et al., 2024]. If r increases, more independent samples can reduce variance in equation 1. Alternatively, keeping r fixed while decreasing the dimensionality of Gi,: improves accuracy by reducing the number of dimensions to estimate. This insight naturally leads to a new degree of freedom in LoRP methods: rather than fixing the dimensionality of Gi,:, we can adjust it with r to balance performance and memory efficiency.

Accordingly, we propose VLoRP (Various-Grained LowRank Projection of Gradients), a LoRP framework that introduces a new degree of freedom, projection granularity, defined as the length of each row segment in the gradient matrix G. Unlike existing LoRP methods that implicitly fix the projection unit to one gradient row, VLoRP reshapes the gradient matrix before projection, thereby jointly controlling projection granularity and rank. This enables a more flexible memory-performance trade-off and can be implemented with only lightweight reshaping operations, incurring no additional computational overhead.

To support Adam-style optimization under projectedgradient storage, we further propose ProjFactor, a memoryefficient Adam-like optimizer that avoids access to full-rank gradients. Leveraging VLoRP, we systematically study different granularity-rank configurations under a fixed memory budget, and observe that finer-grained projections consistently outperform coarser alternatives. We also provide theoretical guarantees: VLoRP with SGD achieves an O(1/T) convergence rate, and ProjFactor admits a monotonically decreasing Lyapunov function under the Hamiltonian descent framework [Maddison et al., 2018; Chen et al., 2023; Liang et al., 2024].

Our contributions are summarized as follows:

1. We introduce VLoRP, which extends existing LoRP methods by jointly controlling projection granularity and rank.
2. We propose ProjFactor, a memory-efficient Adam-like optimizer for VLoRP.
3. We show that, under the same memory budget, finer projection granularity is more effective than increasing rank.
4. We provide convergence guarantees for VLoRP with both SGD and ProjFactor.

2 Exploring Various Granularities of Low-Rank Gradient Projection

2.1 Understanding LoRPs Through the Lens of Stochastic Approximation

Given the gradient matrix G ∈ Rn×m and the projection matrix P ∈ Rm×r , LoRP approaches project the gradient into a low-dimensional subspace spanned by columns of projection matrix P, and project it back to the original space when gradient information is required for parameter update:

where we use Gs and Go to represent the projected gradient in the subspace and the project-backed rank-r approximation of the original gradient, respectively, and γ denotes the scaling coefficient and is usually set as 1. In practice, only the matrix Gs ∈ Rn×r needs to be stored with r ≪min{m, n}. Several strategies exist for constructing the projection matrix P. The classical choice is to extract the top singular vectors of the gradient matrix via Singular Value Decomposition (SVD) [Zhao et al., 2024]. However, recent studies [Hao et al., 2024] demonstrate that a matrix whose columns are sampled independently from a high-dimensional isotropic Gaussian distribution preserves the relevant geometric structure equally well-a result anticipated by the Johnson-Lindenstrauss lemma [Matousek, 2008; Dasgupta and Gupta, 2003; Indyk and Motwani, 1998] and corroborated by our experiments. Thus, we adopt Gaussian-random projections as the default. Nevertheless, the proposed framework readily accommodates projection matrices obtained by any other method.

Particularly, if each element of P is i.i.d. sampled from Gaussian distribution N(0, 1r ), then Go can be reformulated into

where vj ∼ N (0, Im) is the j-th column vector of √rP , and Gi,: ∈ R1×m is the i-th row of the gradient matrix G. Equation (3) highlights that: (a) Go can be represented as an average over r rank-one estimation, each involving one random direction vj; (b) Each row 1r ∑rj=1 (v⊤ j G⊤i,: )v⊤ j aligns with the formulation of the averaged forward gradients [Wengert, 1964; Silver et al., 2021; Baydin et al., 2022; Ren et al., 2022], belonging to the broader family of stochastic approximation methods [Robbins and Monro, 1951; Nevel'son and Has' minskii, 1976; Spall, 1992].

Observation 1. Equation equation 3 indicates that LoRP can be interpreted as a row-wise application of forward gradient estimation with r samples for each gradient matrix.

Motivation. When training an LLMparametrized by θ with an objective function L, the stochastic approximation of the gradient can be applied at different levels. Consider two extreme cases: (1) We can flatten the entire parameter set θ into a single vector θ′ ∈ RD , where D is the total number of parameters, and estimate the gradient using a single random vector v ∈ RD . This approach aligns with the spirit of the MeZO algorithm [Malladi et al., 2023]; (2) Alternatively, the estimation can be applied element-wise: for each onedimensional parameter ξ ∈ θ, the partial derivative ∇ξL can be approximated as 1 r′ ∑r′ i=1 vi∇ξLvi, where vi ∼ N(0, 1). Concatenating all such approximations yields an estimate of the full gradient matrix G, resembling the coordinate gradient estimation (CGE) approach [Chen et al., 2024a].

These cases demonstrate that stochastic approximation can be applied to parameters of varying sizes, leading to different methods. Naturally, according to Observation 1, the Projection Granularity, defined by the row size of G, can also be adjusted in LoRPs. This insight motivates our exploration of low-rank gradient projections with varying granularities.

2.2 Various-Grained Low-Rank Projection of Gradient

In this section, we present the details of our framework, termed VLoRP, which introduces a novel degree of freedom-namely, the Projection Granularity-beyond the conventional hyperparameter rank to balance memory efficiency and training performance in LoRP approaches.

The core idea is straightforward: because gradient projection is typically applied row-wisely, one can reshape the gradient matrix to modify the row size and correspondingly adjust the shape of the projection matrix generated, leading to the adjustment of the Projection Granularity naturally. Concretely, given an LLM, we introduce the granularity factor c, which reshapes any gradient G ∈ Rn×m into ˜G ∈ Rnc×(m/c) . Analogous to the rank r, which globally sets the projection rank for all matrix-shaped parameters, c similarly serves as a global hyperparameter that controls the granularity of projection. Then, we formally have:

where the projection matrix ˜ P is of shape mc ×r, with each element i.i.d. sampled fromN(0, 1r ). Under this setting, Equation (2) emerges as a special case with c = 1. Decreasing c < 1 horizontally compresses G, thereby coarsening the Projection Granularity, whereas increasing c > 1 does the opposite. In practice, c is chosen to be a power of two, and both mc and nc must be integers.

The overall framework is illustrated on the left of Figure 1: at each iteration, VLoRP first reshapes the original gradient matrix G (from Rn×m to Rnc×(m/c) ) to adjust the Projection Granularity. The reshaped ˜G is then projected into the subspace to obtain ˜Gs which needs to be stored. During the parameter update, ˜Gs will be projected back to the original space to obtain ˜Go , after which we reshape ˜Go back to update the parameter. Notably, in practice, the only substantive distinction between VLoRP and existing LoRP variants is an additional pair of tensor-reshaping operations, which imposes no extra computational overhead.

2.3 Optimization: Adam-basedMemory-Efficient Schemes

Since our problem setting focuses on memory-efficient finetuning, the original full-rank gradients are not stored in principle. However, the Adam optimizer requires access to the full-rank gradients to update its states, making it memory inefficient if directly applied in this context. Therefore, in this section, we first investigate two potential Adam-based optimization schemes for VLoRP, and subsequently introduce a memory-efficient variant, ProjFactor, for the superior one. We enable gradient accumulation [Smith et al., 2018] for the projected gradient by default, which means at each update stage, we can only access ˜Gs = ∑Ki=1 ˜Gsi /K, where K is the number of accumulation substeps.

Optimization Schemes: SS vs. OS Broadly, with access only to ˜Gs , there are two typical schemes for storing the optimization states of Adam and performing the associated updates for VLoRP, as illustrated on the left of Figure 2: (1) Subspace Scheme (SS) retains the optimization states ms and vs and computes the adapted gradient ˜∆st = ˜mst /√˜ vs t within the subspace, while (2) Original Space Scheme (OS) projects ˜Gs back to the original space, where the optimization states ˜mo and ˜ vo are stored and the adapted gradient is computed directly. Mathematically, the update rules for both schemes are defined as follows 2 :

OS Performs Better For momentum-based SGD optimizers, where the second moment v is excluded, the Subspace Scheme (SS) and Original Space Scheme (OS) are equivalent, since ˜mot = ˜mst ˜ P ⊤. However, when the nonlinear second moment is incorporated, SS and OS become different. An intuition is that the dynamics of OS are closer to Adam's, as (1) both OS and Adam adapt the gradient directly in the original space, and (2) previous works [Zhao et al., 2024; Hao et al., 2024] have shown that the gradient of LLMs exhibits a low-rank structure, implying that ˜G˜ P ˜ P ⊤ is capable of preserving the primary information of ˜G. To substantiate it, we finetune a LLaMA2-7B on three different benchmarks and compare the loss curves of SS, OS, and the vanilla Adam. As depicted in Figure 2 (right, top rows, panels a-c), while all three schemes reduce the loss, SS exhibits a noticeably slower and less effective convergence compared to OS and Adam. In contrast, OS closely tracks the loss curve of Adam.

2 Appendix A provides a comprehensive explanation of notations.

ProjFactor: a Memory-Efficient Optimizer for VLoRP While OS achieves superior performance, it generally requires more memory during training compared to SS, as it does not reduce the memory usage of optimization states, which still occupy O(nm) space. To address this, we propose ProjFactor which (1) maintains ˜ms instead of ˜mo in the subspace and project it back to the original space when calculating ∆ot , which is justified by the fact ˜mo = ˜ms ˜ P ⊤; (2) follows the spirit of Adafactor [Shazeer and Stern, 2018] by applying rank-1 decomposition on ( ˜Gst ˜ P ⊤)⊙2 , which only requires the storage of two vectors, significantly reducing memory consumption while preserving effective second-moment estimates. Compared to Adafactor, which applies the rank-1 decomposition directly on the gradient, our method applies the rank-r approximation of ˜G first before the squaring and decomposition. The final algorithm for VLoRP optimized with Projfactor is outlined in Appendix B. We showcase the performance comparison between ProjFactor and OS in Figure 2(d-f), which aligns with our expectations.

3 Convergence Analysis

Incorporating Projection Granularity into LoRP approaches offers clear benefits, but it is crucial to ensure that variousgrained projections do not compromise the convergence of model training. This section provides such theoretical guarantees. The detailed proofs are deferred to Appendix D.

Convergence of VLoRP under SGD We consider a loss function L : Rn×m → R defined over matrix parameters W ∈ Rn×m. Fix a memory budget M, a granularity factor c, rank r, and G = ∇WL(W), we have

Theorem 1. Let L be an L-smooth function with respect to the matrix-shaped parameter W. Assume the parameter updates are given by Wt+1 = Wt - η Got , where the step size is defined as η = M (m+c+M)L := C. Then, for any T ≥ 1:

where W ∗ is a global minimizer of L.

Theorem 1 confirms that using random projection-based gradient estimation with VLoRP in conjunction with SGD achieves an O(1/T) convergence rate, regardless of the granularity factor c.

Convergence of VLoRP under ProjFactor To analyze the convergence of the proposed Adam-like memory-efficient optimization, ProjFactor, we adopt the Hamiltonian descent framework, following the line of work [Maddison et al., 2018; Chen et al., 2023; Liang et al., 2024; Nguyen et al., 2024]. The infinitesimal updates of Projfactor are defined as follows:

where a, b are constants. The corresponding Lyapunov function (Hamiltonian) is defined as

Then we have the following convergence guarantee:

Theorem 2. Suppose the functions in system equation 6 are continuously differentiable. Under mild assumptions (formally stated in Assumption 1 in Appendix D.2), we have

1. For (Wt, ˜mst , ˜ vo rt , ˜ vo ct ) satisfying equation 6, we have d dt H(Wt, ˜mst , ˜ vo rt , ˜ vo ct ) ≤ 0.
2. Any bounded solution (Wt, ˜mst , ˜ vo rt , ˜ vo ct )t of equation 6 converges to a stationary point of L(W) as t →∞.

Theorem 2 presents a Hamiltonian interpretation of Projfactor, indicating that the Lyapunov function, namely, the "energy" of the system, decreases monotonically over time. In addition, it ensures that ProjFactor stabilizes at a local optimum, provided the step sizes are sufficiently small.

4 Experimental Results

4.1 Fine-Grained Gradient Projection Improves Performance

VLoRP introduces a novel degree of freedom that extends beyond the rank parameter in LoRP approaches. This raises a critical question: Which level of the Projection Granularity is optimal for gradient projection?

To address this question, it is essential to first establish a fair basis for comparison: for a fixed granularity factor c, increasing the rank typically enhances performance while incurring additional memory costs; Similarly, for a fixed rank r, employing a finer-grained projection (larger c) also increases memory requirements, since the shape of the projected gradient ˜Gs is [nc, r]. We thereby denote a specific pair of (c, r) as a "configuration" of VLoRP, and define the "Memory Budget"Mas the product of c and r-for configurations sharing the sameM, the size of ˜Gs , which is needed to be stored, is the same.

Based on this, we conducted experiments on three representative benchmarks (Commonsense Reasoning, MMLU, GSM8K) to investigate how different configurations of projection granularity and rank affect finetuning performance under an identical memory budget.

The results are illustrated in Figure 3(a-c), where the yaxis represents performance (higher is better) and the x-axis indicates varying projection granularities defined by c = 45-i and r = 4i-1 . Thus, an increase in i corresponds to a smaller c and a larger r, with a smaller c representing coarser projection granularity. Notably, varying i does not change the product of c and r, thereby maintaining a constant overall memory cost. We observed that VLoRP configurations with finer projection granularity (i.e., larger c and smaller r) always yield higher performances. Specifically, the finest-grained VLoRP configuration (c = 44 , r = 40 ) achieves the highest accuracy scores of 61.40, 55.83, and 29.42 on three benchmarks, respectively, among all tested configurations sharing the same memory budget. These results suggest that projection granularity may have a more critical role than rank in balancing memory efficiency and performance. Furthermore, in Figure 3(d), we present the performance of VLoRP across different granularity factors under varying memory budgets.

It is evident that finer-grained projections generally outperform coarser ones.

Beyond the performance advantages, finer-grained projection also offers several additional benefits: 1. Efficient Matrix Multiplication. Given the granularity factor c, the memory budget M and the reshaped gradient ˜G ∈ Rnc×(m/c) , the gradient projection operation ˜G˜ P leads to a computational complexity O ( nmM c ); 2. Faster Generation of Projection Matrix. Obviously, the efficiency of this generation process depends on the dimensions of the projection matrix ˜ P of the shape mc × Mc . As c increases, the projection granularity becomes finer, leading to a quadratic reduction in the size of the generated matrix and hence more efficient generation. The training runtimes reported in Table 4 empirically validate these, showing that the configuration with the finer granularity achieves the faster training speed. This improvement is the result of the two benefits described above; 3. Lower Numerical Error. As c increases, the numerical stability of the random projection tends to improve. As discussed above, a larger c reduces the number of floatingpoint arithmetic operations in matrix multiplications, which can help mitigate the accumulation of numerical errors. This is particularly important in low-precision training scenarios such as float16 or bfloat16, where limited mantissa precision can lead to instability. Furthermore, this numerical stability advantage is further amplified when using gradient accumulation or Adam-based optimizers, as these methods involve iteratively accumulating optimization states. We provide an analytical experiment in Appendix C.3 to support this claim.

4.2 Comprehensive Experiments

In this section, we evaluate the effectiveness of VLoRP across multiple finetuning tasks, demonstrating its competitiveness with state-of-the-art baselines. More experimental studies and implementation details can be found in Appendix C.

Datasets We present a comprehensive evaluation of our proposed approaches using three benchmarks: (1) Commonsense Reasoning, which covers 8 reasoning tasks including BoolQ [Clark et al., 2019], PIQA [Bisk et al., 2020], SIQA [Sap et al., 2019], Hellaswag [Zellers et al., 2019], WinoGrande [Sakaguchi et al., 2019], ARC-e and ARCc [Clark et al., 2018], and OBQA [Mihaylov et al., 2018]; (2) The MMLU benchmark [Hendrycks et al., 2020], which encompasses a wide range of subjects including Humanities, STEM, Social Sciences, and Other fields; (3) Finally, the GSM8K dataset [Cobbe et al., 2021], which consists of 8.5K high-quality math problems.

Baselines We compare VLoRP with 7 state-of-the-art baselines, covering full-parameter finetuning, LoRA, and LoRP methods. Specifically, we compare with Untuned, which represents the basic performance of the pre-trained model without finetuning, Adam [Kingma and Ba, 2014], Adafactor [Shazeer and Stern, 2018], LoRA [Hu et al., 2022], Galore [Zhao et al., 2024], fira [Chen et al., 2024b], and APOLLO [Zhu et al., 2024].

Experimental Settings We test Llama2-7B, Llama3-3B, Qwen3-4B models with the bfloat16 data type. We ensure that for each tested method, every token in the training set is encountered at least once. Gradient accumulation and activation checkpointing [Chen et al., 2016] are enabled.

Results and Analysis We report the performance of all methods on: (1) commonsense reasoning tasks in Table 1, (2)

the MMLU benchmark in Table 2, and (3) the GSM8k task in Figure 4. Generally, every fine-tuning strategy confers substantial performance gains over the untuned model. Among the fine-tuning baselines, full-parameter updates with Adam or Adafactor consistently achieve the strongest results across tasks; nevertheless, the inherent nature of full-parameter fine-tuning precludes memory efficiency. Notably, Adafactor delivers comparable even superior results to Adam, consistent with previous findings [Shazeer and Stern, 2018; Hao et al., 2024]. Moreover, the proposed VLoRP not only outperforms all memory-efficient fine-tuning baselines but in many instances achieves results comparable to full-parameter fine-tuning. Given the difficulty of the benchmarks experimented with, this result is decidedly non-trivial.

Configurations Runtime

c = 4-2 , r = 46 183s/iter

c = 40 , r = 44 165s/iter

c = 42 , r = 42 161s/iter

c = 44 , r = 40 157s/iter

Table 4: Training speed for different configurations. We tested with Llama2-7B model, bfloat16 data type, 512 batch size, and 1024 input length.

Memory Analysis We illustrate the memory usage of various methods in Table 3, emphasizing that VLoRP achieves the most significant memory reduction compared to baseline approaches. Once gradient accumulation is enabled, certain LoRP-based schemes-most notably Fira and APOLLO-require substantially more VRAM than other low-rank approaches. This overhead stems from their reliance on the full-rank gradient during weight update: they compress only the optimiser states, leaving the gradient tensor itself untouched. By contrast, VLoRP intrinsically lowers memory usage by operating exclusively on projected gradients. Specifically, for a parameter matrix W ∈ Rn×m, the total memory required for VLoRP is O(mn+2nM+n+m), whereM = c · r is the memory budget. On the other hand, LoRA requires significantly more storage, O(mn + 4mM+ 4nM), due to its need for additional low-rank matrices and optimization states.

5 Conclusion

In this work, we introduce VLoRP, a memory-efficient finetuning method for LLMs, which implements low-rank gradient projections with varying granularities. By adjusting the projection granularity alongside rank, VLoRP offers a more nuanced control over the trade-off between memory efficiency and performance. Convergence analysis and extensive empirical results confirm the effectiveness of our approach, making it a promising solution for practical deployment in memory-constrained settings.

Acknowledgements

This material is based upon work supported by the Air Force Office of Scientific Research under award number FA238624-1-4011, and this research is partially supported by the Singapore Ministry of Education Academic Research Fund Tier 1 (Award No: T1 251RES2509).