【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()
# ... (加载数据，进行训练) ...
代码解释：
- 加载预训练模型：使用
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'，然后像使用原始的预训练模型一样进行部署，无需额外的计算开销。