【llm对话系统】大模型 Llama 源码分析之 LoRA 微调

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()# ... (加载数据，进行训练) ...

代码解释：

1. 加载预训练模型：使用 transformers 库加载预训练的 Llama 模型和分词器。
2. 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等。不同的模型需要根据实际情况配置。
3. 获取支持 LoRA 的模型：使用 get_peft_model 函数将原始的 Llama 模型转换为支持 LoRA 的模型。
4. 打印可训练参数：使用 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 = rself.lora_alpha = lora_alpha# 初始化 A 和 Bif r > 0:self.lora_A = nn.Parameter(torch.randn(r, in_features))self.lora_B = nn.Parameter(torch.zeros(out_features, r)) # B 初始化为全 0self.scaling = self.lora_alpha / self.rdef forward(self, x: torch.Tensor):result = F.linear(x, self.weight, bias=self.bias) # W @ xif self.r > 0:result += (self.lora_B @ self.lora_A @ x.transpose(-2, -1) # (B @ A) @ x).transpose(-2, -1) * self.scalingreturn result

代码解释：

1. _find_and_replace 函数：遍历模型的每一层，找到需要应用 LoRA 的线性层 (例如，q_proj, v_proj)，并将其替换为 Linear 层。
2. Linear 类：继承自 nn.Linear，并添加了 LoRA 的参数 lora_A 和 lora_B。
   • lora_A 初始化为随机值。
   • lora_B 初始化为全 0，这是为了保证在训练开始时，LoRA 部分的输出为 0，不影响预训练模型的原始行为。
   • scaling 是一个缩放因子，用于控制 LoRA 模块的权重。
3. 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'，然后像使用原始的预训练模型一样进行部署，无需额外的计算开销。