大模型平民化技术之LORA

news2024/12/22 9:31:42

1. 引言

在这篇博文中，我将向大家介绍LoRA技术背后的核心原理以及相应的代码实现。

LoRA 是 Low-Rank Adaptation 或 Low-Rank Adaptors 的首字母缩写词，它提供了一种高效且轻量级的方法，用于微调预先训练好的的大语言模型。这包括 BERT 和 RoBERTa 等掩码语言模型，以及 GPT、Llama 和 Mistral 等因果推断模型。

闲话少说，我们直接开始吧！

2. 优势分析

LoRA的主要优点之一是他们的效率。通过使用更少的参数，LoRA显著降低了模型训练过程中计算复杂性和显存使用量。这可以让我们在消费级的GPU上来训练大模型，并且可以便利地将我们训练好的LoRA权重（以兆为单位）分发给其他人。

此外，LoRA可以提升模型的泛化性。通过限制模型的复杂度，可以有助于防止在训练数据有限场景下的过拟合现象；由于LoRA至少保留了初始模型的能力，在处理一些新的，未见过的数据时更具有弹性。

最后，LoRA可以无缝地集成到现有的神经网络架构中。这种集成允许以最小的额外训练成本对预训练模型进行微调和调整，使其非常适合迁移学习应用。

3. 工作原理

LoRA 的基本思想是保持预训练矩阵（即原始模型的参数）冻结（即处于固定状态），并且只在原始矩阵中添加一个小的增量，其参数量比原始矩阵少很多。

例如，考虑矩阵 W，它可以是全连接层的参数，也可以是来Transformer中计算自注意力机制的矩阵之一：
在这里插入图片描述

显然，如果 Worig 的维数为 n×m，而假如我们只是初始化一个具有相同维数的新的增量矩阵进行微调，虽然我们也实现类似的功能，但是我们的参数量将会加倍。 LoRA使用的Trick就是通过训练低维矩阵 B 和 A ，通过矩阵乘法来构造 ΔW ，来使 ΔW 的参数量低于原始矩阵。
在这里插入图片描述
这里我们不妨定义秩 r，它明显小于基本矩阵维度 r≪n 和 r≪m。则矩阵 B 为 n×r，矩阵 A 为 r×m。将它们相乘会得到一个维度为 nxm的W 矩阵，但构建的参数量减小了很多。

此外，我们希望我们的增量ΔW在训练开始时为零，这样微调就会从原始模型一样开始。因此，B 通常初始化为全零，而 A 初始化为随机值（通常呈正态分布）。

4. 举个栗子

我们不妨来看个直观的栗子，如下图所示：
在这里插入图片描述
想象一下，我们的基本维数是 1024，我们选择了 LoRA 的秩r为 4，则对于上述过程：
● 权重W的参数量为1024X1024≈1M
● A和B的参数量一致，均为rX1024≈4K，这样二者之和为8K
● 这样使用LoRA技术，在上述例子中我们仅仅需要训练0.8%的参数就可以更新我们的参数矩阵

5. LoRA指令速查

主要可查阅微软的官方文档Github: 戳我，由于封装的很好，目前该库页整合至 HuggingFace Parameter-Efficient Fine-Tuning (PEFT) 。

如果模型要将特定层替换成 LoRA，需要调整模型的结构，但调用很简单：

# ===== Before =====
# layer = nn.Linear(in_features, out_features)
# ===== After ======
import loralib as lora
# Add a pair of low-rank adaptation matrices with rank r=16
layer = lora.Linear(in_features, out_features, r=16)

在训练之前要把原本的LLM模型 Freeze 住，并且设定只有 LoRA 的参数是可训练的

import loralib as lora
model = BigModel()
# This sets requires_grad to False for all parameters without the string "lora_" in their names
lora.mark_only_lora_as_trainable(model)
# Training loop
for batch in dataloader:

保存模型时也可以只储存LoRA所训练的权重，这特性将方便大家分享自己的权重

# ===== Before =====
# torch.save(model.state_dict(), checkpoint_path)
# ===== After =====
torch.save(lora.lora_state_dict(model), checkpoint_path)

推理时读取 LoRA 或是原本 LLM 的权重时，要将 strict 设定为 False

# Load the pretrained checkpoint first
model.load_state_dict(torch.load('ckpt_pretrained.pt'), strict=False)
# Then load the LoRA checkpoint
model.load_state_dict(torch.load('ckpt_lora.pt'), strict=False)

6. SD-LoRA应用

近年来生成式 AI 从 DALLE 再到 Stable-diffusion，都显示了现在的 AI 可以生成高质量以及高分辨率的图片，但是让人诟病的还是需要大量的运算资源才能够训练得了这种高分辨率的模型，因为要训练一个高分辨率的扩散模型是需要相当多内存的，即便 Stable-diffusion 将原本的 Pixel-level Diffusion 变成 Latent Diffusion Model 已经大幅降低训练的内存，但仍无法在单一张 11 GB 的 GPU 上训练，但现在不一样了，有人将 LoRA 技术整合到 Stable-diffusion，推出了 Stable Diffusion LoRA！
在这里插入图片描述
整合 LoRA 至 Stable-diffusion 直接带来了以下的好处：