ICML 2024 Oral || DaLore: 大模型加速优化新策略

深度学习与图网络 · 公众号 · · 2024-07-07 22:54

正文

注明：本文与Deep LoRA核心思想相似，都发现了权重矩阵的低秩性，并提出在全参数下利用低秩子空间进行训练的优化策略。不同的是， - Deep LoRA 侧重于微调，需要额外添加微调矩阵，然后更新权重梯度有变化的参数。 - GaLore则采用另外一种思路：直接在全部参数矩阵基础上降维，然后再使用优化器，预训练和微调阶段均适用，在大量减少内存消耗的前提下保证性能。
ICML 2024 Oral || Deep LoRA: 基于子空间不变性的模型压缩与优化

1. 基本信息和摘要

论文题目

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

Arxiv: https://arxiv.org/pdf/2403.03507
Code: https://github.com/jiaweizzhao/GaLore

作者

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

作者研究单位

California Institute of Technology
Meta AI
University of Texas at Austin
Carnegie Mellon University

拟解决问题

大型语言模型（LLMs）全参数训练过程中存在显著 内存挑战 。
LoRA虽然能够缓解内存挑战，但 表现欠佳 ，尤其是在预训练阶段。原因如下：1）参数搜索空间被限制在低秩子空间，而最优的参数解未必是低秩，2）且会影响原来梯度训练轨迹/动态。

GaLore是一种新的训练策略， 全参数训练 的同时 减少内存消耗 ，同时适用于训练和微调过程。

下图展示了在单个设备上，不使用激活检查点和内存卸载的情况下，预训练LLaMA 7B模型内存消耗。

摘要

大型语言模型（LLMs）在多个领域展现出了令人印象深刻的性能，包括对话式AI和语言翻译。然而，它们的预训练和微调不仅需要巨大的计算量，而且对内存的需求也很高。这些内存需求包括数十亿的可训练参数，以及它们的梯度和优化器状态，这些可能比参数存储本身还要大。例如， 从头开始预训练一个LLaMA 7B模型至少需要58 GB的内存。 这使得在消费级GPU上进行训练变得不可行。

为了解决这一挑战，本文提出了GaLore，这是一种训练策略，它利用了 权重矩阵梯度的慢变化低秩结构 ， 而不是试图将权重矩阵本身近似为低秩 。

1）本文首先从理论上说明了在训练期间梯度矩阵变得低秩。

2）然后，提出了GaLore，它计算两个投影矩阵和，将梯度矩阵投影到低秩形式。这样，依赖于逐元素梯度统计的优化器状态的内存成本可以大幅降低。

GaLore允许全参数学习，同时比LoRA等常见低秩适应方法更节省内存。本文的方法能够减少优化器状态的内存使用（高达65.5%），同时在预训练LLaMA 1B和7B架构以及微调RoBERTa模型时保持效率和性能。

2. 方法

传统全参数优化

, 共需要内存消耗

具体如下：

是在时间步的权重矩阵, 是初始权重矩阵, 是学习率。是在时间步经过处理的梯度，用于更新权重。
是一个逐元素的状态保持梯度正则化器（例如，Adam 优化器中的）。是在时间步反向传播得到的原始梯度矩阵。

其中，Adam正则化过程 , 需要：

GaLore

, 共需要内存消耗

GaLore算法定义了如下的梯度更新规则：

其中，