Gradient Weight-normalized Low-rank Projection for Efficient LLM Training

📄 arXiv: 2412.19616v2 📥 PDF

作者: Jia-Hong Huang, Yixian Shen, Hongyi Zhu, Stevan Rudinac, Evangelos Kanoulas

分类: cs.LG, cs.AI

发布日期: 2024-12-27 (更新: 2025-01-05)

备注: Accepted by the 39th AAAI Conference on Artificial Intelligence (AAAI-25) [Main Technical Track]

🔗 代码/项目: GITHUB

💡 一句话要点

提出梯度权重归一化低秩投影(GradNormLoRP)，高效训练大型语言模型。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 参数高效微调 低秩投影 梯度归一化 内存优化

📋 核心要点

  1. 现有全量微调LLM计算资源需求大，参数高效微调方法性能不及全量微调，且内存效率不足。
  2. 提出GradNormLoRP，通过权重归一化改善梯度条件，并对权重和梯度矩阵进行低秩近似，降低内存占用。
  3. 实验表明，GradNormLoRP显著降低优化器内存使用，并在微调任务中优于现有低秩方法，如LoRA。

📝 摘要（中文）

大型语言模型(LLMs)在各种任务中表现出色，但对计算资源的需求不断增加，尤其是在下游任务中广泛使用全量微调时，这带来了重大挑战。为了解决这个问题，开发了参数高效微调(PEFT)方法，但它们通常不如全量微调，并且在内存效率方面存在不足。本文提出了一种新的方法，即梯度权重归一化低秩投影(GradNormLoRP)，该方法在提高参数和内存效率的同时，保持与全量微调相当的性能。GradNormLoRP对权重矩阵进行归一化，以改善梯度条件，从而促进优化过程中的更好收敛。此外，它将低秩近似应用于权重和梯度矩阵，从而显著减少训练期间的内存使用。大量实验表明，我们的8位GradNormLoRP可将优化器内存使用量减少高达89.5%，并支持在NVIDIA RTX 4090等消费级GPU上预训练大型LLM（如LLaMA 7B），而无需额外的推理成本。此外，GradNormLoRP在微调任务中优于现有的低秩方法。例如，当在所有GLUE任务上以秩为8微调RoBERTa模型时，GradNormLoRP的平均得分为80.65，超过了LoRA的79.23。这些结果表明，GradNormLoRP是高效LLM预训练和微调的一种有前途的替代方案。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）微调过程中计算资源消耗过高，尤其是内存占用过大的问题。现有参数高效微调（PEFT）方法虽然降低了参数量，但在性能上通常不如全量微调，并且在内存效率方面仍然存在瓶颈。

核心思路：论文的核心思路是通过梯度权重归一化和低秩投影来提高训练效率和内存效率。权重归一化改善梯度条件，加速收敛；低秩投影则减少权重和梯度矩阵的存储空间，从而降低内存占用。这样设计的目的是在保证模型性能的同时，显著降低训练成本。

技术框架：GradNormLoRP方法主要包含两个关键步骤：首先，对权重矩阵进行归一化处理，以改善梯度流动和优化过程。其次，对权重矩阵和梯度矩阵应用低秩近似，从而减少需要存储和计算的参数数量。整个过程无额外的推理成本。

关键创新：GradNormLoRP的关键创新在于将梯度权重归一化与低秩投影相结合，同时应用于权重和梯度矩阵。这种组合方式既改善了优化过程，又显著降低了内存占用，从而能够在消费级GPU上训练大型LLM。与现有方法（如LoRA）相比，GradNormLoRP在性能和内存效率上都具有优势。

关键设计：GradNormLoRP的关键设计包括：1) 权重归一化的具体实现方式，可能涉及到对权重矩阵的某种范数进行约束或调整。2) 低秩近似的秩的选择，需要在性能和内存占用之间进行权衡。3) 优化器的选择和参数设置，以确保在归一化和低秩近似的情况下，模型能够有效收敛。论文中可能还涉及一些正则化项或损失函数的设计，以进一步提高模型的泛化能力。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果表明，8位GradNormLoRP可将优化器内存使用量减少高达89.5%，并支持在NVIDIA RTX 4090等消费级GPU上预训练LLaMA 7B等大型LLM。在GLUE任务上微调RoBERTa模型时，GradNormLoRP的平均得分（80.65）超过了LoRA（79.23），表明其在性能上也优于现有低秩方法。

🎯 应用场景

GradNormLoRP可应用于各种需要高效训练大型语言模型的场景，例如在资源受限的环境下进行模型微调，或者在消费级GPU上预训练大型模型。该方法可以降低训练成本，加速模型开发周期，并促进LLM在更广泛的应用领域中的部署，例如自然语言处理、机器翻译、文本生成等。

📄 摘要（原文）

Large Language Models (LLMs) have shown remarkable performance across various tasks, but the escalating demands on computational resources pose significant challenges, particularly in the extensive utilization of full fine-tuning for downstream tasks. To address this, parameter-efficient fine-tuning (PEFT) methods have been developed, but they often underperform compared to full fine-tuning and struggle with memory efficiency. In this work, we introduce Gradient Weight-Normalized Low-Rank Projection (GradNormLoRP), a novel approach that enhances both parameter and memory efficiency while maintaining comparable performance to full fine-tuning. GradNormLoRP normalizes the weight matrix to improve gradient conditioning, facilitating better convergence during optimization. Additionally, it applies low-rank approximations to the weight and gradient matrices, significantly reducing memory usage during training. Extensive experiments demonstrate that our 8-bit GradNormLoRP reduces optimizer memory usage by up to 89.5% and enables the pre-training of large LLMs, such as LLaMA 7B, on consumer-level GPUs like the NVIDIA RTX 4090, without additional inference costs. Moreover, GradNormLoRP outperforms existing low-rank methods in fine-tuning tasks. For instance, when fine-tuning the RoBERTa model on all GLUE tasks with a rank of 8, GradNormLoRP achieves an average score of 80.65, surpassing LoRA's score of 79.23. These results underscore GradNormLoRP as a promising alternative for efficient LLM pre-training and fine-tuning. Source code: https://github.com/Jhhuangkay/Gradient-Weight-normalized-Low-rank-Projection-for-Efficient-LLM-Training