GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection

📄 arXiv: 2403.03507v2 📥 PDF

作者: Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, Yuandong Tian

分类: cs.LG

发布日期: 2024-03-06 (更新: 2024-06-02)

备注: ICML 2024 (Oral)

💡 一句话要点

提出GaLore以解决大规模语言模型训练中的内存效率问题

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

关键词: 大规模语言模型 内存优化 低秩适应 梯度低秩投影 深度学习

📋 核心要点

  1. 现有的低秩适应方法在训练大型语言模型时存在内存效率低和性能不足的问题。
  2. GaLore通过允许全参数学习并优化内存使用，提供了一种新的训练策略，克服了传统方法的局限性。
  3. 实验结果显示，GaLore在优化器内存上减少了82.5%，并在24GB显存的消费级GPU上成功预训练了7B模型。

📝 摘要（中文）

训练大型语言模型（LLMs）面临显著的内存挑战，主要源于权重和优化器状态的不断增长。常见的内存减少方法，如低秩适应（LoRA），通过在每层的冻结预训练权重中添加可训练的低秩矩阵来减少可训练参数和优化器状态。然而，这些方法在预训练和微调阶段通常表现不佳，因为它们限制了参数搜索到低秩子空间并改变了训练动态。本文提出了梯度低秩投影（GaLore），一种允许全参数学习但比常见低秩适应方法更具内存效率的训练策略。我们的研究表明，GaLore在优化器状态上减少了高达65.5%的内存使用，同时在LLaMA 1B和7B架构上进行C4数据集的预训练时保持了效率和性能。

🔬 方法详解

问题定义：本文旨在解决大型语言模型训练中的内存效率问题，现有的低秩适应方法如LoRA在性能和内存使用上存在不足，限制了模型的训练能力和效果。

核心思路：GaLore的核心思路是通过梯度低秩投影实现全参数学习，同时显著降低内存使用，避免了传统方法对低秩子空间的限制，从而提升了训练动态和效率。

技术框架：GaLore的整体架构包括优化器状态的低秩投影和全参数更新机制，主要模块包括梯度计算、低秩投影和参数更新，确保在内存使用上达到最优。

关键创新：GaLore的最大创新在于其允许全参数学习的同时，通过低秩投影显著减少内存使用，与传统的低秩适应方法相比，提供了更好的性能和效率。

关键设计：在参数设置上，GaLore采用了8位精度来进一步减少内存消耗，优化器的设计确保了在训练过程中内存使用的最小化，同时保持了模型的训练效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，GaLore在优化器内存使用上减少了高达82.5%，并在总训练内存上减少了63.3%，相较于BF16基线，显著提升了训练效率。此外，GaLore首次在24GB显存的消费级GPU上成功预训练了7B模型，展示了其在实际应用中的可行性。

🎯 应用场景

GaLore的研究成果在大型语言模型的训练中具有广泛的应用潜力，尤其是在资源受限的环境中，如消费级GPU。其内存优化策略可以帮助研究人员和开发者在不牺牲模型性能的情况下，进行更大规模的模型训练，推动自然语言处理和人工智能领域的发展。

📄 摘要（原文）

Training Large Language Models (LLMs) presents significant memory challenges, predominantly due to the growing size of weights and optimizer states. Common memory-reduction approaches, such as low-rank adaptation (LoRA), add a trainable low-rank matrix to the frozen pre-trained weight in each layer, reducing trainable parameters and optimizer states. However, such approaches typically underperform training with full-rank weights in both pre-training and fine-tuning stages since they limit the parameter search to a low-rank subspace and alter the training dynamics, and further, may require full-rank warm start. In this work, we propose Gradient Low-Rank Projection (GaLore), a training strategy that allows full-parameter learning but is more memory-efficient than common low-rank adaptation methods such as LoRA. Our approach reduces memory usage by up to 65.5% in optimizer states while maintaining both efficiency and performance for pre-training on LLaMA 1B and 7B architectures with C4 dataset with up to 19.7B tokens, and on fine-tuning RoBERTa on GLUE tasks. Our 8-bit GaLore further reduces optimizer memory by up to 82.5% and total training memory by 63.3%, compared to a BF16 baseline. Notably, we demonstrate, for the first time, the feasibility of pre-training a 7B model on consumer GPUs with 24GB memory (e.g., NVIDIA RTX 4090) without model parallel, checkpointing, or offloading strategies.