1. 引言
微调 (Fine-tuning) 是将预训练大模型 (LLM) 应用于下游任务的常用方法。然而,直接微调大模型的所有参数通常需要大量的计算资源和内存。LoRA (Low-Rank Adaptation) 是一种高效的微调方法,它通过引入少量可训练参数,固定预训练模型的权重,从而在保持性能的同时大大减少了计算开销。
本文将深入分析 LoRA 的原理,并结合 Llama 源码解读其实现逻辑,最后探讨 LoRA 的优势。
2. LoRA 原理
LoRA 的核心思想是:预训练模型中已经包含了大量的低秩 (low-rank) 特征,微调时只需要对这些低秩特征进行微调即可。
具体来说,LoRA 假设权重更新矩阵 ΔW 也是低秩的。对于一个预训练的权重矩阵 W ∈ R^(d×k),LoRA 将其更新表示为:
W' = W + ΔW = W + BA
其中:
W
是预训练的权重矩阵。ΔW
是权重更新矩阵。B ∈ R^(d×r)
和A ∈ R^(r×k)
是两个低秩矩阵,r
远小于d
和k
,r
被称为 LoRA 的秩 (rank)。
在训练过程中,W
被冻结,只有 A
和 B
是可训练的。
直观理解:
可以将 W
看作一个编码器,将输入 x
编码成一个高维表示 Wx
。LoRA 认为,在微调过程中,我们不需要完全改变这个编码器,只需要通过 BA
对其进行一个低秩的调整即可。
3. Llama 中 LoRA 的实现
虽然 Llama 官方代码没有直接集成 LoRA,但我们可以使用一些流行的库 (例如 peft
by Hugging Face) 来实现 Llama 的 LoRA 微调。peft
库提供了 LoraConfig
和 get_peft_model
等工具,可以方便地将 LoRA 应用于各种 Transformer 模型。
3.1 使用 peft
库实现 Llama 的 LoRA 微调
以下是一个使用 peft
库实现 Llama 的 LoRA 微调的简化示例:
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import get_peft_model, LoraConfig, TaskType
# 加载预训练的 Llama 模型和分词器
model_name = "meta-llama/Llama-2-7b-hf" # 假设使用 Llama 2 7B
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# LoRA 配置
config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
inference_mode=False,
r=8, # LoRA 的秩
lora_alpha=32, # LoRA 的缩放因子
lora_dropout=0.1, # Dropout 比例
target_modules=["q_proj", "v_proj"], # 需要应用 LoRA 的模块
)
# 获取支持 LoRA 的模型
model = get_peft_model(model, config)
# 打印可训练参数的比例
model.print_trainable_parameters()
# ... (加载数据,进行训练) ...
代码解释:
- 加载预训练模型:使用
transformers
库加载预训练的 Llama 模型和分词器。 - LoRA 配置:创建一个
LoraConfig
对象,指定 LoRA 的配置参数:task_type
:任务类型,这里是因果语言模型 (Causal Language Modeling)。r
:LoRA 的秩。lora_alpha
:LoRA 的缩放因子,用于控制 LoRA 模块的权重。lora_dropout
:Dropout 比例。target_modules
: 指定需要应用 LoRA 的模块, 通常是注意力层中的q_proj
,v_proj
, 还可以是k_proj
,o_proj
,gate_proj
,up_proj
,down_proj
等。不同的模型需要根据实际情况配置。
- 获取支持 LoRA 的模型:使用
get_peft_model
函数将原始的 Llama 模型转换为支持 LoRA 的模型。 - 打印可训练参数:使用
model.print_trainable_parameters()
可以查看模型中可训练参数的比例,通常 LoRA 的可训练参数比例非常小。
3.2 peft
库中 LoRA 的实现细节 (部分)
peft
库中 LoraModel
类的部分代码 (为了清晰起见,已进行简化):
class LoraModel(torch.nn.Module):
# ...
def _find_and_replace(self, model):
# ... (遍历模型的每一层) ...
if isinstance(module, nn.Linear) and name in self.config.target_modules:
new_module = Linear(
module.in_features,
module.out_features,
bias=module.bias is not None,
r=self.config.r,
lora_alpha=self.config.lora_alpha,
lora_dropout=self.config.lora_dropout,
)
# ... (将原模块的权重赋值给新模块) ...
# ...
class Linear(nn.Linear):
def __init__(
self,
in_features: int,
out_features: int,
r: int = 0,
lora_alpha: int = 1,
lora_dropout: float = 0.0,
**kwargs,
):
super().__init__(in_features, out_features, **kwargs)
# LoRA 参数
self.r = r
self.lora_alpha = lora_alpha
# 初始化 A 和 B
if r > 0:
self.lora_A = nn.Parameter(torch.randn(r, in_features))
self.lora_B = nn.Parameter(torch.zeros(out_features, r)) # B 初始化为全 0
self.scaling = self.lora_alpha / self.r
def forward(self, x: torch.Tensor):
result = F.linear(x, self.weight, bias=self.bias) # W @ x
if self.r > 0:
result += (
self.lora_B @ self.lora_A @ x.transpose(-2, -1) # (B @ A) @ x
).transpose(-2, -1) * self.scaling
return result
代码解释:
_find_and_replace
函数:遍历模型的每一层,找到需要应用 LoRA 的线性层 (例如,q_proj
,v_proj
),并将其替换为Linear
层。Linear
类:继承自nn.Linear
,并添加了 LoRA 的参数lora_A
和lora_B
。lora_A
初始化为随机值。lora_B
初始化为全 0,这是为了保证在训练开始时,LoRA 部分的输出为 0,不影响预训练模型的原始行为。scaling
是一个缩放因子,用于控制 LoRA 模块的权重。
forward
函数:F.linear(x, self.weight, bias=self.bias)
计算原始的线性变换W @ x
。(self.lora_B @ self.lora_A @ x.transpose(-2, -1)).transpose(-2, -1) * self.scaling
计算 LoRA 部分的输出(B @ A) @ x
,并乘以缩放因子。- 将两者相加,得到最终的输出。
4. LoRA 的优势
- 高效的参数利用:LoRA 只需微调少量的参数 (A 和 B),而冻结了预训练模型的大部分参数,大大减少了训练时的内存占用和计算开销。
- 快速的训练速度:由于可训练参数较少,LoRA 的训练速度通常比全量微调快得多。
- 防止过拟合:LoRA 的低秩约束起到了一定的正则化作用,有助于防止过拟合。
- 性能相当:在许多任务上,LoRA 可以达到与全量微调相当的性能。
- 易于部署:训练完成后,可以将
W
和BA
相加,得到新的权重矩阵W'
,然后像使用原始的预训练模型一样进行部署,无需额外的计算开销。