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

## Abstract

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. This raises a question: whether it is possible 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? In this paper, we propose a new plug-and-play training framework for LLMs called Fira, as the first attempt to achieve this goal. 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. Based on this observation, 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 enable full-rank training. In this way, we can preserve the low-rank constraint in the optimizer while achieving full-rank training for better performance. Moreover, we find that there are sudden gradient rises during the optimization process, potentially causing loss spikes. To address this, we further put forward a norm-growth limiter to smooth the gradient via regulating the relative increase of gradient norms. Extensive experiments on the pre-training and fine-tuning of LLMs show that Fira outperforms both LoRA and GaLore, achieving performance that is comparable to or even better than full-rank training. For instance, our Fira can reduce the memory usage of optimizer states by 61.1%, while achieving improved performance for pre-training on the LLaMA 1B architecture. Notably, for pre-training on the LLaMA 7B architecture, our method uses an 8 × 8\times 8 × smaller rank 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 (Achiam et al., ; Sima et al., ; Feng et al., ). While the substantial increase in model size contributes significantly to these advancements, it also introduces considerable memory bottlenecks, especially for optimizer states (Zhao et al., ).
For instance, pre-training a LLaMA-7B model from scratch (^1^11Training the model with a single batch size and maximum sequence length of 2048 under BF16 precision.) requires at least 58 GB memory, allocated as follows: 14GB for loading parameters, 14GB for weight gradients, 28GB for Adam (Kingma & Ba, ) optimizer states, and 2GB for activations (Zhao et al., ).
Notably, the optimizer states consume even more memory than the parameters themselves. To address this, low-rank training has demonstrated its effectiveness to reduce the memory usage of the optimizer states by conducting training in a low-rank subspace (Zhao et al., ; Hu et al., ).

The current low-rank training methods can be broadly divided into two categories: weight matrix-based and gradient matrix-based low-rank decomposition.
For the weight matrix decomposition methods, the most representative one is Low-Rank Adaptation (LoRA) (Hu et al., ), 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 sub-optimal outcomes (Zhang et al., ; Xia et al., ).
Although the variant ReLoRA (Lialin et al., ) 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 (Zhao et al., ).
For the gradient matrix decomposition based methods, the typical one is the gradient low-rank projection (GaLore) proposed recently (Zhao et al., ). 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?

Figure: (a) LoRA
Refer to caption: x1.png

In light of this, we propose a new memory-efficient training framework for LLMs, called Fira, which, to the best of our knowledge, is the first to achieve full-rank training while consistently maintaining a low-rank 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 commonly-used adaptive optimizer (e.g., Adam).
As a result, the adaptive optimizer fails to correct the full-rank raw gradient according to the optimizer states.
Without this correction, adaptive optimization algorithms would degrade into simple SGD, leading to significantly reduced optimization performance (Kingma & Ba, ; Zhang et al., ). This point is further validated in Section 4.1 and Section 5.4.
Fortunately, we observe an interesting phenomenon during LLM training: the *scaling factor* (^2^22The scaling factor $\phi_{t}(R_{t})$ is defined as $\frac{||\psi(R_{t})||}{||R_{t}||}$, where $||R_{t}||$ is the norm of the raw gradient, $||\psi(R_{t})||$ is the norm of the gradient corrected by the gradient correction function $\psi$ of the optimizer (e.g., Adam).) of the optimizer (e.g., Adam) is similar from low-rank training to full-rank training.
As illustrated in Figure 3, if we sort the weight matrices by their average scaling factors, we can obtain a similar rank order.
A detailed quantitative analysis of this similarity is presented in Appendix A.5 and A.6.

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.
Additionally, we observe sudden increases in the gradient during training, which results in a spike in training loss (as depicted in Figure 4).
This issue could lead to substantial parameter updates, causing the loss function to reach a much higher value and undermining prior optimization efforts (Goodfellow et al., ; Zhang et al., ). Despite the use of gradient clipping techniques (Pascanu et al., ), this issue may not be adequately resolved, as shown in Figure 4.
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 mitigates the problem of loss spikes.

Our main contributions can be summarized as follows:

- 1.
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.
- 2.
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.
- 3.
Extensive experiments across various parameter counts (60M, 130M, 350M, 1B, 7B) validate the effectiveness of Fira in both pre-training and fine-tuning tasks. Our framework not only outperforms both LoRA and GaLore, but also achieves performance comparable to or better than full-rank training.

## 2 Related Work

Low-rank Adaptation. Low-Rank Adaptation (LoRA) has been introduced by Hu et al. () 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 (Zhang et al., ; Wen & Chaudhuri, ; Xia et al., ; Zhang et al., ; Dettmers et al., ).
For instance, ReLoRA (Lialin et al., ) is proposed to extend the application of LoRA from fine-tuning to pre-training. 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 (Xia et al., ).

Gradient Projection. Recent works (Zhang et al., ; Xia et al., ; Valipour et al., ) have indicated that LoRA may yield sub-optimal performance since its low-rank constraints in parameters.
Inspired by traditional projected gradient descent methods (Chen & Wainwright, ; Chen et al., ), GaLore (Zhao et al., ) 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 full-parameter 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 (Chen et al., ; Ren et al., ). However, most of these methods achieve memory efficiency by compromising either time or precision. Gradient checkpointing (Chen et al., ) is proposed to reduce memory usage by trading increased computational time for the re-computation of activations. Quantization (Dettmers et al., ) reduces memory consumption by using lower-bit data types, but at the cost of model precision. Memory offloading (Zhang et al., ; Ren et al., ) 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}\in\displaystyle\mathbb{R}^{m\times n}$. The full-rank gradient can be represented as $G_{t}=\nabla_{W}f_{t}(W_{t})\in\displaystyle\mathbb{R}^{m\times n}$, where $f$ is the objective function. Then the regular full-rank training can be expressed as follows:

$$ ${}W_{t+1}=W_{t}-\eta\psi_{t}(G_{t}),$ (1) $$

where $\eta$ is the learning rate, and $\psi_{t}$ is the gradient correction function of the optimizer (for vanilla SGD, $\psi_{t}(G_{t})=G_{t}$). Instead of vanilla SGD, adaptive optimizers (e.g., Adam (Kingma & Ba, ), AdamW (Loshchilov & Hutter, )) 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 (Kingma & Ba, ) requires storing the optimizer states $M$ and $V$, which consume $2mn$ of memory. The gradient correction process is as follows:

$$ ${}M_{t}=\beta_{1}M_{t-1}+(1-\beta_{1})G_{t},$ (2) $$

$$ ${}V_{t}=\beta_{2}V_{t-1}+(1-\beta_{2})G_{t}^{2},$ (3) $$

$$ ${}\psi_{t}(G_{t})=\frac{\sqrt{1-\beta_{2}^{t}}}{1-\beta_{1}^{t}}\cdot\frac{M_{ t}}{\sqrt{V_{t}}+\epsilon},$ (4) $$

where all matrix operations are element-wise. $\beta_{1}$ and $\beta_{2}$ are Adam’s hyper-parameters, and $\epsilon$ is a small constant (e.g., $1\times 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 (Hu et al., ) and Galore (Zhao et al., ), have been proposed to reduce memory usage in recent years.

### 3.2 Low-Rank Adaptation

The basic idea behind LoRA (Hu et al., ) 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}\in\displaystyle\mathbb{R}^{m\times n}$ and incorporates two low-rank matrices $A_{t}$ and $B_{t}$ for training as:

$$ ${}W_{t}=W_{0}+B_{t}A_{t},$ (5) $$

where $B_{t}\in\displaystyle\mathbb{R}^{m\times r}$, $A_{t}\in\displaystyle\mathbb{R}^{r\times 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}$.

### 3.3 Gradient Low-Rank Projection

In contrast to LoRA, GaLore (Zhao et al., ) utilizes a projection matrix $P_{t}\in\displaystyle\mathbb{R}^{m\times r}$ to project the full-rank gradient $G_{t}\in\displaystyle\mathbb{R}^{m\times n}$ to a low-rank gradient $R_{t}=P^{\top}_{t}G_{t}\in\displaystyle\mathbb{R}^{r\times n}$ ($m\leq n$)(^3^33For simplicity, we assume $m\leq n$, following (Zhao et al., ). If $m>n$, $R_{t}=G_{t}Q_{t}\in\displaystyle\mathbb{R}^{m\times r},Q_{t}\in\displaystyle \mathbb{R}^{n\times r}$.). By doing so, the memory usage of optimizer states could be reduced. The parameter update in GaLore can be formulated as:

$$ ${}W_{t+1}=W_{t}-\eta P_{t}\psi_{t}(R_{t}),$ (6) $$

where
the projection matrix $P_{t}$ can be obtained through singular value decomposition (SVD) of $G_{t}$ and can be updated every $T$ steps:

$$ ${}G_{t}=U\Sigma V^{\top}\approx\sum_{i=1}^{r}\sigma_{i}u_{i}v_{i}^{\top},\quad P _{t}=[u_{1},u_{2},\ldots,u_{r}],$ (7) $$

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 (Zhao et al., ).
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 (Zhao et al., ). 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.

Figure: Figure 2: Training loss of different methods for pre-training LLaMA 60M on C4 dataset with $r/d_{model}$ = 16/256 and T = 200.
Refer to caption: x4.png

Challenge Analysis.
Given the difficulty of incorporating trainable low-rank weights into LoRA to achieve full-rank weight training (Zhao et al., ), we focus on investigating how to achieve full-rank gradient training by extending the gradient projection method, Galore, in this paper.
In GaLore, the projection matrix $P_{t}\in\displaystyle\mathbb{R}^{m\times r}$ projects the full-rank gradient $G_{t}\in\displaystyle\mathbb{R}^{m\times n}$ of the full-rank weight $W_{t}\in\displaystyle\mathbb{R}^{m\times n}$, to the low-rank subspace gradient $R_{t}=P^{\top}_{t}G_{t}\in\displaystyle\mathbb{R}^{r\times n}$.
The gradient outside this subspace can be represented as: $(I-P_{t}P^{\top}_{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 of $G_{t}$ for correction in optimizers. This would lead to significant information loss especially when $r\ll 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 6. In Figure 6, 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 Eq. (6) as follows:

$$ ${}W_{t+1}=W_{t}-\eta P_{t}\psi_{t}(R_{t})-\eta(G_{t}-P_{t}R_{t}).$ (8) $$

We denote the update strategy in Eq. (8) as GaLore-add. However, as illustrated in Figure 2, GaLore-add exhibits almost no improvement compared to updates using Eq. (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 sub-optimal outputs. Besides, in GaLore-add, $P_{t}\psi_{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:

$$ ${}\phi_{t}(R_{t})=\frac{||\psi_{t}(R_{t})||}{||R_{t}||},$ (9) $$

where the scaling factor $\phi_{t}$ represents the magnitude of the correction applied by the adaptive optimizer to the gradient norm.
Based on the scaling factor $\phi_{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.
As shown in Figure 3, sorting weight matrices by their average scaling factors results in an almost similar order.

Based on this observation, we can use the scaling factors in low-rank training to replace those in full-rank training, even though the absolute magnitude of the scaling factors may vary between low-rank and full-rank training.
This is because the absolute magnitude can be regarded as the learning rate, to which adaptive optimizers (e.g., Adam) are not sensitive (Zhao et al., ; Liu et al., ). Instead, the discrepancies (or relative orders) in the correction magnitudes of different parameters (or weight matrices) are of greater significance.

Norm-based Scaling. Inspired by this, we propose a norm-based scaling method that utilizes the scaling factor of a weight matrix in low-rank training as a substitute for the corresponding factor in full-rank training:

$$ ${}W_{t+1}=W_{t}-\eta P_{t}\psi_{t}(R_{t})-\eta\phi_{t}(R_{t})(G_{t}-P_{t}R_{t}).$ (10) $$

By Eq. (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.

Figure: (a) $r/d_{model}$ = 16/256
Refer to caption: x5.png

Furthermore, we can use a more fine-grained average of the scaling factor in Eq.(10), by considering each column in the weight matrix:

$$ ${}\phi_{t}(R_{t})_{i}=\frac{||\psi(R_{t,:,i})||}{||R_{t,:,i}||},\quad i=1,2, \dots,n,$ (11) $$

where $R_{t,:,i}$ is the $i$-th column of $R_{t}$, and $\phi_{t}(R_{t})_{i}$ is the $i$-th scaling factor.

### 4.2 Norm-Growth Limiter

We find that there are suddenly sharp increases of the gradient during training, which could introduce loss spikes. As shown in Figure 4, 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 which transforms abrupt gradient spikes into gradual, smooth increases.

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$ (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 gradient uses a new projection matrix, leading to significant misalignment. Furthermore, as shown in Figure 2, GaLore-add also exhibits training spikes, 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 (Zhang et al., ). Due to space constraints, further analysis is provided in Appendix A.7.

Figure: Figure 4: Training loss and gradient norm of three variants of Fira for pre-training LLaMA 60M on C4 dataset.
Refer to caption: x8.png

Addressing Loss Spikes. To address this issue, a straightforward solution is to use gradient clipping techniques (Pascanu et al., ) 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 be also verified in Figure 4 and Table 4.
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 $\gamma$ when the gradient norm increases:

$$ ${}\text{if }\frac{\displaystyle||S_{t}||}{\displaystyle||S_{t-1}||}>\gamma \text{ then }S_{t}\leftarrow\frac{S_{t}}{\displaystyle||S_{t}||}\cdot\gamma \displaystyle||S_{t-1}||,$ (12) $$

where $\gamma$ is a threshold ensuring that the rate of gradient growth does not exceed this value. $S_{t}=\phi_{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 4, Fira with our proposed limiter improves the optimization performance without significant spikes.

Figure: Algorithm 1 Fira with Adam

**Table 1: Comparison between Fira, GaLore, and LoRA. Denote $W_{t}\in\mathbb{R}^{m\times n}$ ($m\leq n$), rank $r$.**
|  | Fira | GaLore | LoRA |
| --- | --- | --- | --- |
| Weights | $mn$ | $mn$ | $mn+mr+nr$ |
| Optimizer States | $mr+2nr+1$ | $mr+2nr$ | $2mr+2nr$ |
| Full-Rank Gradients | ✓ | ✗ | ✓ |
| Full-Rank Weights | ✓ | ✓ | ✗ |
| Pre-Training | ✓ | ✓ | ✗ |
| Fine-Tuning | ✓ | ✓ | ✓ |

### 4.3 Overall Algorithm

We present the overall algorithm of Fira with Adam in Algorithm 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 which 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.4.

It’s worth noting that Fira only introduces one parameter $||S_{t-1}||$ for each weight matrix in the optimizer state, which is negligible, as shown in Table 1.
Besides, in addition to the original hyper-parameters of optimizers and gradient projection methods, Fira only adds one hyper-parameter $\gamma$ in the norm-growth limiter.
The hyper-parameter $\gamma$ is set to 1.01 across all experiments, which consistently yields satisfactory results.

## 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 Eq. (10) as Fira-matrix, and denote our method additionally using the column-wise strategy of Eq. (11) as Fira.

### 5.1 Memory-Efficient Pre-training

Experimental Setup. We follow the settings in Galore (Zhao et al., ) to conduct the pre-training experiments. We compare Fira with GaLore (Zhao et al., ), LoRA (Hu et al., ), ReLoRA (Lialin et al., ), and full-rank training baselines. Adam optimizer is used for training all baselines and our method on the C4 dataset in the BF16 format. The settings of these baselines can be found in Zhao et al. (). The dataset C4 is a colossal, cleaned version of Common Crawl’s web crawl corpus, which is widely used in LLM pre-training (Raffel et al., ). Following Zhao et al. (), we utilize LLaMA-based architectures equipped with RMSNorm and SwiGLU activations (Zhang & Sennrich, ; Shazeer, ; Touvron et al., ). As in Zhao et al. (), 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 A.2.
We use 8 A100 80G GPUs to conduct pre-training experiments.

**Table 2: Comparison of different methods for pre-training LLaMA models of various sizes on the C4 dataset. We report validation perplexity ($\downarrow$) with a memory estimate of total parameters and optimizer states. Results and memory estimates of all baselines are taken from Zhao et al. (). $r$ refers to the rank and $d_{model}$ is the full-rank dimension of models.**
|  | 60M | 130M | 350M | 1B |
| --- | --- | --- | --- | --- |
| Full-Rank | 34.06 (0.36G) | 25.08 (0.76G) | 18.80 (2.06G) | 15.56 (7.80G) |
| Fira | 31.06 (0.24G) | 22.73 (0.52G) | 16.85 (1.22G) | 14.31 (4.38G) |
| GaLore | 34.88 (0.24G) | 25.36 (0.52G) | 18.95 (1.22G) | 15.64 (4.38G) |
| LoRA | 34.99 (0.36G) | 33.92 (0.80G) | 25.58 (1.76G) | 19.21 (6.17G) |
| ReLoRA | 37.04 (0.36G) | 29.37 (0.80G) | 29.08 (1.76G) | 18.33 (6.17G) |
| $r/d_{model}$ | 128 / 256 | 256 / 768 | 256 / 1024 | 512 / 2048 |
| Training Tokens | 1.1B | 2.2B | 6.4B | 13.1B |

Result Analysis. As shown in Table 2, Fira consistently outperforms low-rank training baselines by a large margin under the same rank constraint, and even surpasses full-rank training.
Following Zhao et al. (), we estimate the memory reduction of the optimizer states via the same memory estimation method introduced in Zhao et al. ().
From Table 2, 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 achieving 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 (Zhou et al., ).

### 5.2 Scaling up to LLaMA 7B Pre-training.

Figure: Figure 5: Pre-training LLaMA 7B with different methods on the C4 dataset.
Refer to caption: x9.png

To validate the scalability of our method, we scale up by pre-training the LLaMA 7B model with the full-rank dimension $d_{\text{model}}=4096$.
We compare Fira with the GaLore baseline, which generally achieves the best performance among low-rank training baselines, as shown in Table 2.
As illustrated in Figure 5, our method demonstrates a
significant improvement over GaLore for pre-training LLaMA 7B, while using an $8\times$ smaller rank.
This highlights Fira’s effectiveness, suggesting it could be a viable solution for large-scale LLM pre-training.

**Table 3: Accuracy ($\uparrow$) of various fine-tuning methods on eight commonsense reasoning datasets with LLaMA 7B. Results for all baseline methods, except GaLore, are taken from Hu et al. ().**
| Method | Memory | BoolQ | PIQA | SIQA | HellaSwag | WinoGrande | ARC-e | ARC-c | OBQA | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Fira | 0.26G | 69.4 | 82.6 | 78.0 | 76.8 | 81.2 | 82.2 | 64.4 | 80.8 | 76.9 |
| Prefix | 0.05G | 64.3 | 76.8 | 73.9 | 42.1 | 72.1 | 72.9 | 54.0 | 60.6 | 64.6 |
| Series | 0.42G | 63.0 | 79.2 | 76.3 | 67.9 | 75.7 | 74.5 | 57.1 | 72.4 | 70.8 |
| Parallel | 1.49G | 67.9 | 76.4 | 78.8 | 69.8 | 78.9 | 73.7 | 57.3 | 75.2 | 72.2 |
| LoRA | 0.35G | 68.9 | 80.7 | 77.4 | 78.1 | 78.8 | 77.8 | 61.3 | 74.8 | 74.7 |
| GaLore | 0.26G | 69.5 | 82.0 | 75.1 | 32.2 | 18.0 | 80.7 | 65.8 | 78.0 | 62.7 |

### 5.3 Memory-Efficient Fine-Tuning

Experimental Setup. Following Hu et al. (), we perform the fine-tuning task to compare Fira with LoRA, GaLore, and other baseline methods, including Prefix-tuning (Prefix) (Li & Liang, ), Series Adapter (Series) (Houlsby et al., ), and Parallel Adapter (Parallel) (He et al., ), 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 Hu et al. (), 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 A.3. We adopt RTX 4090 GPUs for fine-tuning experiments.

Result Analysis.
As shown in Table 3, our Fira achieves the highest performance on 4 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 four 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 $\phi_{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 to avoid loss spikes instead of our proposed norm-growth limiter.

**Table 4: Ablation study on the C4 dataset.**
| Method | Perplexity ($\downarrow$) |
| --- | --- |
| Fira-w.o.-scaling | 37.06 |
| Fira-matrix | 31.52 |
| Fira-w.o.-limiter | 32.22 |
| Fira-gradient-clipping | 31.22 |
| Fira | 31.06 |

Table 4 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 and Fira-gradient-clipping, indicating the effectiveness of our proposed norm-growth limiter in addressing the issue of training loss spikes.

### 5.5 Performance under varying ranks

Figure: Figure 6: Validation perplexity of Fira and GaLore for varying ranks when pre-training LLaMA 60M on the C4 dataset with $d_{model}=256$.
Refer to caption: x10.png

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 6.
From Figure 6, 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 there are sudden spikes in gradient values during optimization, which could lead to spikes in 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.

## Appendix A Appendix

### A.1 Theoretical Analysis

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 (Liu et al., ). As illustrated in Figure 8, 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 mitigated the variance in Fira efficiently.
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 $\{\bm{g}^{(i)}\}_{i=1}^{n}$, where each $\bm{g}^{(i)}=(g^{(i)}_{1},g^{(i)}_{2},\dots,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 (Liu et al., ), we assume the adaptive learning rate of Adam $\psi(.)=\scriptstyle\sqrt{\frac{1-\beta_{2}^{t}}{(1-\beta_{2})\sum_{i=1}^{t}
\beta_{2}^{t-i}g_{i}^{2}}}$, and $g^{(i)}_{j}\sim\mathcal{N}(0,\sigma^{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(\psi(.))=p(\scriptstyle\sqrt{\frac{1-\beta_{2}^{t}}{(1-\beta_{2})\sum_{i=1}^
{t}\beta_{2}^{t-i}g_{i}^{2}}})$$\approx p(\sqrt{\frac{t}{\sum_{i=1}^{t}g_{i}^{2}}})$ (Nau, ), and then $\psi^{2}(.)$ $\sim$ Scale-inv-$\mathcal{X}^{2}(\rho,\frac{1}{\sigma^{2}})$.

Theorem 1. (Variance of Scaling Factors) In the early stages of training, if $\psi^{2}(\cdot)\sim\text{Scale-inv-}\mathcal{X}^{2}(\rho,\frac{1}{\sigma^{2}})$, and $g_{j}^{(i)}\sim\mathcal{N}(0,\sigma^{2})$(^4^44The assumption of a mean-zero normal distribution is valid at the outset of training, as the weights are sampled from normal distributions with a mean of zero (Balduzzi et al., )) for all $i,j$, then for all $\rho>4$, the scaling factor $\phi^{2}=\frac{\sum_{i=1}^{n}\psi_{i}^{2}(g_{t}^{(i)})^{2}}{\sum_{i=1}^{n}(g_{
t}^{(i)})^{2}}$ satisfies $\operatorname{Var}[\phi^{2}]\leq\operatorname{Var}[\psi^{2}]$. If we approximate $\sqrt{\psi^{2}}$ and $\sqrt{\phi^{2}}$ to the first order, we have $\operatorname{Var}[\phi]\leq\operatorname{Var}[\psi]$.

###### Proof.

We express $\phi^{2}$ as a weighted sum:

$$ $\phi^{2}=\sum_{i=1}^{n}w_{i}\psi_{i}^{2},$ (13) $$

where the weights are defined as:

$$ $w_{i}=\frac{(g_{t}^{(i)})^{2}}{\sum_{j=1}^{n}(g_{t}^{(j)})^{2}}.$ (14) $$

Each $w_{i}$ is a non-negative random variable satisfying $\sum_{i=1}^{n}w_{i}=1$.

In the context of adaptive optimization algorithms like Adam, the squared gradients $\psi_{i}^{2}$ accumulate information from past iterations to adapt the learning rate for each parameter. With $\beta_{2}=0.999$, the moving average of the squared gradients places significant weight on historical data, making $\psi_{i}^{2}$ dependent mainly on past gradients, yielding:

$$ $\psi_{i}^{2}\approx\psi_{i}^{2}(g^{(i)}_{1},\ldots,g^{(i)}_{t-1}).$ (15) $$

Since $\psi_{i}^{2}$ primarily depend on past gradients $g^{(i)}_{1},\ldots,g^{(i)}_{t-1}$, and $w_{i}$ depend solely on the current gradients $g_{t}^{(i)}$, we can consider $\psi_{i}^{2}$ and $w_{i}$ to be independent random variables.

Consequently, we can express the variance of $\phi^{2}$ as:

$$ $\operatorname{Var}[\phi^{2}]=\operatorname{Var}\left[\sum_{i=1}^{n}w_{i}\psi_{ i}^{2}\right].$ (16) $$

Using the law of total variance, we have:

$$ $\displaystyle\operatorname{Var}[\phi^{2}]$ $\displaystyle=\operatorname{E}\left[\operatorname{Var}\left[\sum_{i=1}^{n}w_{i }\psi_{i}^{2}\mid w_{1},\ldots,w_{n}\right]\right]+\operatorname{Var}\left( \operatorname{E}\left[\sum_{i=1}^{n}w_{i}\psi_{i}^{2}\mid w_{1},\ldots,w_{n} \right]\right).$ (17) $$

Figure: Figure 7: Scaling factor $\phi_{t}(R_{t})$ during the early stage of training (1K iterations of total 10K iterations).
Refer to caption: x11.png

Since $\psi_{i}^{2}$ are independent of the $w_{i}$, we find:

$$ $\displaystyle\operatorname{E}\left[\sum_{i=1}^{n}w_{i}\psi_{i}^{2}\mid w_{1}, \ldots,w_{n}\right]$ $\displaystyle=\operatorname{E}[\psi_{i}^{2}]\sum_{i=1}^{n}w_{i}=\operatorname{ E}[\psi_{i}^{2}],$ (18) $\displaystyle\operatorname{Var}\left[\sum_{i=1}^{n}w_{i}\psi_{i}^{2}\mid w_{1} ,\ldots,w_{n}\right]$ $\displaystyle=\sum_{i=1}^{n}w_{i}^{2}\operatorname{Var}[\psi_{i}^{2}].$ (19) $$

Thus, the variance simplifies to:

$$ $\operatorname{Var}[\phi^{2}]=\operatorname{Var}[\psi^{2}]\operatorname{E}\left [\sum_{i=1}^{n}w_{i}^{2}\right].$ (20) $$

The second term, $\operatorname{Var}\left(\operatorname{E}\left[\sum_{i=1}^{n}w_{i}\psi_{i}^{2}
\mid w_{i}\right]\right)$, is zero since $\operatorname{E}[\phi^{2}\mid w_{i}]=\operatorname{E}[\psi_{i}^{2}]$ is constant.

Let $X_{i}=(g_{t}^{(i)})^{2}$, where each $X_{i}\sim\sigma^{2}\chi^{2}_{1}$. Then, we can express the weights as:

$$ $w_{i}=\frac{X_{i}}{\sum_{j=1}^{n}X_{j}}.$ (21) $$

Since $X_{i}/\sigma^{2}\sim\chi^{2}_{1}$, and each $w_{i}$ is the ratio of $X_{i}$ to the sum of all $X_{j}$, the vector $(w_{1},\ldots,w_{n})$ follows a Dirichlet distribution with parameters $\alpha_{i}=\frac{\nu_{i}}{2}=\frac{1}{2}$, where $\nu_{i}=1$ is the degrees of freedom of $\chi^{2}_{1}$.

For a Dirichlet distribution, the expected value of $w_{i}^{2}$ is given by:

$$ $\operatorname{E}[w_{i}^{2}]=\frac{\alpha_{i}(\alpha_{i}+1)}{\left(\sum_{k=1}^{ n}\alpha_{k}\right)\left(\sum_{k=1}^{n}\alpha_{k}+1\right)}.$ (22) $$

Substituting $\alpha_{i}=\frac{1}{2}$ and $\sum_{k=1}^{n}\alpha_{k}=\frac{n}{2}$ yields:

$$ $\operatorname{E}[w_{i}^{2}]=\frac{\frac{1}{2}\cdot\frac{3}{2}}{\frac{n}{2} \cdot\left(\frac{n}{2}+1\right)}=\frac{3}{4}\cdot\frac{4}{n(n+2)}=\frac{3}{n(n +2)}.$ (23) $$

Thus, summing over all $i$ gives:

$$ $\operatorname{E}\left[\sum_{i=1}^{n}w_{i}^{2}\right]=n\cdot\operatorname{E}[w_ {i}^{2}]=\frac{3}{n+2}.$ (24) $$

Finally, substituting this result back into the variance expression:

$$ $\operatorname{Var}[\phi^{2}]=\operatorname{Var}[\psi^{2}]\cdot\frac{3}{n+2}.$ (25) $$

Since $n\geq 1$, it follows that:

$$ $\frac{3}{n+2}\leq 1,$ (26) $$

which implies:

$$ $\operatorname{Var}[\phi^{2}]\leq\operatorname{Var}[\psi^{2}].$ (27) $$

Given $\rho>4$ and $\psi^{2}(\cdot)\sim\text{Scale-inv-}\mathcal{X}^{2}(\rho,\frac{1}{\sigma^{2}})$, the variance of $\psi^{2}(\cdot)$ exists (Liu et al., ).

Since $\psi_{i}^{2}$ and $w_{i}$ are independent and $\sum_{i=1}^{n}\operatorname{E}[w_{i}]=1$:

$$ $\operatorname{E}[\phi^{2}]=\operatorname{E}\left[\sum_{i=1}^{n}w_{i}\psi_{i}^{ 2}\right]=\sum_{i=1}^{n}\operatorname{E}[w_{i}]\operatorname{E}[\psi_{i}^{2}]= \operatorname{E}[\psi_{i}^{2}]\sum_{i=1}^{n}\operatorname{E}[w_{i}]= \operatorname{E}[\psi_{i}^{2}],$ (28) $$

Thus, we have shown that:

$$ $\operatorname{Var}[\phi^{2}]\leq\operatorname{Var}[\psi^{2}],\quad\text{and} \quad\operatorname{E}[\phi^{2}]=\operatorname{E}[\psi_{i}^{2}].$ (29) $$

Follow Liu et al. (), we approximate $\sqrt{\psi^{2}}$ and $\sqrt{\phi^{2}}$ to the first order (Wolter & Wolter, )

$$ $\mathrm{Var}[\psi]\approx\frac{\mathrm{Var}[\psi^{2}]}{4\operatorname{E}[\psi^ {2}]},\quad\text{and}\quad\mathrm{Var}[\phi]\approx\frac{\mathrm{Var}[\phi^{2} ]}{4\operatorname{E}[\phi^{2}]}.$ (30) $$

which implies:

$$ $\operatorname{Var}[\phi]\leq\operatorname{Var}[\psi].$ (31) $$

∎

To further examine our theorem, we conduct simulations to calculate the variance of the scaling factor $\phi$ at ranks within the set $\{1,5,10,50,100\}$. The adaptive learning rate $\psi$ is equivalent to that of $\phi$ when the rank equals 1. As shown in Figure 8, the variance decreases as the rank increases, supporting our above theorem $\operatorname{Var}[\phi]\leq\operatorname{Var}[\psi]$. 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.

###### Proof.

### A.2 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 Zhao et al. (). Table 5 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.

**Table 5: Hyper-parameters of LLaMA architectures for pre-training.**
| Params | Hidden | Intermediate | Heads | Layers | Steps | Data Amount (Tokens) |
| --- | --- | --- | --- | --- | --- | --- |
| 60M | 512 | 1376 | 8 | 8 | 10K | 1.3 B |
| 130M | 768 | 2048 | 12 | 12 | 20K | 2.6 B |
| 350M | 1024 | 2736 | 16 | 24 | 60K | 7.8 B |
| 1 B | 2048 | 5461 | 24 | 32 | 100K | 13.1 B |
| 7 B | 4096 | 11008 | 32 | 32 | 150K | 19.7 B |

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 $\alpha$ is considered a fractional learning rate (Zhao et al., ). Furthermore, a relatively large learning rate may result in spikes of the training loss (Zhao et al., ). To address this issue, for models with a size of less than 1B, we set $\alpha$ to 0.25, while for models exceeding 1B, we adjust $\alpha$ to 0.0625.

### A.3 Detailed Fine-tuning Setting

We fine-tune the pre-trained LLaMA-7B model for commonsense reasoning tasks benchmark designed for LLM fine-tuning, which include eight sub-tasks (Hu et al., ). Table 6 shows the hyper-parameter configurations.

**Table 6: Hyper-parameter configurations of fine-tuning LLaMA-7B for Fira.**
| Hyper-parameters | Setting |
| --- | --- |
| Rank $r$ | 32 |
| $\alpha$ | 64 |
| Dropout | 0.05 |
| Base optimizer | AdamW |
| LR | 1e-4 |
| LR Scheduler | Linear |
| Batch size | 16 |
| warm-up Steps | 100 |
| Epochs | 3 |
| Where | Q,K,V,Up,Down |

### A.4 Plug-and-play framework for Fira

Figure: Algorithm 2 Plug-and-play framework for Fira, Pytorch-like.

### A.5 Case Quantitative Analysis of Scaling Factor Similarities

In this section, we analyze the similarities among the three ranking (i.e., order) sequences $R_{1}$, $R_{2}$, and $R_{3}$ depicted in Figure 3, which highlight the discrepancies in scaling factors across different weight matrices. Each ranking comprises ten distinct items. We will employ Kendall’s Tau correlation coefficient (Abdi, ) and Spearman’s rank correlation coefficient (Sedgwick, ) to evaluate their concordance and divergence.

The three ranking sequences are defined as follows:

$$ $R_{1}=(7,6,1,2,4,8,5,10,9,3)$ (32) $$

$$ $R_{2}=(7,8,2,1,5,4,6,10,9,3)$ (33) $$

$$ $R_{3}=(6,8,2,1,5,4,7,10,9,3)$ (34) $$

#### A.5.1 Kendall’s Tau Correlation Coefficient

Kendall’s Tau, a non-parametric statistic, measures the ordinal association between two rankings. It quantifies the degree to which the presence of one ranking implies a similar ranking in another. The formula for Kendall’s Tau is given by:

$$ $\tau=\frac{C-D}{\frac{n(n-1)}{2}}$ (35) $$

where:

- •
$C$ is the number of concordant pairs, which are pairs of observations where the ranks for both items agree,
- •
$D$ is the number of discordant pairs, where the ranks disagree,
- •
$n$ is the total number of observations.

Kendall’s Tau ranges from -1 to +1, with 1 indicating perfect agreement between the rankings, 0 indicating no correlation, and -1 indicating perfect disagreement. This measure is particularly robust against outliers, making it useful for assessing the strength of relationships in ordinal data. A p-value, calculated to evaluate the statistical significance of the observed correlation, tests the null hypothesis that there is no association between the two variables. A low p-value (typically less than 0.05) suggests rejecting the null hypothesis, indicating that the observed correlation is statistically significant.

#### A.5.2 Spearman Rank Correlation Coefficient

The Spearman rank correlation coefficient is another non-parametric measure that assesses the strength and direction of association between two ranked variables. It is calculated as follows:

$$ $\rho=1-\frac{6\sum d_{i}^{2}}{n(n^{2}-1)}$ (36) $$

where:

- •
$d_{i}$ represents the differences between the ranks of each observation,
- •
$n$ is the number of observations.

Spearman’s coefficient also ranges from -1 to +1, with similar interpretations as Kendall’s Tau. A coefficient of 1 indicates perfect positive correlation, -1 indicates perfect negative correlation, and 0 indicates no correlation. Spearman’s method excels when analyzing datasets that fail to meet the normality assumptions requisite for parametric tests. To assess the significance of the Spearman correlation, a p-value is calculated alongside the coefficient. This p-value tests the null hypothesis of no correlation between the rankings of the variables. A small p-value (often less than 0.05) indicates a statistically significant correlation, providing strong evidence against the null hypothesis.

#### A.5.3 Results

The results of the Spearman and Kendall correlation coefficients are summarized in Table 7:

**Table 7: Spearman and Kendall correlation coefficients with corresponding p-values.**
| Sample | Spearman | Kendall |  |  |
| --- | --- | --- | --- | --- |
| Coefficient | P-value | Coefficient | P-value |  |
| $R_{1}R_{2}$ | 0.8545 | 0.0016 | 0.7333 | 0.0022 |
| $R_{1}R_{3}$ | 0.8303 | 0.0029 | 0.6889 | 0.0047 |
| $R_{2}R_{3}$ | 0.9879 | 9.31e-08 | 0.9556 | 5.51e-06 |

#### A.5.4 Analysis and Conclusions

The correlation analysis indicates strong positive relationships among the samples, as summarized below:

- •
For $R_{1}R_{2}$: The Spearman coefficient is 0.8545 (p = 0.0016) and the Kendall coefficient is 0.7333 (p = 0.0022), both indicating strong and statistically significant correlations.
- •
For $R_{1}R_{3}$: The Spearman coefficient of 0.8303 (p = 0.0029) and the Kendall coefficient of 0.6889 (p = 0.0047) further confirm significant positive associations.
- •
For $R_{2}R_{3}$: The Spearman coefficient of 0.9879 (p = 9.31e-08) and the Kendall coefficient of 0.9556 (p = 5.51e-06) suggest an almost perfect correlation, accompanied by high statistical significance.

In general, a p-value less than 0.05 indicates strong evidence against the null hypothesis, suggesting that the observed correlations are statistically significant and unlikely to have occurred by chance.

Overall, all samples exhibit strong correlations. The consistency across both correlation methods underscores the reliability of these findings, suggesting robust relationships that warrant further exploration.

### A.6 Large-Scale Quantitative Analysis of Scaling Factor Similarities

In the previous section, we conduct a case quantitative analysis of the similarity of scaling factors between low-rank and full-rank training across 10 weight matrices of the LLaMA 60M model.

To further substantiate our findings, we expand the scope of our experiment to include all matrices of LLaMA models ranging from 60M to 1B.
We use the same low-rank setup as in Table 2, and then calculate this similarity between the low-rank and full-rank training using the method described in Appendix A.5.
We train these models and assess the similarity of scaling factors averaged over 10,000 steps using Kendall’s Tau and Spearman Rank correlation coefficients.
Additionally, to evaluate the effectiveness of a column-level 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 matrix for analysis. Specifically, in the LLaMA 1B model, over 10,000 columns are sampled.

**Table 8: Spearman and Kendall correlation coefficients with corresponding p-values at both the Matrix and Column levels for pre-training LLaMA models ranging from 60M to 1B parameters, averaged over 10,000 steps.**
| Size | Matrix Level | Column Level |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Spearman | Kendall | Spearman | Kendall |  |  |  |  |  |
| Coefficient | P-value | Coefficient | P-value | Coefficient | P-value | Coefficient | P-value |  |
| 60M | 0.9972 | 1.70e-62 | 0.9662 | 7.25e-26 | 0.9372 | 0.0 | 0.7942 | 0.0 |
| 130M | 0.9925 | 1.50e-76 | 0.9409 | 8.57e-37 | 0.8698 | 0.0 | 0.6830 | 0.0 |
| 350M | 0.9770 | 2.76e-113 | 0.8848 | 5.36e-65 | 0.9091 | 0.0 | 0.7400 | 0.0 |
| 1B | 0.9469 | 1.04e-83 | 0.8249 | 9.77e-57 | 0.8331 | 0.0 | 0.6513 | 0.0 |

As previously stated, both Spearman and Kendall 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 8, both Spearman and Kendall correlation coefficients indicate a strong positive relationship at the matrix and column levels across all sizes of the LLaMA models. All p-values fall below 0.05, indicating statistically significant correlations. Specifically, the matrix-level Spearman coefficients are consistently high with values close to 1, suggesting a very strong correlation. At the column level, the similarity remains high, with Spearman coefficients exceeding 0.8 and Kendall coefficients surpassing 0.6.

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.

### A.7 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 (Zhang et al., ). 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 4. This arises from 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 (Goodfellow et al., ; Zhang et al., ).

### A.8 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 (Zhao et al., ). 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, we first calculate the total number of optimizer states according to Table 1 and then multiply this total number by 2. Additionally, in the fine-tuning of the LLaMA 7B model, pre-trained weights constitute a substantial proportion of the total memory usage, thereby obscuring the memory differences of other components. Given that this memory usage is nearly identical across all methods, it is excluded from our estimates of memory usage.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Can you provide theoretical analysis explaining why the norm-based scaling effectively approximates full-rank gradient corrections? Is this phenomenon present in any LLM training, or is it conditional?
2. Can the baseline comparisons in pre-training and fine-tuning be kept consistent?
3. As a memory-efficient training method, could you provide comparisons regarding training memory usage? Including more experiments on overhead would be a plus.
4. For the Norm-Growth Limiter, could you provide comparative experiments with other methods that address loss spikes?
soundness: 2
strengths: 1. The observation of the similarity in scaling factors of adaptive optimizers between low-rank and full-rank training is insightful and forms a solid foundation for the proposed method.
2. In Norm-Based Scaling, Fira not only maintains memory efficiency but also enables full-rank updates, giving it an advantage over other methods.
3. The Norm-Growth Limiter is a simple yet effective approach to addressing sudden spikes in training loss.
4. The authors conduct extensive experiments, including ablation studies and scaling analyses, across multiple model sizes and tasks, strengthening the empirical validation of their approach.
summary: The paper proposes a memory-efficient training framework called Fira for the pre-training and fine-tuning of large language models (LLMs). Fira aims to reconcile the memory efficiency of low-rank training with the performance benefits of full-rank updates. Building upon GaLore [1], Fira introduces two key components: Norm-Based Scaling and the Norm-Growth Limiter. In Norm-Based Scaling, Fira leverages the empirical observation that the scaling effect on gradient norms remains similar between low-rank and full-rank training. Thus, it uses scaling factors from low-rank optimizers to approximate gradient corrections in full-rank training. The Norm-Growth Limiter mitigates sudden spikes in training loss by limiting the relative increase in gradient norms between steps, effectively converting sudden spikes into gradual rises. The paper validates Fira in both pre-training by evaluating the perplexity on the C4 dataset and fine-tuning by evaluating performance on the CommonSenseQA tasks across various model sizes (60M, 130M, 350M, 1B).

[1] GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection
weaknesses: 1. Limited theoretical justification. While the empirical observations are compelling, the paper lacks a rigorous theoretical analysis explaining why norm-based scaling effectively approximates full-rank gradient corrections.
2. Inconsistent and incomplete baseline comparisons. The baseline comparisons are not comprehensive enough; more recent memory-efficient fine-tuning methods (such as LISA[2], FLoRA[3] that gradient based memory-efficient training method) should be included for a stronger evaluation. Additionally, the evaluation settings between pre-training and fine-tuning are inconsistent: in pre-training, the comparisons include full-rank, Fira, GaLore, LoRA, and ReLoRA, whereas in fine-tuning, full-rank and ReLoRA are omitted, and instead, prefix tuning, series, and parallel methods are included. This inconsistency leads to some confusion. Since the scope of comparison is memory-efficient training, the experiments should be as consistent as possible.
3. Lack of comparison with existing methods addressing loss spikes. Restricting the magnitude of gradients to address abrupt changes in gradients or loss is a commonly adopted approach. While gradient clipping employs absolute value thresholds, the Norm-Growth Limiter uses a relative threshold. Furthermore, sudden spikes are not unique to memory-efficient or parameter-efficient training; they are common issues even in full training. Several studies have proposed methods based on gradient and layer norms to address loss spikes [4,5,6]; however, the paper does not provide comparisons with these methods.
4. Lack of detailed memory efficiency metrics. As a memory-efficient training method, the paper lacks actual memory usage comparisons. In the comparison with GaLore, the paper emphasizes training performance but overlooks memory efficiency and throughput, which is inconsistent with the experiments provided by GaLore.

[2] LISA: Layerwise Importance Sampling for Memory-Efficient Large Language Model Fine-Tuning
[3] Flora: Low-Rank Adapters Are Secretly Gradient Compressors
[4] What language model to train if you have one million GPU hours
[5] GLM-130b: An open bilingual pre-trained model
[6] 8-bit optimizers via block-wise quantization

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to the weakness
soundness: 3
strengths: - The improvements is significant as shown in Table 2.

- The observations of gradient scaling, that can effectively incorporate the gradient residual is interesting and novel.

- The method is clearly explained and easy to follow.
summary: This work proposes a new memory-efficient training approach built on top of GaLore. It leverages low-rank gradients while appropriately combining the residual gradient components with gradient scaling. Experiments demonstrates FIRA achieves significant improvements comparied with previous baselines in both pre-training and fine-tuning scenarios.
weaknesses: - Intuitively, incorporating the residual gradient can lead to performance improvements, though it is typically not expected to surpass Adam. However, as shown in Table 2, Fira significantly outperforms Adam. Is there any insights why it is even better than adam.

- $\gamma = 1.01$ seems to effectively avoid the gradient spiking, whether the performance is senstive to the choice of $\gamma$.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: Please address the concerns in weakness. I am impressed with the performance of the proposed method. However, I am skeptical that the proposed method works in the way the author suggests, or if it works in any other way. I would be willing to increase my score if the authors address these concerns.
soundness: 3
strengths: The proposed method demonstrates strong empirical performance. While the training efficiency is close to the SOTA method, GaLora, the proposed method significantly outperforms GaLora even using 8X smaller ranking. Moreover, it can even outperform full-rank training.
summary: This paper proposes a new method to achieve efficient training of LLMs by reducing the memory usage for optimizer states. Specifically, it designs a norm-based scaling method that approximates full-rank training by using the scaling impact of low-rank optimizers as substitutes for that of original full-rank optimizers. It also proposes a norm-growth limiter to address the issue of loss spikes.
weaknesses: 1. The authors support the use of low-rank scaling factors to approximate full-rank scaling factors by demonstrating that the order of the average scaling factors of the weight matrix is highly similar. However, similarity in order is not equivalent to similarity in value. Therefore, the validity of this approximation method is questionable.

2. Since the proposed method aims to approximate full-rank training, we would assume that the performance of full-rank training is the upper bound. However, the proposed method even outperforms full-rank training. It is unclear why the proposed method outperforms.

3. If the scaling factor for the full rank gradient can be approximated by a low rank scaling factor, why not just apply the low rank scaling factor to the full rank gradient? This strategy has the same memory usage in the optimizer state as the proposed method. Could the authors conduct experiments to validate this strategy? This would be a good support for the effectiveness of the scaling factor approximation.

Minor: 

4. The phenomenon of loss spiking does not seem to be specific to low-order training, therefore, the proposed norm growth limiter is not intended for low-rank training. 

5. The proposed method is only applicable to adaptive optimizers. I know that Adam is very important for efficient training, so this is not a major concern.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 1
questions: See **Weakness**.
soundness: 2
strengths: 1. Compared to the Galore, this paper proposed to utilize the scaled gradient information (Algorithm 1, line 15) resides in the orthogonal complement of the Galore gradient, which achieves improved training performance. The utilization of the orthogonal complement gradient information is new.
2. The authors proposed a new 'gradient clipping' scheme, i.e. the norm-growth limiter in Eq. (12) to alleviate the spikes arouse in during the Fira training process.
3. The authors show evident outperformance of Fira compared to the standard Galore within a similar parameter amount (Table 1 and Table 2).
summary: This work considers training LLMs with full-rank gradients and weights within the low-rank constraints (i.e. imposing low-rank structures on gradient or weights). The main goal of this work is to improve Galore (a low-rank gradient projection optimizer [1]) by incorporating the orthogonal complement information of the Galore projected gradient, without maintaining its optimizer states (1st and 2nd order momentums).

Specifically, the authors observe that (Figure 3), when training LlaMA-60M on C4 dataset with Galore, the scaling factor (the ratio between the norm of the corrected gradient an original gradient) of each matrix exihibits similar rank orders across different choices of GaLore-rank $r$. This motivates the authors to extend GaLore to Fira, a full-rank weight & gradient optimizer under low-rank constraint that consists of two components: 1) a norm-based scaling method, which replace the full-rank scaling factor with the low-rank counterpart; 2) a norm-growth limiter to smooth the training gradient by restricting the gradient norm's increase. The authors showed that proposed Fira outperforms the full-rank baselines on both pretraining and finetuning datasets across various settings.

[1]. *Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: Memory-efficient llm training by gradient low-rank projection. ICML 2024.*
weaknesses: My mainly concerns lies in the motivation of using the norm-based scaling Eq. (10).  

The authors observed that, the average scaling factors (averaged every 20 steps) of the weight matrices exihibit similar orders across different Galore-rank-$r$ training settings (Figure 3 and Table 7). The authors also mentioned that 'the absolute magnitude of the scaling factors can be regarded as the learning rate to which Adam is not sensitive' (Page 5, lines 266-267).

However, from my perspective, several logical gaps still remain between this observation and the 'utilization of norm-based scaling in Eq. (10) at every training step': 

1. To approximate the effect of $(G_t-P_tR_t)$, according to the anaysis and discussions in Page 5 line 238-240, we want to scale it to approximate $\psi_t(G_t-P_tR_t)$ or $\psi_t(G_t-P_t\psi_t(R_t))$, that is, the full-rank gradient dynamic with optimizer states of the orthogonal complement or the full-gradient. But the similarity of the scaling factor orders of $\phi_t(R_t)$ with different Galore-ranks (Figure 3 and Table 7) does not sufficiently implies $\phi_t(R_t)\approx \phi_t(G_t-P_tR_t)$ or $\phi_t(R_t)\approx \phi_t(G_t-P_t\psi_t(R_t))$. 
2. The scaling factor orders shown in Figure 3 and Table 7 only consider the top-10 weights matrices. Thus, this evidence does not sufficiently imply this 'stable order phenomenon' holds for all optimizable weight matrices. 
3. Even if the order of scaling factors of all matrices are the same, this can not sufficiently imply the magnitude of the scaling factors can be absorbed into the learning rate (as mentioned in Page 5 line 267). This is because different matrices have different scaling factors, and the equivalent learning rates of all matrices are different. 

According to Table 4, it seems that the performance gain mainly comes from this norm-based scaling. Therefore, it is necessary to identify the motivation and mechanism of how this works and what is the essential effect of this scaling effect.

Meta-review: This paper proposes a new strategy for training LLMs, where they use the full rank update but maintain optimizer stats in a low-rank fashion to save memory. Two main techniques: Norm-Based Scaling and the Norm-Growth Limiter are proposed to form the actual updates using low-rank stats. Initially reviewers have concerns about both the empirical comparison to previous methods and the theoretical justification of the proposed method. After the rebuttal and author-reviewer discussions, the former is mostly addressed and most of the reviewers agree that the paper demonstrates good results. However, the later (theoretical justification) was not fully addressed: even after rebuttal and discussions, reviewers still think that the paper requires a major revision to provide a more convincing analysis for the proposed method. Further, some reviewers pointed out that Norm-Growth Limiter is similar to gradient clipping which seems to be also used in other places and not specific to low-rank training. Due to these concerns, we decide to recommend rejection and encourage the authors to improve the justification of their proposed method based on the reviewers’ feedback.

Decision: Reject