Lora 小结

Lora 是微软的研究人员为了解决大语言模型微调而开发的一项技术，论文地址LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS.

1. ABSTRACT

自然语言处理的一个重要范例包括对通用领域的大规模数据的大规模预训练和对特定任务或领域的微调。当我们预训练更大的模型时，全量的微调，即重新训练模型所有的参数，变得不太可行。使用 GPT-3 175B 作为例子，不是微调后的模型的独立实例，每个都是 175B 的参数，这是令人望而却步的昂贵。作者提出了 Low-Rank Adaptation, 简称 LoRA, 它冻结预训练模型的权重，并将训练的秩分解矩阵插入到 Transformer 架构的没一层，这大大减少了下游任务的训练参数。对比 GPT-3 175B 用 Adam 微调，LoRA 能减少 10000 倍训练参数量，并且 GPU 内存需要减少 3 倍。在 RoBERTa, DeBERTa，GPT- 2 和 GPT- 3 上，LoRA 表现比微调好，尽管更少的训练参数，更高的训练吞吐量；并且与用 adapter 不同，没有额外的推理延迟。作者还对语言模型适应性中 rank-deficiency 现象进行了实验研究，来揭示 LoRA 的有效性。微软还发布了一个软件包，方便将 LoRA 与 PyTorch 模型集成，并提供了 RoBERTa、DeBERTa 和 GPT- 2 的实现和模型 checkpoint，地址为 https://github.com/microsoft/LoRA。

2. LoRA 原理

利用上图来说明 LoRA 的原理，具体可以看论文4.1 LOW-RANK-PARAMETRIZED UPDATE MATRICES。

1. 在原始 PLM (Pre-trained Language Model) 旁边增加一个旁路，做一个降维再升维的操作，来模拟所谓的“instrisic dimension。
2. 训练的时候固定 PLM 的参数，只训练降维矩阵 $A$ 和升维矩阵 $B$。而模型的输入输出维度不变，输出时将 $BA$ 的参数与 固定的 PLM 参数叠加。
3. 用随机高斯分布初始化 $A$，用 0 矩阵初始化 $B$，保证训练的开始此旁路矩阵依然是 0 矩阵。

假设要在下游任务微调一个预训练语言模型（如 GPT-3），则需要更新预训练模型参数，公式表示如下：
$$
W_0 + \Delta W
$$
其中 $W_0$ 是预训练模型初始化的参数，$\Delta W$ 就是需要更新的参数。如果是全参数微调，则它的参数量是 $W_0$, 比如 GPT-3 175B，则 就是 175B 参数。对应 LLM, 参数量是非常大的。而对于 LoRA 来说，只需要微调 $\Delta W$。

假设预训练的矩阵为 $W_0 \in \mathbb{R}^{d\times k}$，它的更新可表示为：
$$
W_0 + \Delta W = W_0 + BA, \quad B\in \mathbb{R}^{d\times r}, \ A\in \mathbb{R}^{r\times k}
$$
其中秩 $r \ll \min(d, k)$, 如果 r 远远小于 d,k，那么参数量又会减少很多。

在 LoRA 的训练过程中，$W_0$ 是固定不变的，只有 $B,A$ 会更新，这两个才是训练参数。

在前向过程中，$W_0$ 和 $\Delta W$ 都会乘以输入 $x$，然后相加得到输出，即：
$$
h = W_0x+\Delta W x = W_0x+BA x
$$
LoRA 的这种思想有点类似于残差连接，同时使用这个旁路的更新来模拟 Full Fine-Tuning 的过程。并且，Full Fine-Tuning 可以被看做是 LoRA 的特例，即 r = k 时。推理时，没有其它额外的计算，也就没有额外的推理延迟，只要计算 $W=W_0+\Delta W$。

如上图 1 所示，对 $A$ 进行随机高斯初始化，$B$ 初始化为 0，因此 $\Delta W =BA$ 初始时为 0，用 $\frac{\alpha}{r}$ 对 $\Delta Wx$ 放缩，其中 $\alpha$ 是 $r$ 中的一个常数。在使用 Adam 进行优化时，调整 $\alpha$ 的过程跟调整学习率差不多，就是初始化的时候恰当的设置。因此，我们只需将 $\alpha$ 设置为我们尝试的第一个 $r$，并且不再调整它。

3. 问题描述

假设给定一个预训练好的自回归语言模型 $P_{\mathrm{\Phi}}(y|x)$，参数化为 $\mathrm{\Phi}$。例如，$\mathrm{P}{\mathrm{\Phi}}(y|x)$ 可以是一个通用的多任务学习器如 GPTm，其基于 Transformer 架构。考虑将这个预训练模型迁移到下游任务条件文本生成，如摘要生成，机器阅读理解，自然语言转 SQL（NL2SQL)。每个下游任务用文本 - 目标对的训练集表示：$\mathcal{Z}={(x_i, y_i)}{i=1,\cdots,N}$，这里 $x_i$ 和 $y_i$ 都是 tokens 的序列。例如，

1. 在 NL2SQL 中，$x_i$ 是自然语言的 query 和 $y_i$ 是对应的 SQL 命令；
2. 对于摘要，$x_i$ 是文章的内容而 $y_i$ 是其摘要

在全量的微调中，模型初始化为预训练的权重 $\mathrm{\Phi}0$ 和通过沿着梯度方向重复地更新到 $\Phi_0+\Delta\Phi$ 来最大化条件语言模型的目标函数：$$ \max {\Phi} \sum_{(x, y) \in \mathcal{Z}} \sum_{t=1}^{|y|} \log \left(P_{\Phi}\left(y_{t} \mid x, y_{<t}\right)\right) \tag{1}
$$

全量微调中一个主要的缺点是对于每个下游任务，LLM 需要学习一组不同的参数 $\Delta \Phi$，其维度等于 LLM 原始的参数 $|\Phi_0|$。因此，如果预训练模型非常大（比如 GPT-3 175B 就是 $|\Phi_0|\approx 175\ \text{Billion}$），如果可实现的方案中，存储和部署大量独立的微调模型的实例具有非常大挑战性。

在本文中，作者采用更有效的参数方法，这里特定任务的参数增量为 $\Delta \Phi = \Delta \Phi (\Theta)$，

进一步编码为一个更小的参数 $\Theta$，其中 $|\Theta| \ll ||\Theta_0|$。因此目标函数中找到 $\Delta \Theta$ 变成优化 $\Theta$:
$$
\max {\Theta} \sum{(x, y) \in \mathcal{Z}} \sum_{t=1}^{|y|} \log \left(p_{\Phi_{0}+\Delta \Phi(\Theta)}\left(y_{t} \mid x, y_{<t}\right)\right) \tag{2}
$$
这段主要提出了使用低秩表示来编码 $\Delta \Theta$，这对计算和内存都是非常高效的。当预训练模型是 GPT-3 175B，训练参数 $|\Theta|$ 能视为原参数量 $|\Theta_0|$ 的 0.01%。

更具体来说，使用 LORA 不再像以前微调一样更新整个 LLM 的参数，只对只对 module 的增加一个旁路，这个旁路像上图 1 所示一样，最重要的是这里的 $r \ll d$，比如 r =8, d=2048，这样 AB 两个矩阵的参数就远远小于原始的参数 $|\Theta_0|$，训练时只优化这个旁路，原始参数矩阵不变，最后将这个变化量加到原来的参数上，就完成了 LORA 的训练过程。

4.HuggingFace 中 peft 使用 LoRA

huggingface 集成了 LoRA 训练的包peft，官方简单使用 LORA 微调 BigScience 机器翻译模型。

from transformers import AutoModelForSeq2SeqLM
from peft import get_peft_config, get_peft_model, LoraConfig, TaskType
model_name_or_path = "bigscience/mt0-large"
tokenizer_name_or_path = "bigscience/mt0-large"

peft_config = LoraConfig(
    task_type=TaskType.SEQ_2_SEQ_LM,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1
)

model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
# output: trainable params: 2359296 || all params: 1231940608 || trainable%: 0.19151053100118282

这里主要介绍下 LoraConfig

比如对 chatglm-6b 设置 LoraConfig 如下：

from peft import (LoraConfig,
                get_peft_model,
                get_peft_model_state_dict,
                set_peft_model_state_dict
                )

config = LoraConfig(r=args.lora_r,
                        lora_alpha=32,
                        target_modules=['query_key_value'],
                        lora_dropout=0.1,
                        bias='none',
                        task_type='CAUSAL_LM',
                        inference_mode=False

                        )

model = AutoModelForSeq2SeqLM.from_pretrained("THUDM/chatglm-6b"
                                                  trust_remote_code=True
                                                  )
model = get_peft_model(model, config)
#接下来就可以用 deepspeed 或者 Trainer 训练

参考

[1] LoRA：大语言模型参数高效性微调方法

[2] LORA：大模型轻量级微调

[3] alpaca-lora

[4] 大模型训练——PEFT 与 LORA 介绍