GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

Introduction

Large Language Models (LLMs) have shown impressive performance across multiple disciplines, including conversational AI and language translation. However, pre-training and fine-tuning LLMs require not only a huge amount of computation but is also memory intensive. The memory requirements include not only billions of trainable parameters, but also their gradients and optimizer states (e.g., gradient momentum and variance in Adam) that can be larger than parameter storage themselves (Raffel et al., 2023; Touvron et al., 2023; Chowdhery et al., 2022). For example, pre-training a LLaMA 7B model from scratch with a single batch size requires at least 58 GB memory (14GB for trainable parameters, 42GB for Adam optimizer states and weight gradients, and 2GB for activationsThe calculation is based on LLaMA architecture, BF16 numerical format, and maximum sequence length of 2048.). This makes the training not feasible on consumer-level GPUs such as NVIDIA RTX 4090 with 24GB memory.

In addition to engineering and system efforts, such as gradient checkpointing (Chen et al., 2016), memory offloading (Rajbhandari et al., 2020), etc., to achieve faster and more efficient distributed training, researchers also seek to develop various optimization techniques to reduce the memory usage during pre-training and fine-tuning.

However, many recent works demonstrate the limitation of such a low-rank reparameterization. For fine-tuning, LoRA is not shown to reach a comparable performance as full-rank fine-tuning (Xia et al., 2024). For pre-training from scratch, it is shown to require a full-rank model training as a warmup (Lialin et al., 2023), before optimizing in the low-rank subspace. There are two possible reasons: (1) the optimal weight matrices may not be low-rank, and (2) the reparameterization changes the gradient training dynamics.

We demonstrate that GaLore works well in both LLM pre-training and fine-tuning. When pre-training LLaMA 7B on C4 dataset, 8-bit GaLore, combined with 8-bit optimizers and layer-wise weight updates techniques, achieves comparable performance to its full-rank counterpart, with less than 10% memory cost of optimizer states.

Notably, for pre-training, GaLore keeps low memory throughout the entire training, without requiring full-rank training warmup like ReLoRA. Thanks to GaLore’s memory efficiency, for the first time it is possible to train LLaMA 7B from scratch on a single GPU with 24GB memory (e.g., on NVIDIA RTX 4090), without any costly memory offloading techniques (Fig. 1).

GaLore is also used to fine-tune pre-trained LLMs on GLUE benchmarks with comparable or better results than existing low-rank methods. When fine-tuning RoBERTa-Base on GLUE tasks with a rank of 4, GaLore achieves an average score of 85.89, outperforming LoRA, which achieves a score of 85.61.

As a gradient projection method, GaLore is independent of the choice of optimizers and can be easily plugged into existing ones with only two lines of code, as shown in Algorithm 1. Our experiment (Fig. 3) shows that it works for popular optimizers such as AdamW, 8-bit Adam, and Adafactor. In addition, its performance is insensitive to very few hyper-parameters it introduces. We also provide theoretical justification on the low-rankness of gradient update, as well as the convergence analysis of GaLore.

Related Works

Hu et al. (2021) proposed Low-Rank Adaptation (LoRA) to fine-tune pre-trained models with low-rank adaptors. This method reduces the memory footprint by maintaining a low-rank weight adaptor for each layer. There are a few variants of LoRA proposed to enhance its performance (Renduchintala et al., 2023; Sheng et al., 2023; Xia et al., 2024), supporting multi-task learning (Wang et al., 2023), and further reducing the memory footprint (Dettmers et al., 2023). Lialin et al. (2023) proposed ReLoRA, a variant of LoRA designed for pre-training, but requires a full-rank training warmup to achieve comparable performance as the standard baseline.

Subspace Learning

Recent studies have demonstrated that the learning primarily occurs within a significantly low-dimensional parameter subspace (Larsen et al., 2022; Gur-Ari et al., 2018). These findings promote a special type of learning called subspace learning, where the model weights are optimized within a low-rank subspace. This notion has been widely used in different domains of machine learning, including meta-learning and continual learning (Lee & Choi, 2018; Chaudhry et al., 2020).

Projected Gradient Descent

GaLore is closely related to the traditional topic of projected gradient descent (PGD) (Chen & Wainwright, 2015; Chen et al., 2019). A key difference is that, GaLore considers the specific gradient form that naturally appears in training multi-layer neural networks (e.g., it is a matrix with specific structures), proving many of its properties (e.g., Lemma 3.1, Theorem 3.2, and Theorem 3.6). In contrast, traditional PGD mostly treats the objective as a general blackbox nonlinear function, and study the gradients in the vector space only.

Memory-Efficient Optimization

There have been some works trying to reduce the memory cost of gradient statistics for adaptive optimization algorithms (Shazeer & Stern, ; Anil et al., ; Dettmers et al., 2021). Adafactor (Shazeer & Stern, ) achieves sub-linear memory cost by factorizing the second-order statistics by a row-column outer product. GaLore shares similarities with Adafactor in terms of utilizing low-rank factorization to reduce memory cost, but GaLore focuses on the low-rank structure of the gradients, while Adafactor focuses on the low-rank structure of the second-order statistics. GaLore can reduce the memory cost for both first-order and second-order statistics, and can be combined with Adafactor to achieve further memory reduction. Quantization is also widely used to reduce the memory cost of optimizer states (Dettmers et al., 2021; Li et al., 2023). Furthermore, Lv et al. (2023) proposed fused gradient computation to reduce the memory cost of storing weight gradients during training.

In contrast to the previous memory-efficient optimization methods, GaLore operates independently as the optimizers directly receive the low-rank gradients without knowing their full-rank counterparts.

GaLore: Gradient Low-Rank Projection

2 Low-Rank Property of Weight Gradient

While low-rank updates are proposed to reduce memory usage, it remains an open question whether the weight matrix should be parameterized as low-rank. In many situations, this may not be true. For example, in linear regression ${\bm{y}}=W{\bm{x}}$, if the optimal $W^{*}$ is high-rank, then imposing a low-rank assumption on $W$ never leads to the optimal solution, regardless of what optimizers are used.

Surprisingly, while the weight matrices are not necessarily low-rank, the gradient indeed becomes low-rank during the training for certain gradient forms and associated network architectures:

Let $m\leq n$ without loss of generality. The gradient update:

with constant $A$ and PSD matrices $B$ and $C$ and randomly initialized $W_{0}$ leads to low-rank gradient with high probability:

Here $\nu_{1}=\lambda_{\min}(C)$ is the smallest eigenvalues of $C$ and $\lambda_{1}\leq\ldots\leq\lambda_{n}$ are eigenvalues of $B$. Furthermore, if $\lambda_{2}>\lambda_{1}$ and $\nu_{1}>0$, then $G_{t}$ converges to rank-$1$ exponentially.

Note that in Lemma 3.1, we assume a parametric form (Eqn. 6) of the gradient. This is not a limiting assumption. It not only holds for simple linear network with objective $\varphi(W)=\|{\bm{y}}-W{\bm{x}}\|^{2}_{2}$, but also hold in more general nonlinear networks known as “reversible networks” (Tian et al., 2020), including deep ReLU networks:

Note that for softmax objective with small logits, we can also prove a similar structure of backpropagated gradient, and thus Theorem 3.2 can also apply.

For $K$-way logsoftmax loss $\varphi({\bm{y}};{\bm{f}}):=-\log\left(\frac{\exp({\bm{y}}^{\top}{\bm{f}})}{{\bm{1}}^{\top}\exp({\bm{f}})}\right)$, let $\hat{\bm{f}}=P^{\perp}_{\bm{1}}{\bm{f}}$ be the zero-mean version of network output ${\bm{f}}$, where $P^{\perp}_{\bm{1}}:=I-\frac{1}{K}{\bm{1}}{\bm{1}}^{\top}$, then we have:

where $\gamma({\bm{y}},{\bm{f}})\approx 1$ and ${\bm{y}}$ is a data label with ${\bm{y}}^{\top}{\bm{1}}=1$.

With this lemma, it is clear that for a reversible network ${\bm{f}}:=\mathcal{N}({\bm{x}})=J_{l}({\bm{x}})W_{l}{\bm{f}}_{l-1}({\bm{x}})$, the gradient $G_{l}$ of $W_{l}$ has the following form:

which is consistent with the form $G_{l}=A-BW_{l}C$. For a detailed introduction to reversibility, please check the Appendix A.2.

3 Gradient Low-rank Projection (GaLore)

Since the gradient $G$ may have a low-rank structure, if we can keep the gradient statistics of a small “core” of gradient $G$ in optimizer states, rather than $G$ itself, then the memory consumption can be reduced substantially. This leads to our proposed GaLore strategy:

Gradient low-rank projection (GaLore) denotes the following gradient update rules ($\eta$ is the learning rate):

Different from LoRA, GaLore explicitly utilizes the low-rank updates instead of introducing additional low-rank adaptors and hence does not alter the training dynamics.

In the following, we show that GaLore converges under a similar (but more general) form of gradient update rule (Eqn. 6). This form corresponds to Eqn. 8 but with a larger batch size.

A function ${\bm{h}}(W)$ has (Lipschitz) $L$-continuity, if for any $W_{1}$ and $W_{2}$, $\|{\bm{h}}(W_{1})-{\bm{h}}(W_{2})\|_{F}\leq L\|W_{1}-W_{2}\|_{F}$.

Suppose the gradient has the following form (Eqn. 8 with batchsize $>1$):

where $B_{i}$ and $C_{i}$ are PSD matrices, $A_{i}$, $B_{i}$ and $C_{i}$ have $L_{A}$, $L_{B}$ and $L_{C}$ continuity with respect to $W$ and $\|W_{t}\|\leq D$. Let $R_{t}:=P_{t}^{\top}G_{t}Q_{t}$, $\hat{B}_{it}:=P_{t}^{\top}B_{i}(W_{t})P_{t}$, $\hat{C}_{it}:=Q_{t}^{\top}C_{i}(W_{t})Q_{t}$ and $\kappa_{t}:=\frac{1}{N}\sum_{i}\lambda_{\min}(\hat{B}_{it})\lambda_{\min}(\hat{C}_{it})$. If we choose constant $P_{t}=P$ and $Q_{t}=Q$, then GaLore with $\rho_{t}\equiv 1$ satisfies:

As a result, if $\min_{t}\kappa_{t}>L_{A}+L_{B}L_{C}D^{2}$, $R_{t}\rightarrow 0$ and thus GaLore converges with fixed $P_{t}$ and $Q_{t}$.

Setting $P$ and $Q$. The theorem tells that $P$ and $Q$ should project into the subspaces corresponding to the first few largest eigenvectors of $\hat{B}_{it}$ and $\hat{C}_{it}$ for faster convergence (large $\kappa_{t}$). While all eigenvalues of the positive semidefinite (PSD) matrix $B$ and $C$ are non-negative, some of them can be very small and hinder convergence (i.e., it takes a long time for $G_{t}$ to become ). With the projection $P$ and $Q$, $P^{\top}B_{it}P$ and $Q^{\top}C_{it}Q$ only contain the largest eigen subspaces of $B$ and $C$, improving the convergence of $R_{t}$ and at the same time, reduces the memory usage.

While it is tricky to obtain the eigenstructure of $\hat{B}_{it}$ and $\hat{C}_{it}$ (they are parts of Jacobian), one way is to instead use the spectrum of $G_{t}$ via Singular Value Decomposition (SVD):

GaLore for Memory-Efficient Training

For a complex optimization problem such as LLM pre-training, it may be difficult to capture the entire gradient trajectory with a single low-rank subspace. One reason is that the principal subspaces of $B_{t}$ and $C_{t}$ (and thus $G_{t}$) may change over time. In fact, if we keep the same projection $P$ and $Q$, then the learned weights will only grow along these subspaces, which is not longer full-parameter training. Fortunately, for this, GaLore can switch subspaces during training and learn full-rank weights without increasing the memory footprint.

We allow GaLore to switch across low-rank subspaces:

Following the above procedure, the switching frequency $T$ becomes a hyperparameter. The ablation study (Fig. 5) shows a sweet spot exists. A very frequent subspace change increases the overhead (since new $P_{t}$ and $Q_{t}$ need to be computed) and breaks the condition of constant projection in Theorem 3.6. In practice, it may also impact the fidelity of the optimizer states, which accumulate over multiple training steps. On the other hand, a less frequent change may make the algorithm stuck into a region that is no longer important to optimize (convergence proof in Theorem 3.6 only means good progress in the designated subspace, but does not mean good overall performance). While optimal $T$ depends on the total training iterations and task complexity, we find that a value between $T=50$ to $T=1000$ makes no significant difference. Thus, the total computational overhead induced by SVD is negligible ($<10\%$) compared to other memory-efficient training techniques such as memory offloading (Rajbhandari et al., 2020).

2 Memory-Efficient Optimization

GaLore can also apply to other optimizers (e.g., Adafactor) that have similar update rules and require a large amount of memory to store gradient statistics.

To achieve the best memory-performance trade-off, we only use one project matrix $P$ or $Q$, projecting the gradient $G$ into $P^{\top}G$ if $m\leq n$ and $GQ$ otherwise. We present the algorithm applying GaLore to Adam in Algorithm 2.

With this setting, GaLore requires less memory than LoRA during training. As GaLore can always merge $\Delta W_{t}$ to $W_{0}$ during weight updates, it does not need to store a separate low-rank factorization $BA$. In total, GaLore requires $(mn+mr+2nr)$ memory, while LoRA requires $(mn+3mr+3nr)$ memory. A comparison between GaLore and LoRA is shown in Table 1.

As Theorem 3.6 does not require the projection matrix to be carefully calibrated, we can further reduce the memory cost of projection matrices by quantization and efficient parameterization, which we leave for future work.

3 Combining with Existing Techniques

GaLore is compatible with existing memory-efficient optimization techniques. In our work, we mainly consider applying GaLore with 8-bit optimizers (Dettmers et al., 2021) and per-layer weight updates (Lv et al., 2023).

Dettmers et al. (2022) proposed 8-bit Adam optimizer that maintains 32-bit optimizer performance at a fraction of the original memory footprint. We apply GaLore directly to the existing implementation of 8-bit Adam.

Per-layer weight updates.

In practice, the optimizer typically performs a single weight update for all layers after backpropagation. This is done by storing the entire weight gradients in memory. To further reduce the memory footprint during training, we adopt per-layer weight updates to GaLore, which performs the weight updates during backpropagation (Lv et al., 2023).

4 Hyperparameters of GaLore

In addition to Adam’s original hyperparameters, GaLore only introduces very few additional hyperparameters: the rank $r$ which is also present in LoRA, the subspace change frequency $T$ (see Sec. 4.1), and the scale factor $\alpha$.

Scale factor $\alpha$ controls the strength of the low-rank update, which is similar to the scale factor $\alpha/r$ appended to the low-rank adaptor in Hu et al. (2021). We note that the $\alpha$ does not depend on the rank $r$ in our case. This is because, when $r$ is small during pre-training, $\alpha/r$ significantly affects the convergence rate, unlike fine-tuning.

Experiments

We evaluate GaLore on both pre-training and fine-tuning of LLMs. All experiments are conducted on NVIDIA A100 GPUsThe implementation of GaLore is available here.

To evaluate its performance, we apply GaLore to train LLaMA-based large language models on the C4 dataset. C4 dataset is a colossal, cleaned version of Common Crawl’s web crawl corpus, which is mainly intended to pre-train language models and word representations (Raffel et al., 2023). To best simulate the practical pre-training scenario, we train without data repetition over a sufficiently large amount of data, across a range of model sizes up to 7 Billion parameters.

Architecture and hyperparameters.

We follow the experiment setup from Lialin et al. (2023), which adopts a LLaMA-basedLLaMA materials in our paper are subject to LLaMA community license. architecture with RMSNorm and SwiGLU activations (Touvron et al., 2023; Zhang & Sennrich, 2019; Shazeer, 2020). For each model size, we use the same set of hyperparameters across methods, except the learning rate. We run all experiments with BF16 format to reduce memory usage, and we tune the learning rate for each method under the same amount of computational budget and report the best performance. The details of our task setups and hyperparameters are provided in the appendix.

Fine-tuning on GLUE tasks.

GLUE is a benchmark for evaluating the performance of NLP models on a variety of tasks, including sentiment analysis, question answering, and textual entailment (Wang et al., 2019). We use GLUE tasks to benchmark GaLore against LoRA for memory-efficient fine-tuning.

1 Comparison with low-rank methods

We first compare GaLore with existing low-rank methods using Adam optimizer across a range of model sizes.

Our baseline method that applies Adam optimizer with full-rank weights and optimizer states.

Low-Rank

We also evaluate a traditional low-rank approach that represents the weights by learnable low-rank factorization: $W=BA$ (Kamalakara et al., 2022).

LoRA

Hu et al. (2021) proposed LoRA to fine-tune pre-trained models with low-rank adaptors: $W=W_{0}+BA$, where $W_{0}$ is fixed initial weights and $BA$ is a learnable low-rank adaptor. In the case of pre-training, $W_{0}$ is the full-rank initialization matrix. We set LoRA alpha to 32 and LoRA dropout to 0.05 as their default settings.

ReLoRA

Lialin et al. (2023) is a variant of LoRA designed for pre-training, which periodically merges $BA$ into $W$, and initializes new $BA$ with a reset on optimizer states and learning rate. ReLoRA requires careful tuning of merging frequency, learning rate reset, and optimizer states reset. We evaluate ReLoRA without a full-rank training warmup for a fair comparison.

For GaLore, we set subspace frequency $T$ to 200 and scale factor $\alpha$ to 0.25 across all model sizes in Table 2. For each model size, we pick the same rank $r$ for all low-rank methods, and we apply them to all multi-head attention layers and feed-forward layers in the models. We train all models using Adam optimizer with the default hyperparameters (e.g., $\beta_{1}=0.9$, $\beta_{2}=0.999$, $\epsilon=10^{-8}$). We also estimate the memory usage based on BF16 format, including the memory for weight parameters and optimizer states. As shown in Table 2, GaLore outperforms other low-rank methods and achieves comparable performance to full-rank training. We note that for 1B model size, GaLore even outperforms full-rank baseline when $r=1024$ instead of $r=512$. Compared to LoRA and ReLoRA, GaLore requires less memory for storing model parameters and optimizer states. A detailed training setting of each model and our memory estimation for each method are provided in the appendix.

2 GaLore with Memory-Efficient Optimizers

We demonstrate that GaLore can be applied to various learning algorithms, especially memory-efficient optimizers, to further reduce the memory footprint. We apply GaLore to AdamW, 8-bit Adam, and Adafactor optimizers (Loshchilov & Hutter, 2019; Dettmers et al., 2022; Shazeer & Stern, ). We consider Adafactor with first-order statistics to avoid performance degradation.

We evaluate them on LLaMA 1B architecture with 10K training steps, and we tune the learning rate for each setting and report the best performance. As shown in Fig. 3, applying GaLore does not significantly affect their convergence. By using GaLore with a rank of 512, the memory footprint is reduced by up to 62.5%, on top of the memory savings from using 8-bit Adam or Adafactor optimizer. Since 8-bit Adam requires less memory than others, we denote 8-bit GaLore as GaLore with 8-bit Adam, and use it as the default method for the following experiments on 7B model pre-training and memory measurement.

3 Scaling up to LLaMA 7B Architecture

Scaling ability to 7B models is a key factor for demonstrating if GaLore is effective for practical LLM pre-training scenarios. We evaluate GaLore on an LLaMA 7B architecture with an embedding size of 4096 and total layers of 32. We train the model for 150K steps with 19.7B tokens, using 8-node training in parallel with a total of 64 A100 GPUs. Due to computational constraints, we only compare 8-bit GaLore ($r=1024$) with 8-bit Adam with a single trial without tuning the hyperparameters. As shown in Table 3, after 150K steps, 8-bit GaLore achieves a perplexity of 14.65, which is comparable to 8-bit Adam with a perplexity of 14.61.

4 Memory-Efficient Fine-Tuning

GaLore not only achieves memory-efficient pre-training but also can be used for memory-efficient fine-tuning. We fine-tune pre-trained RoBERTa models on GLUE tasks using GaLore and compare its performance with a full fine-tuning baseline and LoRA. We use hyperparameters from Hu et al. (2021) for LoRA and tune the learning rate and scale factor for GaLore. As shown in Table 4, GaLore achieves better performance than LoRA on most tasks with less memory footprint. This demonstrates that GaLore can serve as a full-stack memory-efficient training strategy for both LLM pre-training and fine-tuning.

5 Measurement of Memory and Throughput

While Table 2 gives the theoretical benefit of GaLore compared to other methods in terms of memory usage, we also measure the actual memory footprint of training LLaMA models by various methods, with a token batch size of 256. The training is conducted on a single device setup without activation checkpointing, memory offloading, and optimizer states partitioning (Rajbhandari et al., 2020).

Training 7B models on consumer GPUs with 24G memory. As shown in Fig. 4, 8-bit GaLore requires significantly less memory than BF16 baseline and 8-bit Adam, and only requires 22.0G memory to pre-train LLaMA 7B with a small per-GPU token batch size (up to 500 tokens). This memory footprint is within 24GB VRAM capacity of a single GPU such as NVIDIA RTX 4090. In addition, when activation checkpointing is enabled, per-GPU token batch size can be increased up to 4096. While the batch size is small per GPU, it can be scaled up with data parallelism, which requires much lower bandwidth for inter-GPU communication, compared to model parallelism. Therefore, it is possible that GaLore can be used for elastic training (Lin et al., ) 7B models on consumer GPUs such as RTX 4090s.

Specifically, we present the memory breakdown in Fig. 1. It shows that 8-bit GaLore reduces 37.92G (63.3%) and 24.5G (52.3%) total memory compared to BF16 Adam baseline and 8-bit Adam, respectively. Compared to 8-bit Adam, 8-bit GaLore mainly reduces the memory in two parts: (1) low-rank gradient projection reduces 9.6G (65.5%) memory of storing optimizer states, and (2) using per-layer weight updates reduces 13.5G memory of storing weight gradients.

Throughput overhead of GaLore. We also measure the throughput of the pre-training LLaMA 1B model with 8-bit GaLore and other methods, where the results can be found in the appendix. Particularly, the current implementation of 8-bit GaLore achieves 1019.63 tokens/second, which induces 17% overhead compared to 8-bit Adam implementation. Disabling per-layer weight updates for GaLore achieves 1109.38 tokens/second, improving the throughput by 8.8%. We note that our results do not require offloading strategies or checkpointing, which can significantly impact training throughput. We leave optimizing the efficiency of GaLore implementation for future work.

Ablation Study

We observe that both too frequent and too slow changes of subspaces hurt the convergence, as shown in Fig. 5(left). The reason has been discussed in Sec. 4.1 and is more prevalent for small $r$, since in such case, the subspace switching should happen at the right time to avoid wasting optimization steps in the wrong subspace, while for large $r$ the gradient updates cover more subspaces, providing more cushion.

2 How does the rank of subspace affect the convergence?

Within a certain range of rank values, decreasing the rank only slightly affects the convergence rate, causing a slowdown that is close to linear. As shown in Fig. 5(right), training with a rank of 128 using 80K steps achieves a lower loss than training with a rank of 512 using 20K steps. This shows that GaLore can be used to trade-off between memory and computational cost. In a memory-constrained scenario, reducing the rank allows us to stay within the memory budget while training for more steps to preserve the performance.

Conclusion

We propose GaLore, a memory-efficient pre-training and fine-tuning strategy for large language models. GaLore significantly reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for large-scale LLM pre-training and fine-tuning.

We identify several open problems for GaLore, which include (1) applying GaLore on training of other types of models such as vision transformers and diffusion models, (2) further improving memory efficiency by employing low-memory projection matrices, through quantization or special parameterization, and (3) exploring the possibility of elastic data distributed training on low-bandwidth consumer-grade hardware.

We hope that our work will inspire future research on memory-efficient LLM training strategies from the perspective of low-rank gradient projection. We believe that GaLore will be a valuable tool for the community to train large language models with consumer-grade hardware and limited resources.

Impact Statement

This paper aims to improve the memory efficiency of training large language models (LLMs) in order to reduce the environmental impact of LLM pre-training and fine-tuning. By enabling the training of larger models on hardware with lower memory, our approach helps to minimize energy consumption and carbon footprint associated with training LLMs.

References

Appendix A Proofs

Suppose $h_{t,ij}$ is the $ij$ component of $H_{t}$, then from the equation above we have:

To make it more precise, consider the stable rank:

With high probability, $h^{2}_{0,1j}\geq\epsilon^{2}_{0}$, since $|h^{2}_{1i}|\leq c_{0}$ is bounded, we have:

Using Mediant inequality, $\frac{a}{b}\leq\frac{a+c}{b+d}\leq\frac{c}{d}$ for $a,b,c,d>0$, therefore, we know that for $i$-th row ($i\geq 2$), since $\lambda_{i}\geq\lambda_{1}$:

A.2 Reversibility

A network $\mathcal{N}$ that maps input ${\bm{x}}$ to output ${\bm{y}}=\mathcal{N}({\bm{x}})$ is reversible, if there exists $K({\bm{x}};W)$ so that ${\bm{y}}=K({\bm{x}};W){\bm{x}}$, and the backpropagated gradient ${\bm{g}}_{\bm{x}}$ satisfies ${\bm{g}}_{\bm{x}}=K^{\top}({\bm{x}};W){\bm{g}}_{\bm{y}}$, where ${\bm{g}}_{\bm{y}}$ is the backpropagated gradient at the output ${\bm{y}}$. Here $K({\bm{x}};W)$ depends on the input ${\bm{x}}$ and weight $W$ in the network $\mathcal{N}$.

Note that many layers are reversible, including linear layer (without bias), reversible activations (e.g., ReLU, leaky ReLU, polynomials, etc). Furthermore, they can be combined to construct more complicated architectures:

If $\mathcal{N}_{1}$ and $\mathcal{N}_{2}$ are reversible networks, then (Parallel) ${\bm{y}}=\alpha_{1}\mathcal{N}_{1}({\bm{x}})+\alpha_{2}\mathcal{N}_{2}({\bm{x}})$ is reversible for constants $\alpha_{1}$ and $\alpha_{2}$, and (Composition) ${\bm{y}}=\mathcal{N}_{2}(\mathcal{N}_{1}({\bm{x}}))$ is reversible.

From this property, it is clear that ResNet architecture ${\bm{x}}+\mathcal{N}({\bm{x}})$ is reversible, if $\mathcal{N}$ contains bias-free linear layers and reversible activations, which is often the case in practice. For a detailed analysis, please check Appendix A in (Tian et al., 2020). For architectures like self-attention, one possibility is to leverage JoMA (Tian et al., 2024) to analyze, and we leave for future work.

The gradient of chained reversible networks has the following structure: See 3.2

Note that for layered reversible network, we have

Let ${\bm{f}}_{l}:=\mathcal{N}_{l}(\mathcal{N}_{l-1}(\ldots\mathcal{N}_{1}({\bm{x}})))$ and $J_{l}:=K_{L}({\bm{x}})\ldots K_{l+1}({\bm{x}})$, and for linear layer $l$, we can write $\mathcal{N}({\bm{x}})=J_{l}W_{l}{\bm{f}}_{l-1}$. Therefore, for the linear layer $l$ with weight matrix $W_{l}$, we have:

Let $\hat{\bm{f}}:=P^{\perp}_{\bm{1}}{\bm{f}}$ be the zero-mean version of network output ${\bm{f}}$. Then we have ${\bm{1}}^{\top}\hat{\bm{f}}=0$ and ${\bm{f}}=\hat{\bm{f}}+c{\bm{1}}$. Therefore, we have:

Using the Taylor expansion $\exp(x)=1+x+\frac{x^{2}}{2}+o(x^{2})$, we have:

where $\gamma:=(1+\hat{\bm{f}}^{\top}\hat{\bm{f}}/2K+o(\hat{\bm{f}}^{2}/K))^{-1}\approx 1$. ∎

A.3 Convergence of GaLore

Using the same notation, it is clear to show that:

Then we derive the recursive update rule for $g_{t}$:

where $e_{t}:=(a_{t}-a_{t-1})+(S_{t-1}-S_{t})w_{t}$. Left multiplying by $(Q\otimes P)^{\top}$, we have:

Now we bound the norm. Note that since $P$ and $Q$ are projection matrices with $P^{\top}P=I$ and $Q^{\top}Q=I$, we have:

where $E_{t}:=\frac{1}{N}\sum_{i}(A_{it}-A_{i,t-1})+\frac{1}{N}\sum_{i}(B_{i,t-1}W_{t}C_{i,t-1}-B_{it}W_{t}C_{it})$. So we only need to bound $\|E_{t}\|_{F}$. Note that:

Now we estimate the minimal eigenvalue of $\hat{S}_{t-1}$. Let $\underline{\lambda}_{it}:=\lambda_{\min}(P^{\top}B_{it}P)$ and $\underline{\nu}_{it}:=\lambda_{\min}(Q^{\top}C_{it}Q)$, then $\lambda_{\min}((P^{\top}B_{it}P)\otimes(Q^{\top}C_{it}Q))=\underline{\lambda}_{it}\underline{\nu}_{it}$ and for any unit vector ${\bm{v}}$:

And thus $\lambda_{\min}(\hat{S}_{t})\geq\frac{1}{N}\sum_{i}\underline{\lambda}_{it}\underline{\nu}_{it}$. Therefore, $\lambda_{\max}(I-\eta\hat{S}_{t-1})\leq 1-\frac{\eta}{N}\sum_{i}\underline{\lambda}_{i,t-1}\underline{\nu}_{i,t-1}$. Therefore, let $\kappa_{t}:=\frac{1}{N}\sum_{i}\underline{\lambda}_{it}\underline{\nu}_{it}$ and using the fact that $\|r_{t}\|_{2}=\|R_{t}\|_{F}$, we have:

Appendix B Details of Pre-Training Experiment

We introduce details of the LLaMA architecture and hyperparameters used for pre-training. Table 5 shows the most hyperparameters of LLaMA models across model sizes. We use a max sequence length of 256 for all models, with a batch size of 131K tokens. For all experiments, we adopt learning rate warmup for the first 10% of the training steps, and use cosine annealing for the learning rate schedule, decaying to 10% of the initial learning rate.

For all methods on each size of models (from 60M to 1B), we tune their favorite learning rate from a set of $\{0.01,0.005,0.001,0.0005,0.0001\}$, and the best learning rate is chosen based on the validation perplexity. We find GaLore is insensitive to hyperparameters and tends to be stable with the same learning rate across different model sizes. For all models, GaLore use the same hyperparameters, including the learning rate of $0.01$, scale factor $\alpha$ of $0.25$, and the subspace change frequency of $T$ of $200$. We note that since $\alpha$ can be viewed as a fractional learning rate, most of the modules (e.g., multi-head attention and feed-forward layers) in LLaMA models have the actual learning rate of $0.0025$. This is, still, a relatively large stable learning rate compared to the full-rank baseline, which usually uses a learning rate $\leq 0.001$ to avoid spikes in the training loss.

B.2 Memory Estimates

As the GPU memory usage for a specific component is hard to measure directly, we estimate the memory usage of the weight parameters and optimizer states for each method on different model sizes. The estimation is based on the number of original parameters and the number of low-rank parameters, trained by BF16 format. For example, for a 60M model, LoRA ($r=128$) requires $42.7$M parameters on low-rank adaptors and $60M$ parameters on the original weights, resulting in a memory cost of $0.20$G for weight parameters and $0.17$G for optimizer states. Table 6 shows the memory estimates for weight parameters and optimizer states for different methods on different model sizes, as a compliment to the total memory reported in the main text.

Appendix C Details of Fine-Tuning Experiment

We fine-tune the pre-trained RoBERTa-Base model on the GLUE benchmark using the model provided by the Hugging Facehttps://huggingface.co/transformers/model_doc/roberta.html. We trained the model for 30 epochs with a batch size of 16 for all tasks except for CoLA, which uses a batch size of 32. We tune the learning rate and scale factor for GaLore. Table 7 shows the hyperparameters used for fine-tuning RoBERTa-Base for GaLore.

Appendix D Additional Memory Measurements

We empirically measure the memory usage of different methods for pre-training LLaMA 1B model on C4 dataset with a token batch size of 256, as shown in Table 8.