LORA-FA：针对大型语言模型微调的内存高效低秩自适应

paper是香港浸会大学发表在Arxiv 2023的工作

paper title：LORA-FA: MEMORY-EFFICIENT LOW-RANK ADAPTATION FOR LARGE LANGUAGE MODELS FINE-TUNING

ABSTRACT

低秩自适应 (LoRA) 方法可以大大减少微调大型语言模型 (LLM) 所需的可训练参数数量，但是，它仍然需要昂贵的激活内存来更新低秩权重。减少 LoRA 层数或使用激活重新计算可能会损害微调性能或增加计算开销。在这项工作中，我们提出了 LoRA-FA，这是一种内存高效的微调方法，它可以在不降低性能和进行昂贵的重新计算的情况下减少激活内存。 LoRA-FA 选择在每个 LoRA 层中冻结 A 的向下投影权重并更新 B 的向上投影权重。它确保在 LLM 微调期间模型权重的变化位于低秩空间中，同时无需存储全秩输入激活。我们在多种模型类型（RoBERTa、T5、LLaMA）和模型规模上进行了广泛的实验。我们的结果表明，与全参数微调和 LoRA 相比，LoRA-FA 在不同任务中始终能够实现接近的微调精度。此外，与 LoRA 相比，LoRA-FA 可将总内存成本降低高达 1.4 倍。

1 INTRODUCTION

大型语言模型 (LLM) 已成为自然语言处理的基石 (Brown 等人，2020 年；Touvron 等人，2023a；OpenAI，2023 年；Anil 等人，2023 年)，并且对预训练的 LLM 进行微调已被证明非常有效，可以提高其在各种下游任务中的表现 (Liu 等人，2019 年；Wei 等人，2021 年) 并使它们与人类意图保持一致 (Ouyang 等人，2022 年；Bai 等人，2022 年)。然而，对具有完整参数的 LLM 进行微调的成本过高，例如，使用 AdamW（Loshchilov & Hutter，2017）对 LLaMA-65B（Touvron et al, 2023a）模型进行微调需要超过 1TB 的 GPU 内存来存储模型参数、梯度和优化器状态（Rajbhandari et al, 2020）。

为了减少全参数微调的内存，提出了参数高效微调 (PEFT) 方法，仅更新一小部分参数，例如适配器权重 (Houlsby et al, 2019; Hu et al, 2022) 和提示权重 (Li & Liang, 2021; Lester et al, 2021)。在这些方法中，LoRA (Hu et al, 2022) 已证明其性能与全参数微调相当，并且已广泛应用于许多应用 (Dettmers et al, 2023)。

具体而言，LoRA 在线性层的权重旁边添加了一个并行的低秩适配器，如图 1(b) 所示，其中 $W$ 是预训练的权重，$A$ 和 $B$ 是低秩权重。由于 LoRA 冻结了 $W$ 并仅更新较小的矩阵 $A$ 和 $B$，因此与图 1(a) 所示的全参数微调相比，其可训练参数、梯度和优化器状态的内存开销大幅减少（全参数微调可以视为更新 $W$ 并冻结 $A$ 和 $B$）。此外，通过将 $AB$ 的值合并到 $W$ 中，LoRA 不会引入额外的推理延迟。

图 1：(a) 全参数微调 (FT)、(b) LoRA 和 © LoRA-FA 的图示。

然而，LoRA 仍然存在局限性，因为它需要在 LoRA 层中消耗昂贵的激活内存。这是因为在前馈传递过程中需要存储 X 的大型输入激活，并在反向传播传递过程中用于构建 A 的梯度。这意味着与全参数微调相比，LoRA 无法降低激活内存成本。例如，对输入序列长度为 2048 且批大小为 4 的 LLaMA-65B 进行微调需要在所有 LoRA 层中使用超过 50GB 的激活内存（16 位格式）。为了解决这个问题，现有方法选择一部分线性层进行 LoRA 微调（Hu et al, 2022）或使用激活重新计算（Chen et al, 2016），但这会影响微调性能或效率。

在本工作中，我们提出了 LoRA with Frozen-A（LoRA-FA），该方法可以在不增加计算开销的情况下大幅减少 LoRA 的激活内存占用。具体来说，LoRA-FA 选择同时冻结预训练权重 $W$ 和降维权重 $A$，仅更新升维权重 $B$，如图 1© 所示。通过这种方式，我们仅需计算 $B$ 的梯度，而在前向传播过程中只需存储较小的输入 $AX$。

假设 $W\in {\mathbb{R}}^{d\times d},A\in {\mathbb{R}}^{d\times r}$，以及 $B\in {\mathbb{R}}^{r\times d}$，降维权重 $A$ 将 $d$ 维输入 $X$ 映射为 $r$ 维输入 $XA$。由于 $r\ll d$，LoRA-FA 的激活内存需求可以显著降低。例如，LoRA-FA（秩大小 $r=4$）可以将 LLaMA-65B（隐藏维度 $d=8192$）线性层的激活内存需求相比全参数微调减少 2048 倍。同时，LoRA-FA 将可训练参数的数量从 $d^2$ 降至 $dr$，同样减少了 2048 倍。

我们将 $A$ 随机初始化为服从正态分布的矩阵（通常为秩 $r$ 的矩阵），并将 $B$ 初始化为零。这确保了在微调开始前，预训练模型与包含 LoRA-FA 模块的模型预测保持一致。在模型适配过程中，LoRA-FA 固定 $A$ 并更新 $B$ 来提升模型性能，这意味着模型权重的变化 $\Delta W=AB$ 局限于初始 $A$ 的列空间定义的低秩空间。

实验结果表明，LoRA-FA 已足够用于微调大语言模型（LLMs）。此外，LoRA-FA 不改变 LoRA 的前向传播和反向传播计算（除了跳过 $A$ 的梯度计算），因此在微调阶段不会增加计算开销。在推理阶段，与 LoRA 类似，LoRA-FA 可以通过将低秩权重 $AB$ 合并到 $W$ 中实现，与全参数微调模型相比，不会引入额外的推理延迟。

我们对多种模型类型和规模进行了广泛的实验。我们在自然语言理解任务上对 RoBERTa（Liu 等人，2019 年）、在机器翻译任务上对 T5（Raffel 等人，2020 年）以及在 MMLU（Hendrycks 等人，2021 年）基准上对 LLaMA（Touvron 等人，2023a）进行了微调。我们发现，与全参数微调 (FT) 和 LoRA 相比，我们的 LoRA-FA 可以在许多任务中实现非常接近的模型精度。在内存开销方面，与全参数微调和 LoRA 相比，我们的 LoRA-FA 可以将总内存成本分别降低高达 2 倍和 1.4 倍。例如，LoRA-FA 将微调 LLaMA-7B 模型的内存占用从 56GB 减少到 27.5GB。这意味着我们可以使用更低的资源预算（例如，更便宜的 GPU 服务，内存更小）来实现相同的微调性能。此外，我们研究了超参数的影响，结果表明 LoRA-FA 对超参数具有很强的鲁棒性。

总之，LoRA-FA 有几个关键优势：1）通过减少可训练参数和激活的数量来提高内存效率，2）它不会增加微调阶段的计算开销，也不会增加推理阶段的延迟开销，3）与全参数微调相比，它在许多模型和任务中实现了类似的模型性能。

2 BACKGROUND

2.1 LARGE LANGUAGE MODELS

---

我们专注于基于 Transformer 的大型语言模型 (LLM)。Transformer 模型最早在 (Vaswani et al, 2017) 中提出，用于机器翻译任务。后来，不同的 Transformer 模型被用于语言建模 (即预训练)，预训练模型适用于许多下游任务 (Kenton & Toutanova, 2019; Raffel et al, 2020; Brown et al, 2020)。

以仅解码器的 GPT (Brown et al, 2020) 模型为例，它由 L 个堆叠的 Transformer 块组成，每个块有两个子模块：多头注意力 (MHA) 和前馈网络 (FFN)。在 MHA 中，三个线性层将输入转换为查询、键和值，它们被输入到自注意力中进行交互，注意力输出被发送到另一个线性层。 在 FFN 中，我们有两个线性层和它们之间的 GeLU 激活函数。对于 MHA 和 FFN，应用层规范和残差连接来提高模型性能。在 LLM 中，这些线性层中的权重通常占模型参数的大多数，并负责大部分计算浮点运算。

2.2 LOW-RANK ADAPTATION

---

随着微调大型语言模型（LLMs）的全参数方法成本的不断增加，参数高效微调（PEFT）方法（尤其是 LoRA，Hu et al., 2022）被提出，通过仅更新模型参数的一小部分来减轻内存开销，同时实现与全参数微调 LLMs 可比的性能。具体来说，LoRA 在线性层的权重旁添加了一个低秩适配器，其计算公式如下：

$$Y=XW+\alpha XAB$$

其中，$W\in {\mathbb{R}}^{d_{\text{in}}\times d_{\text{out}}}$ 是预训练权重，$d_{\text{in}}$ 是输入维度，$d_{\text{out}}$ 是输出维度。我们忽略偏置项，因为它不会影响分析。$X\in {\mathbb{R}}^{b\times s\times d_{\text{in}}}$ 和 $Y\in {\mathbb{R}}^{b\times s\times d_{\text{out}}}$ 分别为输入和输出张量，$b$ 是批大小，$s$ 是序列长度。对于 LoRA 部分，$A\in {\mathbb{R}}^{d_{\text{in}}\times r}$ 是降维权重，$B\in {\mathbb{R}}^{r\times d_{\text{out}}}$ 是升维权重，$r$ 是秩大小，$\alpha >0$ 是超参数（通常设置为 $1/r$）。

对于如 GPT（Brown et al., 2020）这样的 Transformer 模型，我们通常有 $d_{\text{in}}=d_{\text{out}}=d$，适用于多头注意力（MHA）的四个线性层；而对于前馈网络（FFN）的第一（或第二）线性层，有 $d_{\text{in}}=d,d_{\text{out}}=4d$（或 $d_{\text{in}}=4d,d_{\text{out}}=d$），其中 $d$ 是隐藏维度。默认情况下，我们在 Transformer 块的所有线性层中添加 LoRA 模块以增强微调性能（Zhang et al., 2023b）。

内存复杂度

对于全参数微调，我们需要更新线性层的权重 $W$，其元素数量为 $d_{\text{in}}\times d_{\text{out}}$，GPT 类型模型的权重总数为 $n=12d^{2}L$。对于 LoRA，仅需更新两个低秩矩阵，其元素数量为 $\left(d_{\text{in}}+d_{\text{out}}\right)r$，GPT 的 LoRA 参数总数为 $n_r=18drL$。因此，当秩大小 $r\ll d$ 时，LoRA 可以显著减少可训练参数的数量。

在16位混合精度训练环境下，全参数微调需要 $2n$ 字节用于模型权重，以及 $14n$ 字节用于梯度和优化器状态（32位 AdamW 的状态和参数副本）（Rajbhandari et al., 2020）。相比之下，LoRA 需要 $2n$ 字节用于模型权重，以及 $16n_r$ 字节用于适配器相关的权重、梯度和优化器状态。这意味着当 $n_r\ll n$（或 $r\ll d$）时，LoRA 可以将这一部分的内存开销减少约 8 倍。

然而，在激活内存开销方面，情况有所不同。全参数微调需要存储输入 $X$ 以计算 $W$ 的梯度，而 LoRA 需要存储输入 $X$ 以计算 $A$ 的梯度，同时还需要存储低秩输入 $XA$ 以计算 $B$ 的梯度。具体而言，在 16 位精度下，LoRA 和全参数微调分别需要 $14bsdL+8bsrL$ 字节和 $14bsdL$ 字节的激活内存，其中 $b$ 是批大小，$s$ 是序列长度，$d$ 是隐藏维度，$L$ 是层数。此外，二者在其他组件（如注意力、GeLU 和层归一化）中也会消耗激活内存（Korthikanti et al., 2023）。

因此，与全参数微调相比，LoRA 无法减少（甚至可能增加）激活内存的开销，这不幸地成为了新的内存瓶颈。

减少激活内存的挑战。有两种方法可以降低 LoRA 的激活记忆成本。首先，我们可以将 LoRA 模块添加到少数线性层中，例如 Transformer 模型中的查询和值投影 (Hu et al, 2022)，因此，没有 LoRA 的其他冻结线性层不需要存储其输入激活。然而，这种方法可能会影响微调任务的性能 (Dettmers et al, 2023)，并且还带来了选择使用 LoRA 进行微调的层的难度 (Zhang et al, 2023b)。其次，已经提出了激活重新计算 (Chen et al, 2016; Korthikanti et al, 2023)，以在前馈传递期间仅检查每个 Transformer 块的输入，并在反向传播传递期间从此检查点开始重新计算其他必要的激活。然而，激活重新计算具有非常昂贵的重新计算成本，这会引入大约 1/3 的总计算次数。

3 LORA-FA METHOD

首先，我们介绍了 LoRA-FA 方法的设计，从低秩模型自适应的角度对其进行了解释，并分析了其在减少内存开销方面的优势。其次，我们表明 LoRAFA 可以集成到其他内存优化技术中以提高其利用率。第三，我们讨论了 LoRA-FA 与梯度压缩之间的关系。

3.1 LORA WITH FROZEN-A

---

LoRA 方法更新了两个低秩矩阵 $A$ 和 $B$，并将 $AB$ 作为线性层预训练冻结权重 $W$ 的变化，即 $W+\alpha \Delta W=W+\alpha AB$。如前所述，LoRA 不直接更新 $W$，因此能够显著减少可训练参数的数量，但仍需消耗昂贵的激活内存。

为了应对这一问题，我们提出了 LoRA with Frozen-A（LoRA-FA），在微调过程中冻结 $W$ 和 $A$，仅更新 $B$。具体来说，我们将 $A$ 随机初始化为服从正态分布的矩阵，通常是秩为 $r$ 的矩阵，并将 $B$ 初始化为零，这样 $\Delta W=AB$ 初始为零，因此在微调开始前，模型预测不会受到影响。

低秩模型适配

在微调过程中，如图 1© 所示，我们冻结初始化的 $A$ 和预训练的 $W$，仅更新升维权重 $B$。因此，模型适配期间的权重变化限制在低秩空间内，表示为：

$$\Delta W=AB=Q\bar{B}=\sum _{i=1}^r{Q}_{:,i}{\bar{B}}_{i,:}$$

其中，$A=QR$ 是 $A$ 的 QR 分解，$Q$ 的 $r$ 列（即 $Q_{:,i}$，$i=1,\cdots ,r$）是正交单位向量，当 $A$ 是秩为 $r$ 的矩阵时成立。我们记 $\bar{B}=RB$，推导得 $\Delta W_{:,j}={\sum }_{i=1}^r{\bar{B}}_{ij}{Q}_{:,i}$，因此 $\Delta W$ 的任何一列都是 $r$ 个正交向量的组合。换句话说，权重变化位于由 $A$ 的列空间定义的低秩空间中。

内存复杂度

我们详细分析了 LoRA-FA 的内存复杂度。对于 LoRA-FA 模块，它仅计算 $B$ 的梯度，包含 $d_{\text{out}}\times r$ 个元素。在 GPT 类型模型中，总的可训练参数为 $n_r/2=9drL$，即 LoRA 可训练参数的一半。因此，在 16 位混合精度训练中，模型权重和适配器相关状态的内存开销为 $2n+8n_r$ 字节。

更重要的是，就激活内存而言，LoRA-FA 仅需存储低秩输入 $XA$ 以计算 $B$ 的梯度，这对所有 LoRA-FA 模块来说需 $8bsrL$ 字节（16 位）。与全参数微调相比，LoRA-FA 通过显著减少可训练参数和输入激活的数量，在内存效率上具有明显优势。

3.2 COMBINATION WITH MEMORY OPTIMIZATIONS

---

LoRA-FA 可以自然地与先进的内存优化方法相结合，例如权重量化（Dettmers 等人，2023 年）、权重分片（Rajbhandari 等人，2020 年）和选择性激活重新计算（Korthikanti 等人，2023 年）。

权重量化

如前所述，16 位格式下模型权重的内存开销为 $2n$，其中 $n$ 是模型参数的数量。例如，LLaMA-65B 模型的权重内存开销为 130GB，无法装载到一块 NVIDIA A100 (80GB) GPU 上。在 LoRA-FA 中，由于模型权重在微调过程中是冻结的，我们可以将其量化为更低的位宽，以减少模型权重的内存开销，同时不影响微调性能。例如，8 位（Dettmers et al., 2022a）和 4 位量化方法（Dettmers et al., 2023）可以与 LoRA-FA 结合，将模型权重内存减少至原来的 1/2 或 1/4。

权重切分

在使用数据并行技术在多块 GPU 上训练 LLM 时，可以结合权重切分或 ZeRO 第 3 阶段（Rajbhandari et al., 2020）技术，将模型权重切分到不同的 GPU 上，从而将每块 GPU 的内存开销减少为原来的 $1/g$，其中 $g$ 是 GPU 的数量。与全参数微调中使用 ZeRO 第 3 阶段不同，我们仅对模型权重进行切分，并在前向和反向传播计算中进行权重收集，而不对适配器相关权重、其梯度和优化器状态进行切分。然而，在 LoRA-FA 中，权重切分引入了代价高昂的权重收集通信开销，而数据并行技术仅需通信少量可训练参数的梯度。

选择性激活重计算

Transformer 模型的其他组件（如注意力机制、层归一化、GeLU 和 dropout）也会导致激活内存开销（Korthikanti et al., 2023）。为了解决这个问题，可以使用完整激活重计算来存储每个 Transformer 块的输入。然而，这将使 LoRA-FA 相较于 LoRA 的内存优势丧失，因为在完整激活重计算下，无需为 LoRA 层存储输入。

为了平衡激活开销和重计算开销，我们采用选择性激活重计算，仅对部分模型组件进行重计算。例如，FlashAttention（Dao et al., 2022）可以消除注意力机制中 softmax 输出的内存开销，同时通过减少高带宽内存（HBM）访问加速注意力计算。此外，我们可以通过存储随机生成器状态来精确重现 dropout 掩码，从而实现 dropout 的重计算。

3.3 RELATION TO GRADIENT COMPRESSION

---

我们探讨 LoRA-FA 与低秩梯度压缩（Vogels et al., 2019；Zhang et al., 2023a）之间的关系。对于一个 LoRA-FA 层（为简化，省略 $\alpha$），即 $Y=XW+XAB$，$B$ 的梯度计算为：

$$dB=A^T{X}^{T}dY=A^{T}dW$$

在标准 SGD 的单步更新中，$B$ 的变化为 $\Delta B=-\eta dB$，其中 $\eta$ 是学习率。因此，冻结 $A$ 时，$W$ 的变化为：

$$\Delta W=A\Delta B=-\eta A{A}^{T}dW$$

其中，$dW$ 是 $W$ 的梯度。这表明，LoRA-FA 等价于一种低秩梯度压缩方法，用于全参数微调。在这种方法中，计算得到的权重梯度通过 $A^{T}dW$ 被压缩（以减少梯度通信开销），然后通过 $A\left(A^{T}dW\right)$ 被解压缩。

由于 $A$ 是从正态分布初始化的，我们有：

$$\mathbb{E}\left[A{A}^{T}dW\right]=\mathbb{E}\left[A{A}^T\right]dW=rdW$$

这表明这种方法几乎是无偏的梯度压缩。

然而，梯度压缩在微调大型语言模型（LLMs）时，相较于 LoRA-FA 并无优势，因为梯度压缩仍然需要更新全部参数，带来较大的内存开销，而 LoRA-FA 仅需少量可训练权重，也能在数据并行设置下减少梯度通信。此外，当使用自适应优化方法（如 AdamW）时，这两种方法会表现出不同的行为。