PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian

arXiv:2403.03507v2cs.LG

TL;DR

LLM training requires substantial memory, and existing low-rank adaptation can limit full-rank learning or require warmup. GaLore projects gradients into changing low-rank spaces to reduce optimizer memory while retaining full-parameter learning. It maintains comparable training performance, enables LLaMA 7B pre-training on a 24 GB GPU, and has a switching-frequency trade-off.

  • Problem

    LLM training requires memory for parameters, gradients, and optimizer states, while low-rank adaptation can underperform full-rank training or require full-rank warmup.

  • Method

    GaLore projects matrix gradients into low-rank spaces and periodically switches projection subspaces, allowing full-parameter learning with reduced optimizer-state memory.

  • Results

    GaLore maintains comparable or better performance across LLaMA pre-training and GLUE fine-tuning, including comparable performance with less than 10% optimizer-state memory and an 85.89 versus 85.61 GLUE score.

  • Takeaways & Limitations

    GaLore enables LLaMA 7B pre-training from scratch on a single 24 GB GPU without full-rank warmup or costly memory offloading.

  • Takeaways & Limitations

    Subspace-switching frequency is a task-dependent hyperparameter: frequent changes add overhead, while infrequent changes can trap optimization in less important regions.

Abstract

from arXiv · show

Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer, reducing trainable parameters and optimizer states. However, such approaches typically underperform training with full-rank weights in both pre-training and fine-tuning stages since they limit the parameter search to a low-rank subspace and alter the training dynamics, and further, may require full-rank warm start. In this work, we propose Gradient Low-Rank Projection (GaLore), a training strategy that allows full-parameter learning but is more memory-efficient than common low-rank adaptation methods such as LoRA. Our approach reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit GaLore further reduces optimizer memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the first time, the feasibility of pre-training a 7B model on consumer GPUs with 24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or offloading strategies.

1. Introduction

LLM training is memory-intensive because parameters, gradients, and optimizer states are large. GaLore addresses limitations of low-rank adaptation by projecting gradients into low-rank spaces while retaining full-parameter learning.

  • Motivation: 58 GB is required to pre-train LLaMA 7B from scratch with a single batch, exceeding the 24 GB memory of consumer GPUs.The estimate includes 14 GB for parameters, 42 GB for Adam optimizer states and weight gradients, and 2 GB for activations.
  • Limitations of Existing Methods: LoRA reduces trainable parameters and optimizer states by learning low-rank adaptors while freezing the full-rank pre-trained weight.It reparameterizes W as W0 + BA, with rank r much smaller than the matrix dimensions.
  • Limitations of Existing Methods: LoRA may underperform full-rank fine-tuning and require full-rank warmup before pre-training in a low-rank subspace.The cited reasons are that optimal weights may not be low-rank and reparameterization changes gradient training dynamics.
  • GaLore: GaLore projects the gradient matrix into a low-rank form, reducing optimizer-state memory while allowing full-parameter learning.Projection matrices are updated occasionally, such as every 200 iterations, to limit amortized computational overhead.
  • Results: Comparable performance to the full-rank counterpart is achieved when 8-bit GaLore uses less than 10% of optimizer-state memory during LLaMA 7B pre-training.The result combines 8-bit optimizers and layer-wise weight updates on the C4 dataset.
  • Results: 85.89 average GLUE score is achieved by rank-4 GaLore on RoBERTa-Base, exceeding LoRA’s 85.61 score.The comparison concerns fine-tuning pre-trained language models on GLUE tasks.

2. Related Works

Related work studies low-rank parameter subspaces, projected gradient descent, naturally low-rank neural-network gradients, and memory-efficient optimization. GaLore builds on these areas while specializing projection to matrix-structured gradients in multilayer networks.

  • Low-rank adaptation: LoRA reduces memory by maintaining a low-rank weight adaptor for each layer during fine-tuning.Variants extend performance, multitask learning, and memory-footprint improvements.
  • Subspace learning: Subspace learning optimizes model weights within a low-rank parameter subspace and has been applied in meta-learning and continual learning.The cited studies report that learning primarily occurs within a significantly low-dimensional parameter subspace.
  • Projected gradient descent: GaLore differs from traditional projected gradient descent by analyzing matrix-structured gradients that arise naturally in multilayer neural networks.Traditional PGD generally treats the objective as a black-box nonlinear function and studies gradients in vector space.
  • Low-rank gradient: Neural-network gradients are naturally low-rank, a property previously used to reduce communication costs and training memory.The cited applications span gradient communication and memory-footprint reduction.
  • Memory-efficient optimization: Memory-efficient optimization reduces gradient-statistics memory through adaptive-optimizer designs, quantization, and fused backward-update operations.These approaches target optimizer states or weight-gradient storage through different mechanisms.

3. GaLore: Gradient Low-Rank Projection

GaLore exploits low-rank structure in neural-network gradients rather than weights, reducing optimizer-state memory while retaining full-parameter learning. Its projections can be updated during training to accommodate changing gradient subspaces and preserve the original training trajectory under full rank.

  • 3.1. Background: Full-rank training stores parameter weights, gradients, and stateful optimizer statistics, making optimizer memory a major cost.For Adam, the parameter matrix and two optimizer-state matrices together require 3mn memory.
  • 3.1. Background: Low-rank parameter updates reduce trainable parameters, but constraining the weight matrix can prevent reaching a high-rank optimum.The paper also identifies altered gradient dynamics as a limitation of low-rank reparameterization.
  • 3.2. Low-Rank Property of Weight Gradient: Gradient matrices can become low-rank during training for certain gradient forms and associated network architectures, even when weight matrices are not low-rank.Under the stated conditions, the stable rank decreases as one term decays exponentially; in particular, the paper derives sr(G_t) ≤ n/2 for large t.
  • 3.3. Gradient Low-rank Projection (GaLore): GaLore projects G into the compact form P^⊤GQ, applies optimizer statistics there, and projects the resulting update back to the original parameter space.The projection matrices can be obtained from the leading singular vectors of the gradient via SVD.
  • 3.3. Gradient Low-rank Projection (GaLore): Changing projection subspaces lets GaLore learn full-rank weights when gradient principal subspaces evolve during complex optimization.With fixed projections, weights grow only along those subspaces; GaLore switches subspaces without increasing the memory footprint.
  • 3.3. Gradient Low-rank Projection (GaLore): When the projection rank reaches min(m, n) and ρ_t ≡ 1, GaLore follows the exact training trajectory of the original model, unlike simultaneous full-rank LoRA-factor optimization.This distinction arises because GaLore projects updates rather than introducing additive low-rank adaptors.

4. GaLore for Memory-Efficient Training

GaLore enables full-parameter learning by projecting gradients into changing low-rank subspaces, reducing optimizer-state memory while retaining the ability to learn full-rank weights. Its memory-efficient optimization combines low-rank gradient statistics, optional 8-bit optimizers, and per-layer updates.

  • Composition of Low-Rank Subspaces: GaLore switches across low-rank subspaces during training, allowing full-rank weights without increasing the memory footprint.Projection matrices are recomputed from the current gradient when switching subspaces.
  • Composition of Low-Rank Subspaces: A single fixed low-rank subspace may miss changing gradient trajectories, whereas switching subspaces supports successful LLM pre-training.The authors identify multiple low-rank subspaces as key to successful pre-training.
  • Composition of Low-Rank Subspaces: The switching frequency T trades computational overhead and optimization fidelity against the risk of remaining in an unimportant subspace.The authors report that T between 50 and 1000 makes little difference in practice, with SVD overhead below 10%.
  • Memory-Efficient Optimization: GaLore projects gradients into low-rank form so optimizers such as Adam track statistics only in the compact space.The method uses projection matrices and updates the weight in the original space after compact-space optimization.
  • Memory-Efficient Optimization: GaLore requires (mn + mr + 2nr) memory, compared with (mn + 3mr + 3nr) for LoRA.GaLore can merge updates into the base weight and therefore does not store a separate low-rank factorization.
  • Combining with Existing Techniques: GaLore is compatible with 8-bit optimizers and per-layer weight updates to further reduce training memory.Per-layer updates apply weight updates during backpropagation rather than storing all layer gradients until the end.

5. Experiments

Experiments evaluate GaLore across LLaMA pre-training, RoBERTa fine-tuning, optimizers, model sizes, memory usage, and throughput. GaLore maintains comparable performance while reducing memory, including enabling 7B pre-training within 24GB on a single device.

  • Evaluation setup: GaLore is evaluated on LLaMA-based pre-training over C4 and RoBERTa fine-tuning on GLUE tasks.Experiments use model sizes up to 7B for pre-training and compare GaLore with full fine-tuning and LoRA for fine-tuning.
  • Optimizer compatibility: GaLore applies to AdamW, 8-bit Adam, and Adafactor without significantly affecting convergence.On LLaMA 1B for 10K steps, rank 512 reduces memory by up to 62.5% on top of savings from 8-bit Adam or Adafactor.
  • Memory measurement: 22.0G memory enables pre-training LLaMA 7B with a small per-GPU token batch on a single 24GB consumer GPU.The setup excludes activation checkpointing and memory offloading; checkpointing can raise the per-GPU batch size to 4096.
  • Memory measurement: 63.3% total-memory reduction versus the BF16 Adam baseline comes from 65.5% optimizer-state savings and reduced weight-gradient storage.Compared with 8-bit Adam, GaLore reduces 9.6G in optimizer states and 13.5G in weight gradients.
  • Throughput: 1019.63 tokens/second for 8-bit GaLore incurs 17% throughput overhead compared with 8-bit Adam.Disabling per-layer weight updates reaches 1109.38 tokens/second, improving throughput by 8.8%.

6. Ablation Study

The ablations show that GaLore’s convergence depends on subspace-update frequency and rank. Lower rank can modestly slow convergence, allowing memory savings to be exchanged for additional training steps.

  • Subspace updates: Both overly frequent and overly slow subspace changes hurt convergence.For small ranks, more frequent switching avoids optimization in an unsuitable subspace, whereas larger ranks provide more coverage.
  • Rank and training duration: Within a certain rank range, decreasing rank only slightly slows convergence in an approximately linear trend.Rank 128 trained for 80K steps achieves lower loss than rank 512 trained for 20K steps.
  • Rank and training duration: Reducing rank trades memory for computation by allowing longer training within a fixed memory budget.The paper presents this as a way to preserve performance in memory-constrained settings.

7. Conclusion

GaLore is presented as a memory-efficient strategy for full-parameter LLM pre-training and fine-tuning. The conclusion identifies broader model coverage, lower-memory projection matrices, and elastic consumer-hardware training as open problems.

  • Contribution: GaLore reduces optimizer-state memory by up to 65.5% while maintaining efficiency and performance in large-scale LLM pre-training and fine-tuning.The conclusion frames GaLore as a strategy for memory-efficient training rather than a low-rank parameterization of the weights.
  • Open problems: Open problems include applying GaLore to vision transformers and diffusion models, lowering projection-matrix memory, and testing elastic distributed training on low-bandwidth hardware.These are identified as future research directions.
  • Implications: The authors position GaLore as a potential tool for training large-scale models on consumer-grade hardware with limited resources.The stated motivation is to inspire further research on memory-efficient training through gradient low-rank projection.

Impact Statement

The paper motivates memory-efficient LLM training through the costs of optimizer states and relates GaLore to factorized optimization and reversible-network gradient structure. It also identifies limitations in applying related methods and in extending its theory to self-attention.

  • Relation to prior methods: Adafactor factorizes second-order statistics, whereas GaLore exploits low-rank gradient structure and can reduce memory for both first- and second-order statistics.GaLore can also be combined with Adafactor for further memory reduction.
  • Relation to prior methods: LOMO and AdaLOMO reduce optimizer memory through fused backward operations, but their direct applicability to larger-scale pre-training from scratch is constrained by Adafactor instability.The paper specifically notes increased training instabilities for vanilla Adafactor at larger scales.
  • Theoretical analysis: The theoretical analysis derives gradient structure for chained reversible networks and notes that self-attention analysis via JoMA is left for future work.The result is stated for reversible architectures and includes linear layers and reversible activations as examples.

B.2. Gradient becomes low-rank

The analysis models training gradients through structured dynamics and shows that their stable rank can decrease over time, yielding low-rank gradients under stated assumptions.

  • The gradient follows structured dynamics whose vectorized update is governed by a positive semidefinite operator S.The analysis assumes a parametric gradient form and studies vanilla SGD weight updates.
  • The initial gradient decomposes into components in the minimal eigenspace V1 and its orthogonal complement.The component in V1 evolves according to the smallest eigenvalue, while the residual evolves separately.
  • The stable rank of the gradient component aligned with V1 is bounded by n − N′ under the stated rank condition.Here N′ is the rank of the relevant feature set, with N′ < n.
  • When the residual term becomes negligible, the gradient can be represented as a sum of N′ rank-1 matrices and therefore has rank at most N′.This establishes low-rank behavior after sufficiently long training under the theorem’s assumptions.
  • If the minimal eigenspace is one-dimensional with a decomposable eigenvector, the projected gradient has stable rank 1 and becomes rank-1.The projection onto V1 is proportional to the decomposable eigenvector.

B.3. Gradient Low-rank property for Transformers

For Transformer feed-forward networks, the analysis argues that gradient structure becomes increasingly low-rank over time, while fixed GaLore projections converge under explicit regularity and stability conditions.

  • Gradient structure in Transformers: The Transformer analysis studies the project-up FFN weight W and shows its gradient becomes low-rank over time under stationary-gradient assumptions.The project-up weight has dimensions m × n, where m is the embedding dimension and n is the number of hidden nodes.
  • Gradient structure in Transformers: The transformed matrix V := U ⊤W is decomposed row-wise, with row growth rates differing across token indices.Different growth rates produce different row norms and contribute to low-rank structure.
  • Gradient structure in Transformers: Eventually, V becomes rank-1 because one row’s growth diverges while the others remain finite.The result follows from the relative critical times of the row-growth functions.
  • Gradient structure in Transformers: The gradient of V is even more low-rank than V itself because the derivative growth contains α_l in its exponents.The analysis characterizes this as the gradient being “exponentially more low-rank” than the matrix.
  • Convergence of fixed projections: With fixed projection matrices P and Q, GaLore converges when the minimum κ_t exceeds LA + LBLCD^2.Under the theorem’s continuity and boundedness assumptions, the projected residual R_t tends to zero.

C.1. Architecture and Hyperparameters

The experiments use standardized LLaMA training settings and estimate memory from parameters and optimizer states, while Figure 6 compares training trajectories across model sizes.

  • LLaMA architecture and training settings: All LLaMA models use sequence length 256, a batch size of 131K tokens, 10% learning-rate warmup, and cosine annealing to 10% of the initial rate.These settings are shared across the pre-training experiments.
  • LLaMA architecture and training settings: GaLore uses learning rate 0.01, scale factor α = 0.25, and a subspace change frequency of 200 across model sizes.The authors report that GaLore is insensitive to hyperparameters and stable with the same learning rate across sizes.
  • Memory measurement: Memory estimates include BF16 weight parameters and optimizer states, with low-rank methods additionally accounting for their trained low-rank parameters.The estimates are based on the number of original and low-rank parameters rather than direct component measurement.
  • Memory measurement: GaLore retains full-rank weight memory while low-rank methods reduce weight memory, according to the reported parameter-memory estimates.The table separates full-rank, GaLore, low-rank, LoRA, and ReLoRA estimates.
  • Training progression: Across 130M, 350M, 1B, and 7B models, GaLore closely matches the full-rank training trajectory and initially converges slightly faster than LoRA.These comparisons are reported for the training progression shown in Figure 6.

D.2. Fine-Tuning on SQuAD dataset

The fine-tuning evaluations compare GaLore with low-rank methods across several datasets and report that GaLore outperforms LoRA on SQuAD using BERT-Base.

  • SQuAD evaluation: On SQuAD, GaLore and LoRA are evaluated with the pre-trained BERT-Base model using rank 16 for both methods.The evaluation reports both Exact Match and F1 scores.
  • SQuAD evaluation: GaLore outperforms LoRA on both Exact Match and F1 scores in the SQuAD evaluation.
  • Other fine-tuning evaluations: Additional experiments apply rank-128 GaLore and LoRA to Gemma-2b, Phi-2, and LLaMA-7B on OpenAssistant Conversations and Belle-1M.The Belle-1M evaluation reports testing perplexity.
Loading 2403.03507v2…