Fira:我们能否在低秩约束下实现完整秩训练大语言模型?
机器学习
2025-11-18 v3 人工智能
摘要
低秩训练已成为降低大型语言模型(LLM)训练内存消耗的有前景的方法。以往的方法要么依赖分解权重矩阵(如LoRA),要么 seeks to 分解梯度矩阵(如GaLore),以确保降低内存消耗。然而,两者都限制了训练处于低秩子空间,从而不可避免地导致次优性能。这引发了一个问题:是否可能在保持低秩约束以实现内存效率的同时,实现完整秩训练(即使用完整秩权重的完整秩梯度进行训练),以避免不佳的结果?本文提出了一种名为Fira的全新可插入式LLM训练框架,这是首次尝试实现上述目标。首先,我们观察到在LLM训练过程中,一种有趣的现象:自适应优化器(如Adam)对梯度范数的缩放影响在从低秩训练到完整秩训练之间保持相当一致。基于这一观察,我们提出了一种基于范数的缩放方法,利用低秩优化器的缩放效应来替代原始完整秩优化器的缩放效应,从而实现完整秩训练。如此一来,我们既可以在优化器中保持低秩约束,又能实现完整秩训练以获得更佳性能。此外,我们发现,在优化过程中会出现梯度急剧上升的现象,潜在地引发损失尖峰。为此,我们进一步提出了一种范数增长限制器,通过调节梯度范数相对增长来平滑梯度。对LLM的预训练和微调进行大量实验表明,Fira在性能上优于LoRA和GaLore,达到与甚至优于完整秩训练相当的性能。
引用
@article{arxiv.2410.01623,
title = {Fira: Can We Achieve Full-rank Training of LLMs Under Low-rank Constraint?},
author = {Xi Chen and Kaituo Feng and Changsheng Li and Xunhao Lai and Xiangyu Yue and Ye Yuan and Guoren Wang},
journal= {arXiv preprint arXiv:2410.01623},
year = {2025}
}
备注
NeurIPS 2025, Project page: https://github.com/xichen-fy/Fira