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

• URL: http://arxiv.org/abs/2410.01623v2
• Date: Sat, 12 Oct 2024 08:44:01 GMT
• Title: Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?
• Authors: Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, Guoren Wang,
• Abstract summary: Low-rank training has emerged as a promising approach for reducing memory usage in training Large Language Models. We propose a new plug-and-play training framework for LLMs called Fira, as the first attempt to achieve this goal. We show that Fira outperforms both LoRA and GaLore, achieving performance that is comparable to or even better than full-rank training.
• Score: 40.94505326255136
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• AdaRankGrad: Adaptive Gradient-Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning [9.51289606759621]
  Training and fine-tuning large language models (LLMs) come with challenges related to memory and computational requirements. Various techniques have been developed to tackle these challenges, such as low-rank adaptation (LoRA) We introduce a new method inspired by a phenomenon we formally prove: as training progresses, the rank of the estimated gradient gradually decreases.
  arXiv  Detail & Related papers  (2024-10-23T13:53:26Z)
• SubZero: Random Subspace Zeroth-Order Optimization for Memory-Efficient LLM Fine-Tuning [66.27334633749734]
  As language models grow in size, memory demands for backpropagation increase. Zeroth-order (ZOZO) optimization methods offer a memory-efficient alternative. We show that SubZero enhances fine-tuning and achieves faster results compared to standard ZOZO approaches.
  arXiv  Detail & Related papers  (2024-10-11T17:01:43Z)
• From GaLore to WeLore: How Low-Rank Weights Non-uniformly Emerge from Low-Rank Gradients [86.40635601953446]
  We study the emergence of low-rank structures across different layers of Modern Large Language Models. We present Weight Low-Rank Projection (WeLore) that unifies weight compression and memory-efficient fine-tuning as ONE.
  arXiv  Detail & Related papers  (2024-07-15T21:05:20Z)
• BlockLLM: Memory-Efficient Adaptation of LLMs by Selecting and Optimizing the Right Coordinate Blocks [19.007090250576585]
  BlockLLM is an approach inspired by block coordinate descent. It achieves state-of-the-art performance in both finetuning and pretraining tasks.
  arXiv  Detail & Related papers  (2024-06-25T05:45:12Z)
• SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining [39.56934385513862]
  Training large language models (LLMs) from scratch requires significant computational power and extensive memory capacity. Recent studies have explored low-rank structures on weights for efficient fine-tuning in terms of parameters and memory. We propose to parameterize the weights as a sum of low-rank and sparse matrices for pretraining, which we call SLTrain.
  arXiv  Detail & Related papers  (2024-06-04T11:14:21Z)
• GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection [133.45193150403537]
  Training Large Language Models (LLMs) presents significant memory challenges due to the growing size of weights and GPU states. In this work, we propose Gradient Low-Rank Projection (GaLore) as a memory-efficient training strategy. Our 8-bit GaLore further reduces memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline.
  arXiv  Detail & Related papers  (2024-03-06T07:29:57Z)
• PRILoRA: Pruned and Rank-Increasing Low-Rank Adaptation [65.268245109828]
  We introduce PRILoRA, which linearly allocates a different rank for each layer, in an increasing manner, and performs pruning throughout the training process. We validate the effectiveness of PRILoRA through extensive experiments on eight GLUE benchmarks, setting a new state of the art.
  arXiv  Detail & Related papers  (2024-01-20T20:25:17Z)
• InRank: Incremental Low-Rank Learning [85.6380047359139]
  gradient-based training implicitly regularizes neural networks towards low-rank solutions through a gradual increase of the rank during training. Existing training algorithms do not exploit the low-rank property to improve computational efficiency. We design a new training algorithm Incremental Low-Rank Learning (InRank), which explicitly expresses cumulative weight updates as low-rank matrices.
  arXiv  Detail & Related papers  (2023-06-20T03:03:04Z)
• In-Loop Meta-Learning with Gradient-Alignment Reward [34.1954698584925]
  We present a cheap-to-compute and memory-saving reward, the gradient-alignment reward (GAR), that can guide the optimization. First, we present the application of GAR to choosing the data distribution as a mixture of multiple dataset splits in a small scale setting. Second, we show that it can successfully guide learning augmentation strategies competitive with state-of-the-art augmentation strategies on CIFAR-10 and CIFAR-100.
  arXiv  Detail & Related papers  (2021-02-05T16:27:08Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.