在低秩约束下实现大模型全秩训练,提升性能同时节省内存。
Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?
- 通过梯度范数缩放机制,用低秩优化器替代全秩优化器。
- 实验显示Fira在预训练与微调中性能优于LoRA和GaLore。
- 适合追求高效训练且不牺牲模型性能的研究者使用。
低秩训练已成为降低大语言模型(LLM)训练内存消耗的有前景方法。现有方法或对权重矩阵进行分解(如LoRA),或对梯度矩阵进行分解(如GaLore),但均将训练限制在低秩子空间,导致性能下降。本文提出首个尝试解决此矛盾的方法——Fira,一种可即插即用的训练框架。我们发现自适应优化器(如Adam)在低秩与全秩训练中对梯度范数的缩放影响相似。基于此,提出基于范数的缩放方法,利用低秩优化器的缩放特性替代原全秩优化器,从而在保持低秩优化器约束的同时实现全秩训练。此外,观察到优化过程中存在突增梯度,可能引发损失尖峰,为此设计范数增长限制器,通过调控梯度范数相对增幅来平滑训练过程。大量实验表明,Fira在大模型预训练与微调任务中表现优于LoRA与GaLore,性能达到甚至超越全秩训练水平。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。