Title: Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Model

Memory Efficient Training

Full-rank Training
Abstract: Low-rank training has emerged as a promising approach for reducing memory usage in training Large Language Models (LLMs). Previous methods either rely on decomposing weight matrices (e.g., LoRA), or seek to decompose gradient matrices (e.g., GaLore) to ensure reduced memory consumption. However, both of them constrain the training in a low-rank subspace, thus inevitably leading to sub-optimal performance. To resolve this, we propose a new plug-and-play training framework for LLMs called Fira, as the first attempt to consistently preserve the low-rank constraint for memory efficiency, while achieving full-rank training (i.e., training with full-rank gradients of full-rank weights) to avoid inferior outcomes. 
First, we observe an interesting phenomenon during LLM training: the scaling impact of adaptive optimizers (e.g., Adam) on the gradient norm remains similar from low-rank to full-rank training. In light of this, we propose a \textit{norm-based scaling} method, which utilizes the scaling impact of low-rank optimizers as substitutes for that of original full-rank optimizers to achieve this goal.
Moreover, we find that there are potential loss spikes during training. To address this, we further put forward a norm-growth limiter to smooth the gradient.
Extensive experiments on the pre-training and fine-tuning of LLMs show that Fira outperforms both LoRA and GaLore. Notably, for pre-training LLaMA 7B, our Fira uses $8\times$ smaller memory of optimizer states than Galore, yet outperforms it by a large margin.

Fira: Can We Achieve Full-rank Training of LLMs under Low-rank Constraint?

Abstract

Low-rank training has emerged as a promising approach for reducing memory usage in training Large Language Models (LLMs). Previous methods either rely on decomposing weight matrices (e.g., LoRA), or seek to decompose gradient matrices (e.g., GaLore) to ensure reduced memory consumption. However, both of them constrain the training in a low-rank subspace, thus inevitably leading to sub-optimal performance. To resolve this, we propose a new plug-and-play training framework for LLMs called Fira, as the first attempt to consistently preserve the low-rank constraint for memory efficiency, while achieving full-rank training (i.e., training with full-rank gradients of full-rank weights) to avoid inferior outcomes. First, we observe an interesting phenomenon during LLM training: the scaling impact of adaptive optimizers (e.g., Adam) on the gradient norm remains similar from low-rank to full-rank training. In light of this, we propose a norm-based scaling method, which utilizes the scaling impact of low-rank optimizers as substitutes for that of original full-rank optimizers to achieve this goal. Moreover, we find that low-rank optimizers may lead to potential loss spikes during training. To address this, we further put forward a norm-growth limiter to smooth the gradient. Extensive experiments on the pre-training and fine-tuning of LLMs show that Fira outperforms both LoRA and GaLore. Notably, for pre-training LLaMA 7B, our Fira uses 8 × smaller memory of optimizer states than Galore, yet outperforms it by a large margin.

1 Introduction

In recent years, Large Language Models (LLMs) have achieved remarkable advancements in various domains [2, 30, 10]. While the substantial increase in model size contributes significantly to these advancements, it also introduces considerable memory bottlenecks, especially for optimizer states [43]. For instance, pre-training a LLaMA-7B model from scratch 2 requires at least 58 GB memory, allocated as follows: 14GB for loading parameters, 14GB for weight gradients, 28GB for Adam [16] optimizer states, and 2GB for activations [43]. Notably, the optimizer states consume even more memory than the parameters themselves. To address this, low-rank training has demonstrated

full-gradient of full-parameter

Trainable

update

full-rank gradient

s

of low-r

ank weight

s

Frozen

Trainable

Trainable

update

(a) LoRA

full-gradient of full-parameter

low

-rank

gradient

s

of

full-rank weight

s

Trainable

Trainable

update

update

full-rank gradients

of

full-rank weights

subspace-gradient of full-parameter

Trainable

Trainable

update

update

full-gradient of subspace-parameter

Frozen

Trainable

Trainable

update

(c) Fira (ours)

its effectiveness in reducing the memory usage of the optimizer states by conducting training in a low-rank subspace [43, 13].

The current low-rank training methods can be broadly divided into two categories: weight matrixbased and gradient matrix-based low-rank decomposition. For the weight matrix decomposition methods, the most representative one is Low-Rank Adaptation (LoRA) [13], where its basic idea is to use low-rank matrices as decomposed representations of the pre-trained weights during training, as shown in Figure 1 (a). However, the optimization of LoRA is constrained in a low-rank subspace of the weights. This will inevitably cause the reduction of representation capacity, leading to suboptimal outcomes [40, 35]. Although the variant ReLoRA [18] attempts to extend the application of LoRA from fine-tuning to pre-training, by periodically updating high-rank weights with multiple low-rank updates. It still requires full-rank weight training as a warm-up before low-rank training, thus rendering memory efficiency unachievable [43].

For the gradient matrix decomposition based methods, the typical one is the gradient low-rank projection (GaLore) proposed recently [43]. In contrast to LoRA, GaLore attempts to reduce the memory usage in optimizer states via decomposing the gradient matrix, as shown in Figure 1 (b). While GaLore supports the training of full-rank weights, it leverages only low-rank gradients, restricting them to a low-rank subspace. Consequently, any gradient information outside this subspace is lost, in contrast to training with full-rank gradients. Note that since these methods constrain LLM training to a low-rank subspace, this inevitably leads to sub-optimal results compared to full-rank training (i.e., training with full-rank gradients and full-rank weights). This raises the question: Can we achieve full-rank training for LLMs while consistently maintaining a low-rank constraint?

In light of this, we propose a new memoryefficient training framework for LLMs, called Fira , which, to the best of our knowledge, is the first to achieve f ull-rank tra i ning while consistently maintaining a lowra nk constraint. To achieve this goal, a significant challenge is that the low-rank constraint makes it hard to preserve complete optimizer states (e.g., gradient momentum and variance) of full-rank weights in the commonlyused adaptive optimizer (e.g., Adam). As a result, the adaptive optimizer fails to correct the raw gradients without corresponding op-

Table 1: Cosine Similarity and MSE between lowrank and full-rank scaling factors at the matrix and column levels for pre-training LLaMA models ranging from 60M to 1B parameters, averaged over 10K steps. Detailed analyses are presented in Appendix D.

timizer states. Without this correction, adaptive optimization algorithms would degrade into simple stochastic gradient descent (SGD), leading to significantly reduced optimization performance [16, 39]. This point is further validated in Section 4.1 and Section 5.4.

Fortunately, we observe an interesting phenomenon during LLM training: at weight matrix level, the scaling factors 3 of low-rank gradients exhibit strong similarity to those of full-rank gradients. As illustrated in Table 1, Cosine Similarity and Mean Squared Error (MSE) between low-rank and full-rank scaling factors exhibit significant similarity (Cosine Similarity close to 1, while MSE close to 0). Detailed quantitative analyses of this similarity are presented in Appendix D.

Based on this observation, we put forward a norm-based scaling method that utilizes the scaling factor of a weight matrix in low-rank training to replace the corresponding matrix's scaling factor in full-rank training. In this way, our scaling factor can also play a similar role in correcting the raw gradient, as adaptive optimizers do. Therefore, we can enable full-rank training while preserving the low-rank constraint.

However, we observe potential sudden increases in gradients caused by low-rank optimizers during the early stages of training, which can lead to spikes in training loss, as illustrated in Figure 3. This phenomenon primarily arises because our norm-based scaling method lacks the gradient stability compared to the original full-rank Adam optimizer. Such gradient instability can trigger substantial parameter updates, causing the loss function to reach a much higher value and undermining prior optimization efforts [11, 39]. Despite the use of gradient clipping techniques [25], this issue may not be adequately resolved, as shown in Figure 3. To this end, we propose a norm-growth limiter, which aims to smooth the gradient by restricting the magnitude of the gradient norm's increase. By employing our limiter, we adaptively convert sudden rises in gradients into gradual increases, thereby facilitating a smooth update that effectively mitigates the problem of loss spikes.

Our main contributions can be summarized as follows:

We propose Fira, a plug-and-play memory-efficient training framework of LLMs, constituting the first attempt to enable full-rank training consistently under the low-rank constraint. We will release the source code and package of our Fira into a Python library for easy use.

We design two components in Fira: a norm-based scaling strategy that leverages the scaling effects of low-rank optimizers to facilitate full-rank training, and a norm-growth limiter to address the issue of loss spikes by limiting the growth of gradient norm.

Extensive experiments on the pre-training and fine-tuning of LLMs across various parameter counts (60M, 130M, 350M, 1B, 7B) validate the effectiveness of Fira in both pre-training and fine-tuning tasks, surpassing both LoRA and GaLore.

2 Related Work

Low-rank adaptation. Low-Rank Adaptation (LoRA) has been introduced by [13] as an efficient fine-tuning method for LLMs. The core idea of LoRA is to freeze the pre-trained weights and introduce trainable low-rank matrices as decomposed representations of the pre-trained weights. In this way, the memory usage of training LLMs could be saved. Recently, a variety of methods by extending LoRA have been proposed to further improve the performance [41, 33, 35, 40, 9]. For instance, ReLoRA [18] is proposed to extend the application of LoRA from fine-tuning to pretraining. However, it still requires full-rank warm-up training before low-rank training, which prevents achieving memory efficiency. It is worth noting that while LoRA-based methods reduce memory usage by limiting training to a low-rank parameter subspace, they inevitably reduce representation capacity [35].

Gradient projection. Recent works [40, 35, 32] have indicated that LoRA may yield sub-optimal performance since its low-rank constraints in parameters. Inspired by traditional projected gradient descent methods [6, 4], GaLore [43] has been proposed recently to mitigate this problem. It enables full-parameter learning under low-rank constraints by projecting the gradient into a low-rank subspace, reducing memory usage for optimizer states. However, while GaLore allows memory-efficient fullparameter training, it confines the gradient to a low-rank subspace, discarding the portion outside the subspace and resulting in significant information loss.

System-Based memory-efficient techniques. Many system-based techniques have been developed to reduce memory usage in LLM training [5, 27]. However, most of these methods achieve memory efficiency by compromising either time or precision. Gradient checkpointing [5] is proposed to reduce memory usage by trading increased computational time for the re-computation of activations. Quantization [9] reduces memory consumption by using lower-bit data types but at the cost of model precision. Memory offloading [38, 27] reduces GPU memory usage by using non-GPU memory (e.g., CPU) as an extension. However, it introduces additional communication overhead, such as CPU-GPU transfer time. It's important to note that our proposed method is complementary to these approaches and can potentially be combined with them to further reduce memory usage.

3 Preliminaries

3.1 Regular Full-Rank Training

At time step t , we denote the full-rank weight matrix as W t ∈ R m × n . The full-rank gradient can be represented as G t = ∇ W f t ( W t ) ∈ R m × n , where f is the objective function. Then the regular full-rank training can be expressed as follows:

where η is the learning rate, and ψ t is the gradient correction function of the optimizer (for vanilla SGD, ψ t ( G t ) = G t ). Instead of vanilla SGD, adaptive optimizers (e.g., Adam [16], AdamW [21]) are usually employed to correct the raw gradient for improving the training performance. However, this typically requires additional memory for storing optimizer states used in gradient correction. For instance, Adam [16] requires storing the optimizer states M and V , which consume 2 mn of memory. The gradient correction process is as follows:

where all matrix operations are element-wise. β 1 and β 2 are Adam's hyper-parameters, and ϵ is a small constant (e.g., 1 × 10 -8 ) used for numerical stability. Since this regular full-rank training typically consumes a large amount of memory for training LLMs, many representative low-rank training methods, e.g., LoRA [13] and Galore [43], have been proposed to reduce memory usage in recent years.

3.2 Low-Rank Adaptation

The basic idea behind LoRA [13] is to use low-rank matrices as decomposed representations of the pre-trained weights during training, in order to reduce memory usage. Formally, LoRA freezes the full-rank weight matrix W 0 ∈ R m × n and incorporates two low-rank matrices A t and B t for training as:

where B t ∈ R m × r , A t ∈ R r × n , and the rank r < min( m,n ) . While LoRA reduces memory usage by limiting training to a low-rank subspace of the weight, it inevitably diminishes the representation capacity of the weight matrix W t [35].

3.3 Gradient Low-Rank Projection

In contrast to LoRA, GaLore [43] utilizes a projection matrix P t ∈ R m × r to project the full-rank gradient G t ∈ R m × n to a low-rank gradient R t = P ⊤ t G t ∈ R r × n ( m ≤ n ) 4 . By doing so, the memory usage of optimizer states could be reduced. The parameter update in GaLore can be formulated as:

where the projection matrix P t can be obtained through singular value decomposition (SVD) of G t and can be updated every T step:

where u i is the i -th column vector of the left singular matrix U . By selecting the first r columns of matrix U that correspond to the largest singular values, the projection matrix P t effectively captures the most significant directions in the gradient space, leading to faster convergence [43]. The optimal switching frequency T is usually set to be between 50 to 1000 , and the additional computational overhead introduced by SVD is negligible ( < 10% ), as stated in [43]. Since Galore restricts the gradient in the low-rank subspace, the gradient information outside this subspace is lost, leading to inferior performance.

4 Proposed Method

To achieve full-rank training under low-rank constraints, our framework, named Fira, consists of two important components: (i) a norm-based scaling method, enabling full-rank training by leveraging the scaling effects of adaptive optimizers; (ii) a norm-growth limiter, which restricts the growth of the gradient norm to prevent spikes in training loss. Next, we will elaborate on these two components.

4.1 Norm-Based Scaling

The low-rank constraint makes it challenging to record complete optimizer states for correcting raw gradients in full-rank training. Fortunately, we find an interesting phenomenon in LLM training: the scaling factor at the matrix level remains similar from low-rank training to full-rank training. Based on this observation, we propose a norm-based scaling strategy that approximately corrects the raw gradient, similar to adaptive optimizers, thereby enabling full-rank training.

Challenge analysis. Given the difficulty of incorporating trainable low-rank weights into LoRA to achieve full-rank weight training [43], we focus on investigating how to achieve fullrank gradient training by extending the gradient projection method, Galore, in this paper. In GaLore, the projection matrix P t ∈ R m × r projects the full-rank gradient G t ∈ R m × n of the fullrank weight W t ∈ R m × n , to the low-rank subspace gradient R t = P ⊤ t G t ∈ R r × n . Then, the gradient outside this subspace can be represented as: ( I -P t P ⊤ t ) G t = G t -P t R t . In other words, the full-rank gradient G t can be divided into two terms: P t R t and ( G t -P t R t ) .

In GaLore, the optimizer states only store the information of R t instead of G t to realize the



































low-rank constraint. The term of ( G t -P t R t ) is directly discarded in Galore due to the lack of corresponding optimizer states for correction in optimizers. This would lead to significant information loss especially when r ≪ d model , where d model = min( m,n ) is the full-rank dimension of models (This point can be verified in our experiment section, as illustrated in Figure 5. In Figure 5, the validation perplexity of GaLore significantly increases at r = 4 compared to r = 128 when d model = 256 , indicating a substantial loss of information and decreased training performance). Intuitively, to capture the information of ( G t -P t R t ) , we can directly add it based on Eqn. 6 as follows:

We denote the update strategy in Eqn. 8 as GaLore-add. However, as illustrated in Figure 2, GaLoreadd exhibits almost no improvement compared to updates using Eqn. 6 in GaLore. This phenomenon primarily arises because the term of ( G t -P t R t ) doesn't have corresponding optimizer states for gradient correction. As a result, the optimization of ( G t -P t R t ) uses vanilla SGD, yielding suboptimal outputs. Besides, in GaLore-add, P t ψ t ( R t ) employs the Adam optimizer for training while ( G t -P t R t ) employs vanilla SGD. This gradient misalignment may also account for the lack of noticeable improvement.

Similarity of scaling factor. To tackle this challenge, we propose the concept of the scaling factor , which is defined as follows:

where the scaling factor ϕ t represents the magnitude of the correction applied by the adaptive optimizer to the gradient norm. Based on the scaling factor ϕ t , we observe an interesting phenomenon during LLM training: the scaling factors at the matrix level exhibit a high degree of similarity between low-rank and full-rank training (column level will be introduced later). As shown in Table 1, Cosine Similarity and MSE between low-rank and full-rank scaling factors exhibit significant similarity (Cosine Similarity close to 1, while MSE close to 0). Details analyses and additional experiments of this similarity are presented in Appendix D.

Norm-based scaling. Building on the above observation, we propose a norm-based scaling method that leverages the scaling factor of a weight matrix in low-rank training as a substitute for the corresponding factor in full-rank training:

Through Eqn. 10, we can approximately correct ( G t -P t R t ) as adaptive optimizers do, so as to achieve full-rank training under low-rank constraints.

To further enhance our approach, we introduce a more fine-grained strategy for computing scaling factors in Eqn. 10 by considering each column of the weight matrix individually:

where R t, : ,i is the i -th column of R t , and ϕ t ( R t ) i is the i -th scaling factor. As evidenced in Table 1, scaling factors computed at the column level also demonstrate strong similarity between low-rank and full-rank training, further validating the effectiveness of this fine-grained strategy. Additional theoretical analysis and experimental results on scaling factors are provided in Appendix E.

4.2 Norm-Growth Limiter

However, limited by low-rank constraints, we notice that there are still shortcomings in the normbased scaling approach that need to be improved compared to the full-rank Adam. Specifically, potential sharp increases of the gradient may occur at the early stage of training, leading to spikes of training loss. As shown in Figure 3, Fira-w.o.-limiter (our method without using the proposed norm-growth limiter) experiences spikes in both gradient norm and training loss. In this section, we analyze the reasons for this issue and propose a norm-growth limiter that transforms abrupt gradient spikes into gradual, smooth increases, and then enhances the norm-based scaling method.

Loss spike analysis. There are two main reasons for the spikes: (i) Switching the projection matrix P t in gradient projection methods would cause instability during training. As illustrated in Figure 2, both GaLore and GaLore-add exhibit significant training loss spikes at integer multiples of T = 200 (i.e., the frequency of switching the projection matrix P t ). This instability occurs because, when switching projection matrices P t , the optimizer retains states linked to the previous matrix, while the current input gra-







































Spike 1

Spike 2

Spike 1

Spike 2

dient uses a new projection matrix, leading to significant misalignment. Furthermore, as shown in Figure 2, GaLore-add also exhibits additional training spikes compared to Galore, reinforcing our earlier claim that directly incorporating ( G t -P t R t ) may introduce instability and hinder training; (ii) Maintaining the original direction of the raw gradient ( G t -P t R t ) may be insufficient for handling the sharp loss landscapes in LLM training, unlike Adam [39]. Due to space constraints, further analysis is provided in Appendix G.

Addressing loss spikes. To address this issue, a straightforward solution is to use gradient clipping techniques [25] to avoid loss spikes. However, clipping based on the absolute norm of gradient matrices fails to account for significant differences between them, leading to sub-optimal results. This point can also be verified in Figure 3 and Table 5. To this end, we propose a norm-growth limiter method that constrains the ratio of the current gradient norm to the previous step's norm to a fixed ratio γ when the gradient norm increases:

where γ is a threshold ensuring that the rate of gradient growth does not exceed this value. S t = ϕ t ( R t )( G t -P t R t ) is the corrected gradient by applying our norm-based scaling. This approach

limits the magnitude of gradient norm increases, converting sudden spikes into gradual rises and thus preventing loss spikes. Moreover, by constraining the relative increase of each gradient matrix's norm, our method is more flexible than the absolute norm clipping. As illustrated in Figure 2 and Figure 3, Fira with our proposed limiter improves the optimization performance without significant spikes.

Discussion with other stabilization methods. A variety of approaches have been developed to mitigate loss spikes and enhance training stability, including embedding normalization [17], gradient shrink on embedding layers [36], tensor-wise scaling [8]. While these methods contribute to stabilization, they either primarily concentrate on embedding stabilization (e.g., embedding normalization and gradient shrink), or employ static scaling (e.g., tensor-wise scaling). Compared to our norm-growth limiter, these techniques are insufficient in addressing the adaptive instability that arises from the substantial variations among different weight matrices, thereby inadequately resolving the issue of loss spikes in norm-based scaling. Further detailed analyses and experimental results regarding the norm-growth limiter are discussed in Section 5.4 and Appendix G.

4.3 Overall Algorithm

We present the overall algorithm of Fira with Adam in Algorithm 1 in Appendix A.1. Our main components, the norm-based scaling method, and the norm-growth limiter, are straightforward to implement, requiring only 3 additional lines of code. Moreover, Fira is a plug-and-play framework that can be easily integrated into the training process without requiring significant modifications. The plug-and-play Pytorch-like pseudo-code of Fira is provided in Appendix A.2.

It's worth noting that compared to Galore, Fira only introduces one parameter ∥ S t -1 ∥ for each weight matrix in the optimizer state, which is negligible, as shown in Table 2. Besides, in addition to the original hyper-parameters of optimizers and gradient projection methods, Fira only adds one hyper-parameter γ in the norm-growth limiter, whose choice is not sensitive to the performance of Fira. The hyper-parameter γ is set to 1.01 across all experiments, which consistently yields satisfactory results. Sensitivity analyses of hyper-parameter γ to training performance are provided in Appendix H.

5 Experiments

In this section, we validate the effectiveness of Fira in pre-training and fine-tuning tasks of LLMs. In our experiments, we denote our method using the strategy of Eqn. 10 as Fira-matrix, and denote our method additionally using the column-wise strategy of Eqn. 11 as Fira.

5.1 Memory-Efficient Pre-training

Experimental setup. We follow the settings in Galore [43] to conduct the pre-training experiments. We compare Fira with GaLore [43], LoRA [13], ReLoRA [18], and full-rank training baselines. Adam optimizer is used to train all baselines and our method on the C4 dataset in the BF16 format. The settings of these baselines can be found in [43]. The dataset C4 is a colossal, cleaned version of Common Crawl's web crawl corpus, which is widely used in LLM pre-training [26]. Following [43], we utilize LLaMA-based architectures equipped with RMSNorm and SwiGLU activations [37, 29, 31]. As in [43], our training protocol excludes data repetition and spans a sufficiently large

dataset, encompassing a diverse array of model sizes (60M, 130M, 350M, 1B). To guarantee a fair comparison, we employ the same learning rate 0 . 01 as used in GaLore and maintain the same rank r for each model size. The detailed settings of pre-training are provided in Appendix B.1. We use 8 A100 80G GPUs to conduct pre-training experiments.

Result analysis. As shown in Table 3, Fira consistently outperforms low-rank training baselines by a large margin under the same rank constraint, and even surpasses full-rank training. Following [43], we estimate the memory reduction of the optimizer states via the same memory estimation method introduced in [43]. Detailed memory estimation methods are presented in C.1. Additional experiments on real memory usage and throughput are presented in C.2. From Table 3, our Fira saves 61.1% memory usage of the optimizer states when pre-training the LLaMA 1B architecture compared to full-rank training, while Fira achieves better results. Compared to full-rank training, Fira's superior performance may be attributed to the following reason: the gradient direction in the norm-based scaling method is determined by the current state, rather than by historical gradients in Adam. Therefore, Fira introduces a higher degree of randomness in training, which can enhance the model's ability to escape the local optima, leading to better training performance [44]. It's worth noting that previous work [42, 45] also reveals this phenomenon. Detailed analysis of this phenomenon is presented in Appendix F.

5.2 Scaling up to LLaMA 7B Pre-training

To validate the scalability of our method, we scale up by pre-training the LLaMA 7B model with the full-rank dimension d model = 4096 . We compare Fira with the GaLore baseline, which generally achieves the best performance among low-rank training baselines, as shown in Table 3. As illustrated in Figure 4, our method demonstrates a significant improvement over GaLore for pre-training LLaMA 7B, while using an 8 × smaller rank (memory of optimizer states). This highlights Fira's effectiveness, suggesting it could be a viable solution for large-scale LLM pretraining.

5.3 Memory-Efficient Fine-Tuning

Experimental setup. Following [14], we perform the fine-tuning task to compare Fira with LoRA, GaLore, Flora [12], ReLoRA, Full-rank training, on the LLaMA-7B model for commonsense reasoning tasks. This task consists of eight sub-tasks, each with its own designated training and testing sets. Following the approach of [14], we combine the training datasets from all eight sub-tasks into a unified training set, while evaluating each sub-task individually using its respective testing dataset. In the fine-tuning task, the rank r is set to 32 and the learning rate is set to 1e-4. The detailed



































settings of fine-tuning are provided in Appendix B.2. We adopt RTX 4090 GPUs for fine-tuning experiments.

Result analysis. As shown in Table 4, our Fira achieves the highest performance on 5 out of 8 datasets, demonstrating better or comparable performance compared to the baseline methods. Notably, GaLore struggles to adapt to the HellaSwag and WinoGrande datasets, resulting in a significant decline in scores. In contrast, our Fira adapts to these tasks well and achieves the highest scores on WinoGrande. In terms of memory efficiency, our method uses comparable or even less memory than the low-rank training methods LoRA and GaLore. These results illustrate the effectiveness of our method for the fine-tuning of LLMs.

5.4 Ablation Study

In this section, we conduct an ablation study to assess the effectiveness of each component in our method. We adopt the same settings in Section 5.1 for pre-training the LLaMA 60M model. We design six variants of our method for the ablation study: (1) Fira-w.o.-scaling: our Fira without using the scaling factor to correct the gradient (i.e., setting ϕ t ( R t ) to a fixed value of 1). (2) Fira-matrix: our Fira using the scaling factor at the matrix level instead of at the column level. (3) Fira-w.o.-limiter: our Fira without using norm-growth limiter to avoid

training loss spikes. (4) Fira-gradient-clipping: our Fira using gradient clipping [25] to avoid loss spikes instead of our proposed norm-growth limiter. (5) Fira-gradient-shrink: our Fira using gradient shrink [36] instead of our proposed norm-growth limiter. (6) Fira-tensor-wise-scaling: our Fira using tensor-wise scaling [8] instead of our proposed norm-growth limiter.

Table 5 presents the results. It can be found that Fira outperforms Fira-w.o.-scaling, thereby demonstrating the effectiveness of our proposed norm-based scaling method for gradient correction. This also suggests that directly incorporating the raw gradient outside the subspace without correction will lead to sub-optimal results. Besides, Fira yields better performance than Fira-matrix, illustrating that a more fine-grained consideration of the scaling factor is beneficial. Furthermore, Fira demonstrates improved performance over Fira-w.o.-limiter, Fira-gradient-clipping, Fira-gradient-shrink, and Firatensor-wise-scaling, indicating the effectiveness of our proposed norm-growth limiter in addressing the issue of training loss spikes.

5.5 Performance under Varying Ranks

In this section, we illustrate the advantages of our Fira over Galore under a lower rank. We adjust various rank configurations within the set { 4 , 16 , 64 , 128 } and d model = 256 , and then assess the performance of pre-training the LLaMA 60M model on the C4 dataset as outlined in Section 5.1. The validation perplexity of Fira and GaLore after 10K steps across different ranks is depicted in Figure 5. From Figure 5, we can observe that Fira consistently surpasses GaLore across all rank configurations. Notably, even when the ranks are set very low (4 and 16), Fira still achieves performance comparable



































to full-rank training. In contrast, the performance of GaLore significantly declines in these cases. These results highlight the superiority of our proposed Fira at lower ranks and its effectiveness in reducing memory usage.

6 Conclusion

In this paper, we present a plug-and-play memory-efficient training framework for LLMs, called Fira, as the first attempt to facilitate full-rank training consistently under low-rank constraints. First, we find a notable phenomenon in LLM training: the scaling effect of adaptive optimizers on the gradient norm remains similar between low-rank and full-rank training. Building on this observation, we propose a norm-based scaling method that applies the scaling effect of low-rank optimizers in place of full-rank optimizers to facilitate full-rank training. This allows us to maintain the low-rank constraint within the optimizer while still benefiting from the advantages of full-rank training for improved performance. Additionally, we observe sudden spikes in gradient values during optimization, which can result in corresponding spikes in the training loss. To mitigate this, we propose a norm-growth limiter that smooths gradients by regulating the relative increase in gradient norms. Extensive experiments in both pre-training and fine-tuning of LLMs demonstrate the effectiveness of our proposed Fira.

Acknowledgments and Disclosure of Funding

This work was supported by the NSFC under Grants U2441242.

References

Hervé Abdi. The kendall rank correlation coefficient. Encyclopedia of measurement and statistics , 2:508-510, 2007.

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

David Balduzzi, Marcus Frean, Lennox Leary, JP Lewis, Kurt Wan-Duo Ma, and Brian McWilliams. The shattered gradients problem: If resnets are the answer, then what is the question? In International conference on machine learning , pages 342-350. PMLR, 2017.

Han Chen, Garvesh Raskutti, and Ming Yuan. Non-convex projected gradient descent for generalized low-rank tensor regression. Journal of Machine Learning Research , 20(5):1-37, 2019.

Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 , 2016.

Yudong Chen and Martin J Wainwright. Fast low-rank estimation by projected gradient descent: General statistical and algorithmic guarantees. arXiv preprint arXiv:1509.03025 , 2015.

Israel Cohen, Yiteng Huang, Jingdong Chen, Jacob Benesty, Jacob Benesty, Jingdong Chen, Yiteng Huang, and Israel Cohen. Pearson correlation coefficient. Noise reduction in speech processing , pages 1-4, 2009.

Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. In International Conference on Learning Representations , 2022.

Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized llms. Advances in Neural Information Processing Systems , 36, 2024.

Kaituo Feng, Changsheng Li, Xiaolu Zhang, JUN ZHOU, Ye Yuan, and Guoren Wang. Keypointbased progressive chain-of-thought distillation for llms. In International Conference on Machine Learning , 2024.

Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning , volume 1. MIT Press, 2016.

Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compressors. In Forty-first International Conference on Machine Learning , 2024.

Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022.

Zhiqiang Hu, Lei Wang, Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Ka-Wei Lee. Llm-adapters: An adapter family for parameter-efficient fine-tuning of large language models. In The 2023 Conference on Empirical Methods in Natural Language Processing , 2023.

Dayal Singh Kalra, John Kirchenbauer, Maissam Barkeshli, and Tom Goldstein. When can you get away with low memory adam? arXiv preprint arXiv:2503.01843 , 2025.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Teven Le Scao, Thomas Wang, Daniel Hesslow, Stas Bekman, M Saiful Bari, Stella Biderman, Hady Elsahar, Niklas Muennighoff, Jason Phang, Ofir Press, et al. What language model to train if you have one million gpu hours? In Findings of the Association for Computational Linguistics: EMNLP 2022 , pages 765-782, 2022.

Vladislav Lialin, Sherin Muckatira, Namrata Shivagunde, and Anna Rumshisky. Relora: Highrank training through low-rank updates. In The Twelfth International Conference on Learning Representations , 2024.

Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations , 2019.

Xu-Hui Liu, Yali Du, Jun Wang, and Yang Yu. On the optimization landscape of low rank adaptation methods for large language models. In The Thirteenth International Conference on Learning Representations , 2025.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2019.

Kai Lv, Hang Yan, Qipeng Guo, Haijun Lv, and Xipeng Qiu. Adalomo: Low-memory optimization with adaptive learning rate. arXiv preprint arXiv:2310.10195 , 2023.

Ionut-Vlad Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurti´ c, Thomas Robert, Peter Richtárik, and Dan Alistarh. Microadam: Accurate adaptive optimization with low space overhead and provable convergence. Advances in Neural Information Processing Systems , 37:1-43, 2024.

Robert Nau. Forecasting with moving averages. Fuqua School of Business, Duke University , pages 1-3, 2014.

Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning , pages 1310-1318. Pmlr, 2013.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research , 21(140):1-67, 2020.

Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. { Zero-offload } : Democratizing { billion-scale } model training. In 2021 USENIX Annual Technical Conference (USENIX ATC 21) , pages 551-564, 2021.

Philip Sedgwick. Spearman's rank correlation coefficient. Bmj , 349, 2014.

Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 , 2020.

Chonghao Sima, Katrin Renz, Kashyap Chitta, Li Chen, Hanxue Zhang, Chengen Xie, Ping Luo, Andreas Geiger, and Hongyang Li. Drivelm: Driving with graph visual question answering. arXiv preprint arXiv:2312.14150 , 2023.

Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

Mojtaba Valipour, Mehdi Rezagholizadeh, Ivan Kobyzev, and Ali Ghodsi. Dylora: Parameter efficient tuning of pre-trained models using dynamic search-free low-rank adaptation. arXiv preprint arXiv:2210.07558 , 2022.

Yeming Wen and Swarat Chaudhuri. Batched low-rank adaptation of foundation models. arXiv preprint arXiv:2312.05677 , 2023.

Kirk M Wolter and Kirk M Wolter. Introduction to variance estimation , volume 53. Springer, 2007.

Wenhan Xia, Chengwei Qin, and Elad Hazan. Chain of lora: Efficient fine-tuning of language models via residual learning. arXiv preprint arXiv:2401.04151 , 2024.

Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang, Hanyu Lai, Ming Ding, Zhuoyi Yang, Yifan Xu, Wendi Zheng, Xiao Xia, Weng Lam Tam, Zixuan Ma, Yufei Xue, Jidong Zhai, Wenguang Chen, Zhiyuan Liu, Peng Zhang, Yuxiao Dong, and Jie Tang. GLM-130b: An open bilingual pre-trained model. In The Eleventh International Conference on Learning Representations , 2023.

Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems , 32, 2019.

Haoyang Zhang, Yirui Zhou, Yuqi Xue, Yiqi Liu, and Jian Huang. G10: Enabling an efficient unified gpu memory and storage architecture with smart tensor migrations. In Proceedings of the 56th Annual IEEE/ACM International Symposium on Microarchitecture , pages 395-410, 2023.

Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In International Conference on Learning Representations , 2020.

Longteng Zhang, Lin Zhang, Shaohuai Shi, Xiaowen Chu, and Bo Li. Lora-fa: Memory-efficient low-rank adaptation for large language models fine-tuning. arXiv preprint arXiv:2308.03303 , 2023.

Qingru Zhang, Minshuo Chen, Alexander Bukharin, Nikos Karampatziakis, Pengcheng He, Yu Cheng, Weizhu Chen, and Tuo Zhao. Adalora: Adaptive budget allocation for parameterefficient fine-tuning. arXiv preprint arXiv:2303.10512 , 2023.

Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. arXiv preprint arXiv:2406.16793 , 2024.

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. In International Conference on Machine Learning , 2024.

Pan Zhou, Jiashi Feng, Chao Ma, Caiming Xiong, Steven Chu Hong Hoi, et al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. Advances in Neural Information Processing Systems , 33:21285-21296, 2020.

Hanqing Zhu, Zhenyu Zhang, Wenyan Cong, Xi Liu, Sem Park, Vikas Chandra, Bo Long, David Z Pan, Zhangyang Wang, and Jinwon Lee. Apollo: Sgd-like memory, adamw-level performance. arXiv preprint arXiv:2412.05270 , 2024.

Philip Zmushko, Aleksandr Beznosikov, Martin Takáˇ c, and Samuel Horváth. Frugal: Memoryefficient optimization by reducing state overhead for scalable training. arXiv preprint arXiv:2411.07837 , 2024.

A Fira Implementation

A.1 Algorithm Pseudocode

A.2 Plug-and-play Framework for Fira

B Details of Experiments

B.1 Detailed Pre-Training Setting

This section provides an overview of the LLaMA architectures and the hyper-parameters employed during pre-training. To ensure a fair comparison, we adopt the same settings as [43]. Table 6 presents the hyper-parameters of the LLaMA architectures across various sizes. For all architectures, we utilize a maximum sequence length of 256 and a batch size of 131K tokens. Furthermore, we implement a learning rate warm-up during the initial 10% of training steps and employ cosine annealing for the learning rate schedule, which decreases to 10% of the initial learning rate.

For all methods except Fira and GaLore, we tune the optimal learning rate from the set {0.01, 0.005, 0.001, 0.0005, 0.0001} across model sizes ranging from 60M to 1B, selecting their best validation

perplexity to report. In contrast, both Fira and GaLore employ the same learning rate 0.01 and a subspace change frequency T of 200 without tuning. Additionally, the scale factor α is considered a fractional learning rate [43]. Furthermore, a relatively large learning rate may result in spikes of training loss [43]. To address this issue, for models with a size of less than 1B, we set α to 0.25, while for models exceeding 1B, we adjust α to 0.0625.

B.2 Detailed Fine-tuning Setting

Wefine-tune the pre-trained LLaMA-7B model for commonsense reasoning tasks benchmark designed for LLM fine-tuning, which include eight sub-tasks [14]. Table 7 shows the hyper-parameter configurations.

C Details of Overhead Comparison

C.1 Memory Estimates

Due to the difficulties associated with directly measuring GPU memory usage for a specific component, we estimate the memory requirements for weight parameters and optimizer states across various methods and model sizes, following GaLore [43]. This estimate is derived from the number of parameters and optimizer states in BF16 format. In particular, for the memory of parameters, we multiply the total number of parameters by 2; for the memory of optimizer states and gradients, we first calculate the total number of them according to Table 2 and then multiply this total number by 2.

C.2 Additional Experiments on Real Overhead

Weconduct additional comparisons regarding real memory usage and throughput of different memoryefficient training methods for both pre-training and fine-tuning. As illustrated in Tables 8 and 9, Fira achieves superior memory efficiency compared to full-rank training without significantly reducing throughput. Although Fira's throughput is slightly lower than that of other memory-efficient methods, it delivers exceptional performance. During pre-training, methods like LoRA necessitate maintaining additional higher-rank adapters (e.g. r/d model = 512 / 2048 ) to achieve performance comparable

to full-rank training. However, in practice, maintaining these higher-rank adapters outweighs the benefits of fewer trainable parameters compared to full-rank training. This may be due to the fact that higher-rank adapters significantly increase computational demands during pre-training compared to fine-tuning (e.g. r/d model = 32 / 2048 ), thus leading to more memory and less throughput. Furthermore, since full fine-tuning of LLaMA 7B's memory requirements exceeds the A100's 80GB capacity, we utilize DeepSeed's Zero2 technology to mitigate its memory usage.

D Additional Quantitative Analysis of Scaling Factor Similarity

In this section, we analyze the similarities of scaling factors between low-rank and full-rank training at the matrix and column levels. Specifically, when analyzing similarity, we will form a vector of values for different scaling factors ( ϕ 1 , ϕ 2 , . . . , ϕ k ) ( ϕ i = ∥ ψ ( R ( i ) t ) ∥ ∥ R ( i ) t ∥ ). For the matrix level, we take the different weight matrices of the model as items (i.e., R ( i ) t ∈ R r × n ); for the column level, we take the columns in the weight matrix as items (i.e., R ( i ) t ∈ R r × 1 ). We use the same low-rank setup as in Table 3. The models of different training runs are trained from the same random initialization. The only difference between them is the value of the rank hyperparameter (low-rank vs. full-rank).

D.1 Spearman, Kendall, and Pearson Correlation Coefficients

In this section, we will employ Kendall's Tau correlation coefficient [1] and Spearman's rank correlation coefficient [28], Pearson's correlation coefficient [7] to evaluate the similarities of scaling factors between low-rank and full-rank training. We conduct our experiment by including all matrices of LLaMA models ranging from 60M to 1B. Then, we train these models and assess the similarity of scaling factors averaged over 10,000 steps. Additionally, to evaluate the effectiveness of a columnlevel fine-grained strategy for scaling factors, we perform a column-level quantitative similarity analysis. Due to the computational challenges posed by the large number of columns, we randomly sample 100 columns for each weight matrix for analysis. Specifically, in the LLaMA 1B model, over 10,000 columns are sampled.

Both Spearman, Kendall, and Pearson correlation coefficients range from -1 to +1. A coefficient of 1 signifies a perfect positive correlation, and -1 signifies a perfect negative correlation. The p-value helps us determine whether the observed correlation is statistically significant or if it could have occurred by random chance. For instance, a p-value less than 0.05 means there is less than a 5% probability that the observed correlation happened by chance if there was actually no correlation. Generally, a p-value below 0.05 suggests that a significant correlation exists. As shown in Table 10 , we can observe the significant similarity of scaling factors between low-rank and full-rank LLM training (all coefficients close to 1, while p-values close to 0). Thus, it is likely that the observed behavior is an inherent feature of LLM training, manifesting across a broad range of scenarios. This insight provides a robust experimental basis for our proposed norm-based scaling in Fira and helps explain its effectiveness.

D.2 Similarity Trends.

In this section, we conduct additional experiments on similarity trends of scaling factors through Cosine Similarity and Euclidean Distance using two rank settings: r/d model = 16 / 256 and r/d model = 128 / 256 .

































t

(

Rt

)



t

(

Gt

)

t

(

Rt

)



t

(

Gt

PtRt

)

t

(

Gt

)



t

(

Gt

PtRt

)





























t

(

Rt

)



t

(

Gt

)

t

(

Rt

)



t

(

Gt

PtRt

)

t

(

Gt

)



t

(

Gt

PtRt

)





























t

(

Rt

)



t

(

Gt

)

t

(

Rt

)



t

(

Gt

PtRt

)

t

(

Gt

)



t

(

Gt

PtRt

)



























t

(

Rt

)



t

(

Gt

)

t

(

Rt

)



t

(

Gt

PtRt

)

t

(

Gt

)



t

(

Gt

PtRt

)

As shown in Figure 6, 7, 8, and 9, the similarity exhibits fluctuations during the initial training phase but achieves a relatively steady pattern with high similarity in the later iterations. Under lower rank setting r/d model = 16 / 256 , there is negligible reduction in similarity. Besides, ϕ t ( G t ) and ϕ t ( G t -P t R t ) demonstrate significantly higher similarity owing to their closely aligned dimensions.

E Theoretical Analysis

E.1 Error Upper Bound for the Approximation of Scaling Factors

In Section 4.1, we use the scaling factors of low-rank gradients to approximate that of full-rank gradients. To quantify the effectiveness of this approximation, we will derive its error upper bound theoretically, and verify our analysis experimentally. The error of the approximation κ ( r ) can be written as:

where rank r ≤ n , and R t = P ⊤ t G t ∈ R r × n . To simplify the proof, we consider that r components ( g 1 , . . . , g r ) of low-rank gradients are directly sampled from n components ( g 1 , . . . , g n ) of full-rank gradients. Under these conditions, the error can be rewritten as:

Assumption 1. (Bounded Scaling Factors): we assume that the adaptive optimizer scales each gradient component g i by the scaling factor that lies within known bounds. Specifically, there exist constants c min and c max such that for all i :

Theorem 1. (Error Upper Bound for Approximation) Under the assumption that ψ 2 i ( g i ) g 2 i are bounded between constants c 2 min and c 2 max for all components i , the approximation error κ ( r ) satisfies:

̸

where ∆ ϕ 2 ≜ sup g =0 ψ 2 ( g ) g 2 -inf g =0 ψ 2 ( g ) g 2 = ( c 2 max -c 2 min ) defines the squared scaling factor variation, ∥ G ∥ = ∑ n i =1 g 2 i , and ∥ G r ∥ = ∑ r i =1 g 2 i .

̸

Proof. We first define the following quantities: Total Gradient Norm s n = ∑ n i =1 g 2 i , Partial Gradient Norm (first r components) s r = ∑ r i =1 g 2 i , Remaining Gradient Norm s n -r = s n -s r = ∑ n i = r +1 g 2 i , Total Corrected Gradient Norm S n = ∑ n i =1 ψ 2 i ( g i ) , Partial Adjusted Gradient Norm (first r components) S r = ∑ r i =1 ψ 2 i ( g i ) , Remaining Adjusted Gradient Norm S n -r = S n -S r = ∑ n i = r +1 ψ 2 i ( g i ) .

Then, our estimation error κ ( r ) can be rewritten using these definitions:

First, we rewrite the estimation error κ ( r ) in terms of S r , S n -r , s r , and s n -r :

Then, compute the difference in the numerator:

This implies:

Simplify the numerator, thus, the estimation error becomes:

From our bounded assumption, we have:

Therefore, the maximum possible difference between S n -r s n -r and S r s r is:

In addition, we have:

Finally, the approximation error κ ( r ) is bounded above by:

From this theory, we can find that the error upper bound on the approximation of scaling factors is mainly determined by two aspects, and we can verify them experimentally:

Variability of Scaling Factor ∆ ϕ 2 : This term represents the maximum variation in the scaling factors of different gradient components. For further validation, we designed Fira-onlyscaling , a variant of Fira. It directly applies the low-rank scaling factors to the full-rank gradients by changing the Eqn. 10 from W t +1 = W t -ηP t ψ t ( R t ) -ηϕ t ( R t )( G t -P t R t ) to W t +1 = W t -ηϕ t ( R t ) G t . In this way, we are able to exclude the influence of the original Adam term P t ψ t ( R t ) and better analyze the effectiveness of our approximation. As shown in Table 11, Fira-only-scaling (column-level) gains better performance than Fira-only-scaling (matrix-level) for its more fine-grained consideration of the scaling factor, which also means a smaller maximum variation ∆ ϕ 2 .

Effectiveness of Gradient Sampling ( 1 -∥ G r ∥ ∥ G ∥ ) : This term represents the proportion of the gradients norm contributed by the sampled low-rank r components from full-rank n components. As shown in Table 12, we conducted ablation experiments Fira-only-scalingw.o.-svd , i.e., Fira-only-scaling without SVD in low-rank gradient sampling. As we can see, SVD is capable of sampling more prominent low-rank gradients, which leads to a reduction in the upper bound of error and enhanced performance. Similarly, as shown in Table 13, employing a higher rank enables the sampling of a greater proportion of the gradients norm, resulting in reduced error upper bound and improved performance.

Table 11: Ablation on the level of scaling factors for the variant Fira-only-scaling.

After that, we factor out s n -r :

E.2 Variance of Scaling Factors.

The variance of adaptive learning rates is significantly elevated during the early stage of training, often necessitating a warm-up to mitigate this variance and stabilize training [19]. As illustrated in Figure 10, the scaling factor in Fira exhibits a similar pattern, characterized by substantial variance during the early stage of training, which also necessitates a warm-up. However, the addition of an extra warm-up hyper-parameter for Fira would be inefficient. Therefore, it is crucial to investigate whether the original warm-up would have efficiently mitigated the variance in Fira. In the subsequent theoretical analysis, we show that, during the early training phase, the variance of the scaling factor of Fira is less than or equal to that of the adaptive learning rate. This finding suggests that the existing warm-up strategy is sufficient to mitigate the variance of Fira, thereby eliminating the need for an additional warm-up hyper-parameter.

Consider independent random vectors { g ( i ) } n i =1 , where each g ( i ) = ( g ( i ) 1 , g ( i ) 2 , . . . , g ( i ) t ) . Here, the superscript i indicates the index of the weight matrix to which the vector belongs, while the subscript j (where j ranges from 1 to t ) denotes training iterations with each parameter. Following [19], we assume the adaptive learning rate of Adam ψ ( . ) = √ 1 -β t 2 (1 -β 2 ) ∑ t i =1 β t -i 2 g 2 i , and g ( i ) j ∼ N (0 , σ 2 ) for all i and j in the early stage. Additionally, approximate the distribution of the exponential moving average as the distribution of the simple average, p ( ψ ( . )) = p ( √ 1 -β t 2 (1 -β 2 ) ∑ t i =1 β t -i 2 g 2 i ) ≈ p ( √ t ∑ t i =1 g 2 i ) [24], and then ψ 2 ( . ) ∼ Scale-invX 2 ( ρ, 1 σ 2 ) .

Theorem 2. (Variance of Scaling Factors) In the early stages of training, if ψ 2 ( · ) ∼ Scale-invX 2 ( ρ, 1 σ 2 ) , and g ( i ) j ∼ N (0 , σ 2 ) 5 for all i, j , then for all ρ > 4 , the scaling factor ϕ 2 = ∑ n i =1 ψ 2 i ( g ( i ) t ) 2 ∑ n i =1 ( g ( i ) t ) 2 satisfies Var[ ϕ 2 ] ≤ Var[ ψ 2 ] . If we approximate √ ψ 2 and √ ϕ 2 to the first order, we have Var[ ϕ ] ≤ Var[ ψ ] .

Proof. We express ϕ 2 as a weighted sum:

where the weights are defined as:

Each w i is a non-negative random variable satisfying ∑ n i =1 w i = 1 .

In the context of adaptive optimization algorithms like Adam, the squared gradients ψ 2 i accumulate information from past iterations to adapt the learning rate for each parameter. With β 2 = 0 . 999 ,

























































the moving average of the squared gradients places significant weight on historical data, making ψ 2 i dependent mainly on past gradients, yielding:

Since ψ 2 i primarily depend on past gradients g ( i ) 1 , . . . , g ( i ) t -1 , and w i depend solely on the current gradients g ( i ) t , we can consider ψ 2 i and w i to be independent random variables. Consequently, we can express the variance of ϕ 2 as:

Using the law of total variance, we have:

Since ψ 2 i are independent of the w i , we find:

Thus, the variance simplifies to:

n is constant.

The second term, Var ( E [∑ i =1 w i ψ 2 i | w i ]) , is zero since E[ ϕ 2 | w i ] = E[ ψ 2 i ] Let X i = ( g ( i ) t ) 2 , where each X i ∼ σ 2 χ 2 1 . Then, we can express the weights as:

Since X i /σ 2 ∼ χ 2 1 , and each w i is the ratio of X i to the sum of all X j , the vector ( w 1 , . . . , w n ) follows a Dirichlet distribution with parameters α i = ν i 2 = 1 2 , where ν i = 1 is the degrees of freedom of χ 2 1 .

For a Dirichlet distribution, the expected value of w 2 i is given by:

Substituting α i = 1 2 and ∑ n k =1 α k = n 2 yields:

Thus, summing over all i gives:

Since n ≥ 1 , it follows that:

which implies:

Finally, substituting this result back into the variance expression:

Given ρ > 4 and ψ 2 ( · ) ∼ Scale-invX 2 ( ρ, 1 σ 2 ) , the variance of ψ 2 ( · ) exists [19].

Since ψ 2 i and w i are independent and ∑ n i =1 E[ w i ] = 1 :

Thus, we have shown that:

which implies:

Follow [19], we approximate √ ψ 2 and √ ϕ 2 to the first order [34]

To further examine our theorem, we conduct simulations to calculate the variance of the scaling factor ϕ at ranks within the set { 1 , 5 , 10 , 50 , 100 } . The adaptive learning rate ψ is equivalent to that of ϕ when the rank equals 1. As shown in Figure 11, the variance decreases as the rank increases, supporting our above theorem Var[ ϕ ] ≤ Var[ ψ ] . Furthermore, we observe a surprisingly large variance during the early stage, which corroborated our initial experiments. Consequently, we conclude that our method is efficient without requiring an additional warm-up.

F Potential Reasons for Fira's Enhanced Performance Compared to Full-Rank Optimizer

While our proposed method aims to approximate full-rank training, one might intuitively assume that the performance of full-rank Adam represents the upper bound. However, in previous experiments, we noticed that Fira may achieve performance that is comparable to or even better than full-rank training. This raises the question of why Fira outperforms full-rank Adam in certain scenarios. In this section, we first demonstrate that, despite its goal of approximating full-rank training, Fira is fundamentally

distinct from full-rank Adam. We then further substantiate this claim through experimental evidence. Finally, we explore potential reasons why Fira can achieve enhanced performance compared to full-rank Adam. Notably, the phenomenon where approximate optimizer correction outperforms the original full-rank optimizer has also been observed in prior work, such as [42, 45], consistent with our findings in Fira.

The approximation of low-rank to full-rank training in Fira is limited to scaling factors. Even in the case of full-rank training, Fira is not equivalent to the original Adam. To illustrate this, we introduce a variant of Fira called Fira-only-scaling in Appendix E.1. To isolate the influence of the original Adam term P t ψ t ( R t ) , Fira-only-scaling directly applies low-rank scaling factors to the full-rank gradients by modifying Eqn. 10 from:

As shown in Table 13, higher-rank approximations of scaling factors lead to improved performance, validating the effectiveness of our approximation. But we can also find that full-rank Fira-only-scaling and Adam are not the same. Fira-only-scaling employs a weight-matrix-wise scaling factors ϕ t ( R t ) to adjust the raw gradients, whereas Adam uses element-wise adaptive learning rates.

As a result, in contrast to full-rank Adam, which is parameter-level adaptive, Fira applies adaptive strategy at matrix-level (column-level), while maintaining the original gradient direction within a matrix (column) (i.e., ϕ t ( R t ) G t ). As a result, the gradient direction which is only determined by the current state can introduce a higher degree of randomness in training. Then, the randomness in training can enhance the model's ability to escape the local optima, thus leading to better generalization performance. This insight potentially explains why Fira could match or even surpass the full-rank Adam baseline.

To further argue this point, we conduct additional experiments on the comparisons of perplexity trends. We compare the perplexity trends of Fira, Fira-only-scaling, SGD, and Adam. As illustrated in 12 (a), the performance of vanilla SGD is significantly inferior, highlighting its inadequacy for directly training LLMs. As depicted in 12 (b) and (c), while Adam demonstrates faster convergence during the initial stages, both Fira and Fira-only-scaling achieve superior performance in the later stages. This may occur because Fira applies an adaptive strategy exclusively at the matrix level while preserving the original gradient direction within each weight matrix. As previously analyzed, Fira may introduce a higher degree of randomness in training and a better ability to escape the local optima, and then enable better generalization performance.

























(a) Fira v.s. SGD

















  



  



  







(b) Fira v.s. Adam

















  



  



  







(c) Fira-only-scaling v.s. Adam

G Additional Analysis of Spikes

Maintaining the direction of the raw gradient without correction might be unable to effectively deal with the steep loss landscapes of LLM training like Adam [39]. The steep loss landscapes are likely to cause abrupt increases in raw gradients. When the raw gradients increase abruptly, the gradients' norm after norm-based scaling may also increase abruptly, as illustrated in Figure 3. This arises from to:











































































the fact that the norm-based scaling method only adjusts the average gradient norm of the gradient at the matrix level, failing to make fine-grained adjustments to each parameter, unlike the optimizer Adam. As a result, a significant parameter update may occur, undermining previous optimization efforts, i.e. training loss spikes [11, 39]. As shown in Figure 13, when we directly use Adam to pre-train the LLaMA model, there will be no loss spike. Our norm-growth limiter is mainly aimed at addressing the gradient stability capability that our norm-based scaling method lacks compared to Adam.

For more comprehensive comparisons of our norm-growth limiter, we design two additional gradient stabilization variants to solve the loss spike: Gradient Shrink [36] ( | S t | = | S t | · µ + | S t -1 | · (1 -µ ) ), and Tensor-Wise Scaling [8] ( | S t | = | S t | · µ ), where S t = ϕ t ( R t )( G t -P t R t ) is the corrected gradient by applying our norm-based scaling. As shown in Figure 14 and Table 14, Fira outperforms other gradient stabilization methods. For further analysis, Gradient Shrink fails to solve the loss spike, while Tensor-Wise Scaling solves the loss spike but leads to sub-optimal results.

H Sensitivity Analyses of γ

As shown in Table 15, the performance of Fira is not sensitive to the choice of γ . Provided that γ ≤ 1 . 01 , Fira can effectively mitigate spikes in loss, with only a marginal decrease in performance when γ = 1 . As we can see, the setting of γ = 1 . 01 , employed across all experiments in this paper, is highly effective. This value is neither too small to restrict the normal growth of the gradient nor too large to fail to limit sudden increases in gradient magnitude.

I Discussion on Concurrent Works

Adam-mini [42] employs block-wise learning rates to replace Adam's element-wise learning rate based on the Hessian structure of neural networks. However, it only reduces memory usage for the second-order momentum while keeping the first-order momentum unchanged. In contrast, Fira reduces memory for both first-order and second-order momentum.

SlimAdam [15] uses layer-wise SNR analysis to replace second-moment tensors with their means for memory efficiency. However, like Adam-mini, it only focuses on the second-order momentum and is unable to reduce memory for both first-order and second-order momentum like Fira.

MicroAdam [23] compresses gradient information to save memory with theoretical convergence guarantees, but it is designed for fine-tuning workloads and cannot be adapted to LLM pre-training. However, Fira works for both LLM pre-training and fine-tuning.

AdaLomo [22] enhances low-memory optimization (LOMO) with adaptive learning rates to improve robustness and convergence. Unlike Fira, which focuses on enabling full-rank training under low-rank constraints, AdaLomo focuses on improving LOMO's convergence without involving the low-rank constraints.

FRUGAL [46] addresses gradient information loss outside the subspace in GaLore by directly adding these gradients outside the subspace via SGD or signSGD, without exploring more effective ways to utilize such gradients. In contrast, Fira explores the use of scaling factors to enable more effective updates for gradients outside the subspace. Additionally, FRUGAL provide theoretical convergence guarantees when using SGDM for subspace updates and SGD for updates outside the subspace.

GaRare [20] analyzes and compares the loss landscapes of low-rank and full-rank setting under gradient projection of GaLore, demonstrates that low-rank constraints improve the optimization landscape by avoiding sharp minima. This may help explain Fira's superior performance. Although GaLore's low-rank constraints enable a better optimization landscape, they discard significant gradient information outside the subspace, leading to suboptimal results. In contrast, Fira leverages low-rank constraints for a better optimization landscape while effectively capturing gradient information outside the subspace, resulting in superior performance.

J Additional Experimental Comparison

To further demonstrate the generality of our method, we conduct additional experiments on a wider range of architectures (Gemma-7B, Mistral-7B), datasets (MMLU), and adaptive optimizer (Adagrad).

As shown in Table 16 and 17, Fira remains competitive with all baselines across different model architectures, datasets, and optimizers, confirming its broad applicability.

To provide a comprehensive comparison against a broader range of baseline methods, we further extend our evaluation under the same experimental settings as in Tables 3 and 4.

As shown in Table 18, Fira demonstrates consistent superiority over these baselines when pretraining LLaMA models of different sizes. As shown in Table 19, when fine-tuning LLaMA 7B on commonsense reasoning datasets, Fira achieves the highest performance across most datasets, exhibiting superior average accuracy compared to these baselines.

K Robustness of Norm-based Scaling

As discussed in Adam-mini [42] and APOLLO [45], Adam's element-wise scaling contains significant redundancy. Using a shared scaling for a group of parameters, such as column-wise scaling in Firaonly-scaling, is sufficient to replace Adam's element-wise scaling. In this scenario, by aggregating historical gradient information across multiple parameters, column-wise norm-based scaling of Fira may be more robust than Adam's element-wise scaling when dealing with training data noise.

To empirically assess this robustness, we inject Gaussian noise into the gradients during pre-training and measure the resulting performance degradation.

As shown in Table 20, Fira-only-scaling exhibits slight performance degradation, while Adam's element-wise scaling shows significant performance degradation. This indicates that column-wise scaling is more robust.

L Limitations

This work presents Fira, a novel memory-efficient training framework that enables full-rank training of Large Language Models (LLMs) under low-rank constraints, significantly reducing memory usage while maintaining high performance. However, our current research primarily focuses on LLMs. In future work, we aim to extend the applicability of Fira to other domains, such as Multimodal Large Language Models (MLLMs) and diffusion models, to further broaden its impact across diverse machine learning tasks and architectures.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The abstract and introduction clearly reflect the paper's primary contributions and scope, including the development and evaluation of our proposed method.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The paper offers detailed information required to replicate the primary experimental results, promoting transparency and reproducibility of the findings.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code will be released publicly.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We detail all training and testing procedures in Section 5 and the Appendix B. Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Owing to the high training costs, we did not conduct multiple runs to calculate error bars.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We detail the information on the computer resources in the Section 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research in the paper fully complies with the NeurIPS Code of Ethics, with all aspects of our work carefully reviewed to meet the established ethical standards.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [No]

Justification: The paper has no negative societal impacts.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper poses no such risks, because no data or models were released.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have appropriately credited and referenced the creators or original owners of assets, such as code, data and models, used in our work. And, we have explicitly mentioned and properly respected the licenses and terms of use.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: This paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not include crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: This paper does not include crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces **Fira**, a training framework that enables **full-rank training of LLMs under low-rank memory constraints**. Unlike LoRA and GaLore that restrict training to low-rank subspaces and degrade performance, Fira uses a **norm-based scaling method** to approximate full-rank optimizer behavior using low-rank gradients. It also introduces a **norm-growth limiter** to stabilize training. Experiments show Fira consistently outperforms existing low-rank methods in both pre-training and fine-tuning, with significantly lower memory usage.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper’s central insight—that the scaling factors of low-rank and full-rank gradients are highly similar—is both surprising and impactful. This observation is well-supported by quantitative evidence (cosine similarity, correlation coefficients), and it forms the basis for a practical and effective norm-based scaling method that approximates adaptive optimizer behavior without storing full-rank optimizer states.\n\n2. A common but under-addressed issue in low-rank training is the occurrence of loss spikes, especially when switching projection subspaces. The proposed norm-growth limiter directly tackles this by adaptively smoothing gradient norm increases. Experiments demonstrate that this mechanism outperforms traditional gradient clipping and other heuristics, contributing to much more stable training.\n\n3. The paper conducts extensive experiments across a wide range of LLaMA models (60M–7B) in both pre-training and fine-tuning settings. Fira consistently outperforms LoRA and GaLore in performance and memory usage, including an 8× reduction in optimizer memory when scaling to LLaMA-7B.\n\n4. Fira is easy to integrate into existing training pipelines, requiring only minimal code changes.\n\nWeaknesses:\n\n1. The theoretical grounding for the norm-based scaling method relies on empirical similarity rather than formal guarantees. While the observed correlation is strong, a deeper theoretical explanation would strengthen the contribution.\n\n2. The comparison scope omits other memory-efficient optimizers, such as Adafactor or 8-bit Adam. Including these baselines could better contextualize Fira’s benefits.\n\n3. Only the LLaMA 7B model is used for large-scale validation. While the results are promising, testing on additional architectures (e.g., GPT, Mistral) would improve generality. That said, the limitation is understandable given the high computational cost.\n\nQUESTIONS:\n1. While the empirical similarity between low-rank and full-rank scaling factors is clearly demonstrated, the theoretical justification (e.g., in Appendix E) remains limited and relies on simplifying assumptions. Could the authors provide a more rigorous theoretical explanation or clarify under what conditions such similarity is expected to hold?\n\n2. The norm-based scaling update involves computing the residual term $G_t - P_t R_t$. It appears this term is used during the update step but not explicitly stored. Could the authors confirm whether this computation introduces any additional memory overhead compared to GaLore?\n\n3. In Figure 5, Fira even outperforms full-rank training in some settings. This result appears consistent with findings in [1], which suggest that low-rank constraints can improve the optimization landscape by avoiding sharp minima. Could the authors discuss whether this phenomenon helps explain Fira’s superior performance? If so, it would be helpful to explicitly cite [1] and further analyze this connection.\n\n[1] *On the Optimization Landscape of Low-Rank Adaptation Methods for Large Language Models*, ICLR 2025.\n\nLIMITATIONS:\nThe paper does not explicitly discuss the limitations of the proposed method.\n\n1. The method has only been evaluated on LLaMA architectures, and its generalization to other popular LLM architectures (e.g., GPT, OPT, Mistral) remains unverified.\n2. While the paper demonstrates strong results on LLaMA-7B, larger models have not been tested. It is unclear how Fira performs at even greater scales (e.g., 13B or 65B), where memory and stability challenges may differ.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present Fira, a method to handle low-rank optimizer scaling for full-rank training. The authors observe that the overall magnitude of Adam's update for a layer is very similar whether it's full-rank or low-rank. So Fira replaces Adam's per element update with an update that has the same direction as the full-rank gradient, but the magnitude is derived from the low-rank Adam stats.\n\nFira seems to have two components: a norm-based scaling strategy that leverages the scaling effects of low-rank optimizers to facilitate full-rank training, and a norm-growth limiter to address the issue of loss spikes by limiting the growth of gradient norm.\n\nFira introduces the concept of a scaling factor that represents the magnitude of the correction applied by the adaptive optimizer to the gradient norm. Since at the matrix level, it exhibits a high degree of similarity between low-rank and full-training training, Fira uses it as a norm-based scaling method that leverages the scaling factor of a weight matrix in low-rank.\n\nExtensive experiments in both pre-training and fine-tuning of LLMs demonstrate the effectiveness of Fira.\n\nSTRENGTHS AND WEAKNESSES:\n**Quality**\n\nStrengths: The methodology looks good: the norm-based scaling rule and the limiter are both derived from concrete empirical observations about Adam's per-matrix norms, and the paper backs them up with ablations showing each part is indispensable (perplexity jumps from 31.06 to 37.06 without scaling, and to 32.22 without the limiter).\n\nExperiments cover four model sizes up to 1 B parameters plus a 7 B scaling test, compare against strong baselines (LoRA, GaLore, ReLoRA, full-rank), and report both memory and accuracy metrics. The algorithm is easy to reproduce (three extra lines of PyTorch) and the authors release code.\n\nWeaknesses: Real application evaluation is limited: pre-training uses only the C4 corpus and runs 10 K steps on the 7 B model, so long-horizon stability is untested.\n\n**Clarity**\n\nStrengths: The paper is easy to follow: Figure 1 contrasts LoRA/GaLore/Fira conceptually, Algorithm 1 gives step-by-step pseudocode, and spike analysis is illustrated with paired loss-and-norm plots. Notation is consistent and an appendix supplies proofs and implementation details.\n\n**Significance**\n\nStrengths: Optimizer-state memory is a real bottleneck for 7 B models and up. Reducing it by about 61% on a 1 B LLaMA while matching or beating full-rank perplexity is practically valuable. Because Fira is optimizer-agnostic, it looks very promising to be adopted widely.\n\nWeaknesses: Memory footprint evaluation seems limited to the optimizer part.\n\n**Originality**\nStrengths: the idea that there’s some sort of norm scaling in Adam and that is preserved in low-rank and full-rank seems to be novel. This is proven empirically in Table 1 via cosine similarity.\n\nWeaknesses: Fira is heavily based on Galore’s approach.\n\nQUESTIONS:\n- Does it make sense to compare against system-level baselines such as ZeRO-offload in [23] or quantized optimizers?\n- Is it fair to say that the method reduces the optimizer memory but activation and parameter memory can still dominate the total memory footprint?\n- Downstream impact is somehow limited but it seems enough to demonstrate the applicability of the approach. What limitations do you forecast on vision or SFT or RLHF?\n- How similar is conceptually and its effect the norm-growth limiter to gradient-clipping?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a new low-memory training technique for LLMs called Fira. This method reduces the memory footprint of Adam by storing optimizer states only for a low-rank approximation of the gradient, while still using the full-rank gradient at each step to improve performance. To achieve this, the authors not only use the low-rank gradient approximation with the standard Adam update (as done in low-rank methods like GaLore), but also utilize the remaining part of the full-rank gradient after applying column-wise or matrix-wise scaling to it based on the same low-rank optimizer states. The paper motivates this scaling approach and proposes a strategy to stabilize the optimization, specifically addressing the spikes in the loss curve that occur during training with this procedure. Empirical results demonstrate that Fira is effective compared to other low-memory methods in both pretraining and fine-tuning of LLMs, and it can even outperform full-rank Adam.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper addresses an important direction: low-memory optimization is crucial from a resource-efficiency perspective, yet the effectiveness of many existing methods remains limited in settings that require substantial training, such as pretraining or fine-tuning on highly dissimilar tasks.\n2. The paper makes a reasonable effort to motivate all proposed modifications and supports them with a broad set of ablation experiments.\n3. The proposed method demonstrates strong performance in both pretraining and fine-tuning settings. It not only outperforms other low-memory methods but also achieves better results than standard Adam in most of the cases.\n4. Overall, while I have some concerns regarding the motivation and generalizability of the results, I find the empirical evidence convincing and believe the proposed method is effective.\n\nMain weaknesses:\n1. Motivation for scaling. Since scaling is a central component of the proposed method, I believe its motivation should be explained more clearly in the paper. Could you clarify how the comparison between scaling factors for the low-rank and full-rank gradients is performed? Are both computed within the same training run, and if so, what training procedure is used? My concern is that observing similar scaling factors along the trajectory of a single optimization procedure (either low-rank or full-rank) does not necessarily imply that the scaling factors are similar across the two procedures. As a result, it is not clear whether the proposed modification is expected to make low-rank training behave more like full-rank training.\n2. Fira vs full rank. Continuing the previous point, it appears that the main accuracy gains of Fira do not result from a better approximation of full-rank training, but rather from generally changing the scaling procedure from element-wise to column-wise. This change is beneficial even in the full-rank setting, as indicated by the results for Fira-only-scaling in Table 13. However, the paper does not provide a clear explanation of why this scaling modification is effective. The discussion in the appendix about escaping local optima is vague and lacks depth. Moreover, column-wise scaling is not always better than element-wise scaling. Fira performs better than Fira-only-scaling, which suggests that element-wise scaling is helpful for the main components of the gradient but may be detrimental for others. Since the choice of scaling strategy is such an important factor in Fira’s performance, it should be analyzed in more detail in the paper.\n3. Related works and baselines. From a practical perspective, my main concern is the limited discussion of other low-memory training techniques and the lack of comparison to them. For example, there is a line of research on low-memory Adam modifications, including Adam-mini (https://arxiv.org/abs/2406.16793), AdaLomo (https://arxiv.org/abs/2310.10195), MicroAdam (https://arxiv.org/abs/2405.15593), and SlimAdam (https://arxiv.org/abs/2503.01843 - this one which is a concurrent work but still relevant). These methods are clearly related to the proposed technique and therefore should be discussed in the related work section and included as baselines in the experiments. Additionally, GaLore-add, which is described as the first step of Fira’s modifications, appears very similar to FRUGAL (https://arxiv.org/abs/2411.07837), so this connection should be clearly stated. Overall, the baselines considered in the paper (LoRA, ReLoRA, GaLore) are quite basic. Including comparisons with current state-of-the-art low-memory training methods would be beneficial to better position the proposed method in context. I do not expect the proposed method to necessarily be state of the art, but such a comparison would provide valuable perspective.\n\nAdditional concerns, comments and questions:\n1. Writing. The writing of the paper could be improved. In particular, the motivation for the proposed scaling technique is not clearly described in the main text. Even after reading the appendix, I still have some questions. Additionally, the comparison between matrix and column scaling is discussed somewhat inconsistently, which could be clarified.\n2. Limited experimental results. All experiments in the paper are conducted using the same architecture (LLaMA) and include only one dataset for pretraining and one for fine-tuning. Additional experiments with different architectures and datasets would strengthen the generality of the claims. Moreover, based on the abstract and introduction, the paper appears to position the proposed modification as general for adaptive optimizers, while the experiments are performed only with Adam. The paper should either narrow its claims or provide results using other adaptive optimizers. Additionally, including full-rank experiments for LLaMA 7B pretraining would be beneficial to evaluate whether Fira consistently outperforms full-rank methods on larger-scale models.\n3. Memory and compute. I am a bit confused about the memory usage comparisons in Tables 3 and 8. Could you please explain why Fira is expected to have a significant memory improvement according to Table 3, but shows almost no real improvement in Table 8? Also, why does Fira provide a substantial decrease in throughput compared to GaLore during fine-tuning but not during pretraining?\n4. The scaling in Fira can be applied at the matrix or column level. Is the same true for the norm growth limiter, or is it always applied at the matrix level? If it is always applied at the matrix level, could you explain why?\n\nQUESTIONS:\nAll questions and concerns are detailed in the Strengths and Weaknesses section. The main weaknesses 1–3 are the most critical and will have the greatest impact on my final evaluation after the rebuttal.\n\nLIMITATIONS:\nI believe the limitations section should be more thorough and explicitly mention the limited understanding of the reasons behind the method’s benefits. Additionally, it should acknowledge that the experiments are limited to a single architecture, dataset, and adaptive optimizer.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces Fira, a novel low-rank training framework that achieves the performance of full-rank training while maintaining the memory efficiency of low-rank methods. Fira restores full-rank training quality while keeping GaLore-style low-rank memory use. Starting from GaLore’s gradient-projection scheme, Fira restores GaLore’s discarded gradient component as a new term when updating, scales it with a norm-based factor that emulates adaptive optimizer corrections without extra state, and keeps training stable via a norm-growth limiter that caps each step’s gradient-norm increase. Experiment results show Fira outperforms full-rank, GaLore, LoRA and ReLoRA on LLaMA pre-training (60 M–1 B) and beats GaLore on LLaMA-7B using 8× smaller rank; fine-tuning on eight reasoning tasks also shows consistent gains.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n1. Novel norm-based scaling recovers discarded gradients in GaLore, which improve the performance without extra optimizer state storage.\n2. Simple to implement but effective.\n3. The experiments cover various model sizes (60M to 7B), demonstrate consistent performance gains in both pre-training and fine-tuning across multiple commonsense reasoning tasks, and include thorough and fair ablation studies. And also good result analysis, especially analysis in Appendix F.\n4. Good error-bound analysis in Appendix E.\n5. well-written \n\n## Weaknesses\n1. Although the experimental design is comprehensive, the comparison with more recent training methods is limited.\n\nFor instance:\nLiu, Xu-Hui, et al. \"On the Optimization Landscape of Low Rank Adaptation Methods for Large Language Models.\" The Thirteenth International Conference on Learning Representations. 2025.\n\nHow would this method work with the gradient random projection proposed above? Further discussion and an experiment are needed to draw a convincing conclusion. \n\n\n2. No precise FLOPs since there is additional computing overhead.\n\nQUESTIONS:\n1. The paper emphasizes memory efficiency but provides limited quantitative analysis on computational overhead, especially the cost of additional operations introduced by Fira. Could the authors include a detailed breakdown of computational complexity or measured runtime overhead compared to baseline methods?\n\n2. While comparisons with LoRA, GaLore, and ReLoRA are provided, newer methods involving tensor decompositions or other gradient compression techniques are not discussed. Could the authors elaborate on how Fira compares to other recent low-rank or gradient-efficient methods?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-05"}