使用 LoRA 进行模型微调的步骤

使用 LoRA 进行模型微调的步骤

以下是关于如何在训练过程中应用 LoRA 的详细步骤

1. 安装所需的库

首先，确保已安装 PyTorch 和 Hugging Face 的 transformers 库：

pip install torch transformers

2. 定义 LoRA 层

定义一个 LoRA 模块，用于替换 transformer 中的标准线性层，通常在自注意力机制的 query、key 和 value 投影中使用

import torch
import torch.nn as nn

# 定义 LoRA 模块
class LoRA(nn.Module):
    def __init__(self, input_dim, output_dim, rank=8):
        super(LoRA, self).__init__()
        # LoRA 引入了两个额外的矩阵 W_down 和 W_up
        self.W_down = nn.Linear(input_dim, rank, bias=False)  # 低秩降维
        self.W_up = nn.Linear(rank, output_dim, bias=False)   # 低秩升维

    def forward(self, x):
        # 将低秩适配结果加到原始输出上
        return self.W_up(self.W_down(x))

3. 修改 Transformer 模型以应用 LoRA

将 transformer 中的 query、key 和 value 投影替换为 LoRA 模块

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载预训练模型
model_name = "gpt2"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 定义 LoRA 版的 Transformer 模型
class LoRATransformer(nn.Module):
    def __init__(self, base_model, lora_rank=8):
        super(LoRATransformer, self).__init__()
        self.base_model = base_model

        # 遍历 transformer 层并应用 LoRA
        for name, module in self.base_model.named_modules():
            # 对 self-attention 中的线性层应用 LoRA
            if isinstance(module, nn.Linear) and 'attn' in name:
                input_dim = module.in_features
                output_dim = module.out_features
                lora_module = LoRA(input_dim, output_dim, rank=lora_rank)
                # 替换原始的线性层为 LoRA 模块
                setattr(self.base_model, name, lora_module)

    def forward(self, input_ids, attention_mask=None):
        return self.base_model(input_ids, attention_mask=attention_mask)

# 初始化应用了 LoRA 的模型
lora_model = LoRATransformer(model)

4. 准备训练

在训练之前，冻结预训练模型的参数，仅更新 LoRA 模块中的参数。

# 冻结原始模型参数
for param in model.parameters():
    param.requires_grad = False

# 仅训练 LoRA 层的参数
for name, param in lora_model.named_parameters():
    if 'W_down' in name or 'W_up' in name:
        param.requires_grad = True

5. 开始训练模型

现在可以开始训练模型了，使用优化器来更新 LoRA 模块的参数

# 定义优化器，仅更新 LoRA 层的参数
optimizer = torch.optim.Adam([param for param in lora_model.parameters() if param.requires_grad], lr=1e-4)

# 损失函数（交叉熵损失）
loss_fn = nn.CrossEntropyLoss()

# 示例训练循环
for epoch in range(3):  # 调整 epoch 数量
    for batch in dataloader:  # 假设 'dataloader' 是已经准备好的数据加载器
        inputs = tokenizer(batch['text'], return_tensors='pt', padding=True, truncation=True)
        labels = inputs.input_ids.clone()  # 使用输入 tokens 作为标签

        # 前向传播
        outputs = lora_model(inputs.input_ids)
        logits = outputs.logits
        
        # 计算损失
        loss = loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
        
        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    print(f"第 {epoch+1} 轮训练，损失: {loss.item()}")

6. 保存与加载 LoRA 增强的模型

训练完成后，可以保存 LoRA 层的参数，并在未来加载这些参数进行推理或进一步训练。

# 保存 LoRA 模型的参数
torch.save(lora_model.state_dict(), "lora_model.pth")

# 加载 LoRA 模型的参数
lora_model.load_state_dict(torch.load("lora_model.pth"))

LoRA 专注于注意力层：LoRA 通常应用于 transformer 模型中的自注意力机制（例如 query、key 和 value 投影层）。
参数高效性：与全量微调相比，LoRA 仅引入少量可训练参数，从而显著减少了计算成本。
冻结预训练权重：预训练模型的原始权重保持不变，仅更新 LoRA 模块的低秩矩阵，从而实现高效微调。