arXiv:2410.17881cs.LG2024-10ICLR被引 17

通过自适应降低梯度秩,实现高效低内存大模型训练。

AdaRankGrad: Adaptive Gradient-Rank and Moments for Memory-Efficient LLMs Training and Fine-Tuning

  • 训练中动态降低梯度秩,用在线更新的低秩投影优化
  • 相比现有方法减少显著内存占用,且提升预训练与微调性能
  • 适合资源有限但需全参数微调的场景,如生物模型训练

大语言模型(LLMs)的训练和微调面临内存与计算需求高的挑战,主要源于模型权重和优化器状态规模不断增大。尽管已有诸如低秩适配(LoRA)等技术缓解此问题,但其将参数搜索限制在低秩子空间,常导致训练动态失衡,需全秩预热来弥补。本文首次形式化证明:随着训练推进,层梯度的秩逐渐下降,渐近趋近于1。基于此,我们提出自适应降低梯度秩的方法,在Adam优化步骤中采用高效的在线低秩投影规则。同时设计随机SVD方案以快速获取投影矩阵。该方法支持全参数微调,通过自适应低秩梯度更新显著降低整体内存开销,优于当前最优方法。最后,我们提供了收敛性分析,并验证了其在语言与生物基础模型训练中的有效性。

原文摘要 · Abstract (English)

Training and fine-tuning large language models (LLMs) come with challenges related to memory and computational requirements due to the increasing size of the model weights and the optimizer states. Various techniques have been developed to tackle these challenges, such as low-rank adaptation (LoRA), which involves introducing a parallel trainable low-rank matrix to the fixed pre-trained weights at each layer. However, these methods often fall short compared to the full-rank weight training approach, as they restrict the parameter search to a low-rank subspace. This limitation can disrupt training dynamics and require a full-rank warm start to mitigate the impact. In this paper, we introduce a new method inspired by a phenomenon we formally prove: as training progresses, the rank of the estimated layer gradients gradually decreases, and asymptotically approaches rank one. Leveraging this, our approach involves adaptively reducing the rank of the gradients during Adam optimization steps, using an efficient online-updating low-rank projections rule. We further present a randomized SVD scheme for efficiently finding the projection matrix. Our technique enables full-parameter fine-tuning with adaptive low-rank gradient updates, significantly reducing overall memory requirements during training compared to state-of-the-art methods while improving model performance in both pretraining and fine-tuning. Finally, we provide a convergence analysis of our method and demonstrate its merits for training and fine-tuning language and biological foundation models.

大模型训练低秩优化内存效率自适应梯度

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。