LoRA训练推理部署流程

前言

这是LoRA相关的第三篇文章

上篇文章已经比较全面的讨论过了LoRA的一些配置和网络模型结构，本文主要从以下几个方面进一步看一下LoRA的微调、模型保存、加载和推理流程。

不同层配置不同的LoRA参数

首先来看一下上篇最后的问题 如何修改LoraConfig使得fc1和fc2中的超参不同呢？

当LoRA的配置文件这样配置时，fc1和fc2使用的是同一份超参数，打印结果如下

config = LoraConfig(
    r=4,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["fc1", "fc2"],
)

model = get_peft_model(model, config)
print(f"增加Lora策略的，模型参数量：{count_parameters(model)}")

可以看到fc1和fc2中额外的低秩矩阵r都为4。

base_model.model.fc1.lora_A.default.weight: torch.Size([4, 10])
base_model.model.fc1.lora_B.default.weight: torch.Size([20, 4])
base_model.model.fc2.lora_A.default.weight: torch.Size([4, 20])
base_model.model.fc2.lora_B.default.weight: torch.Size([2, 4])

现在我们修改配置文件
fc1层的配置r=4，dropout=0.1，fc2层的配置r=2，dropout=0.3。整体代码如下

import torch
import torch.nn as nn
from peft import get_peft_model, LoraConfig
from simple_mlp import SimpleMLP


def count_parameters(model):
    """计算可训练的参数"""
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(f"{name}: {param.size()}")
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


model = SimpleMLP()
# print(f"原始的网络结构：\n {model}")
# print(f"原始模型的参数量：{count_parameters(model)}")

# 为fc1层设置r=4
config_fc1 = LoraConfig(
    r=4,
    lora_alpha=16,
    lora_dropout=0.1,
    target_modules=["fc1"],
)

# 为fc2层设置r=2
config_fc2 = LoraConfig(
    r=2,
    lora_alpha=16,
    lora_dropout=0.3,
    target_modules=["fc2"],
)

# 应用LoRA到模型的不同层
model = get_peft_model(model, config_fc1)
model = get_peft_model(model, config_fc2)

print(f"Lora微调的网络结构：\n {model}")
print(f"增加Lora策略的，模型参数量：{count_parameters(model)}")

网络结构如下

Lora微调的网络结构：
 PeftModel(
  (base_model): LoraModel(
    (model): PeftModel(
      (base_model): LoraModel(
        (model): SimpleMLP(
          (fc1): lora.Linear(
            (base_layer): Linear(in_features=10, out_features=20, bias=True)
            (lora_dropout): ModuleDict(
              (default): Dropout(p=0.1, inplace=False)
            )
            (lora_A): ModuleDict(
              (default): Linear(in_features=10, out_features=4, bias=False)
            )
            (lora_B): ModuleDict(
              (default): Linear(in_features=4, out_features=20, bias=False)
            )
            (lora_embedding_A): ParameterDict()
            (lora_embedding_B): ParameterDict()
          )
          (relu): ReLU()
          (fc2): lora.Linear(
            (base_layer): Linear(in_features=20, out_features=2, bias=True)
            (lora_dropout): ModuleDict(
              (default): Dropout(p=0.3, inplace=False)
            )
            (lora_A): ModuleDict(
              (default): Linear(in_features=20, out_features=2, bias=False)
            )
            (lora_B): ModuleDict(
              (default): Linear(in_features=2, out_features=2, bias=False)
            )
            (lora_embedding_A): ParameterDict()
            (lora_embedding_B): ParameterDict()
          )
        )
      )
    )
  )
)
base_model.model.base_model.model.fc1.lora_A.default.weight: torch.Size([4, 10])
base_model.model.base_model.model.fc1.lora_B.default.weight: torch.Size([20, 4])
base_model.model.base_model.model.fc2.lora_A.default.weight: torch.Size([2, 20])
base_model.model.base_model.model.fc2.lora_B.default.weight: torch.Size([2, 2])
增加Lora策略的，模型参数量：164

正文

以下是本文主要的部分，可作为LoRA实战初步入门来学习。

LoRA是一种高效微调预训练大模型的技术，首先就需要一个训练好的模型，用于LoRA微调。
现在我们来造一个无实际用途的网络模型。仅为了说明本文想表达的整理流程。

定义模型文件

文件simple_mlp.py，定义两层Linear，中间使用relu激活函数。模拟从一个10 维空间映射到2维空间。

from torch import nn

class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(10, 20)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(20, 2)

    def forward(self, x):
        x = self.relu(self.fc1(x))
        return self.fc2(x))

训练一个原始模型

训练文件train.py，模拟一个存在预训练模型，将其保存在项目下的model/路径下。

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from simple_mlp import SimpleMLP

class SimpleDataset(Dataset):
    def __init__(self, num_samples=1000):
        self.x = torch.randn((num_samples, 10))
        self.y = torch.randint(0, 2, (num_samples,))
    def __len__(self):
        return len(self.x)
    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]
        
dataset = SimpleDataset()
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 初始化模型
model = SimpleMLP()
# 损失函数
criterion = nn.CrossEntropyLoss()
# 优化器
optimizer = optim.AdamW(model.parameters(), lr=0.001)
# 训练循环
num_epochs = 5
for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

# 保存整个模型
torch.save(model, "./model/model.pth")  # Size:3.4K
# 保存模型的状态字典
torch.save(model.state_dict(), "./model/model_state.pth")  # Size:2.6K

LoRA训练过程

文件lora_train.py，训练主要流程如下：

1、加载原始预训练模型
2、设置LoraConfig超参数
3、保存LoRA的权重

import torch
import torch.nn as nn
from peft import get_peft_model, LoraConfig
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
from simple_mlp import SimpleMLP


# 定义数据集
class SimpleDataset(Dataset):
    def __init__(self, num_samples=1000):
        self.x = torch.randn((num_samples, 10))
        self.y = torch.randint(0, 2, (num_samples,))
    def __len__(self):
        return len(self.x)
    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

dataset = SimpleDataset()
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 按照保存的方式加载模型
model = torch.load("./model/model.pth")

# model = SimpleMLP()
# model.load_state_dict(torch.load("./model/model_state.pth"))
# LoRA配置
peft_config = LoraConfig(
    r=4,
    lora_alpha=8,
    target_modules=["fc1"],
    lora_dropout=0.1,
)
model = get_peft_model(model, peft_config)
# 损失函数
criterion = nn.CrossEntropyLoss()
# 优化器
optimizer = optim.AdamW(model.parameters(), lr=0.001)
# 训练循环
num_epochs = 5
for epoch in range(num_epochs):
    for inputs, labels in dataloader:
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

model.save_pretrained("./checkpoint")

LoRA推理过程

推理文件predict.py，加载原始模型和LoRA微调的参数

import torch
from peft import PeftModel
from sample_input import sample_input
from simple_mlp import SimpleMLP

# 加载原始模型和LoRA权重
model = torch.load("/home/my_code/lora/model/model.pth")
lora_path = "/home/my_code/lora/checkpoint"

model = PeftModel.from_pretrained(model, model_id=lora_path)
with torch.no_grad():
    output = model(sample_input)
    predicted_class = torch.argmax(output, dim=1)
print("Predicted class:", predicted_class.item())

推理数据源，数据模拟文件sample_input.py

import torch

sample_input = torch.randn(1, 10)

# 测试使用固定数据
# sample_input = torch.tensor(
#     [[1.0477, 1.3758, 0.3937, 0.8138, 0.3244, -0.2521, -1.4973, 0.4456, -0.1439, -0.3920]], dtype=torch.float32
# )

LoRA常见保存的文件说明

正常我们会结合tranformers库对模型进行微调，此时Lora保存的checkpoint文件下会有很多文件如下所示：

adapter_config.json：LoRA 配置信息。
adapter_model.safetensors LoRA的权重参数。 
.safetensors 格式是一种更安全的二进制格式，用于存储 PyTorch 模型权重。
optimizer.pt 存储优化器的状态，用于恢复训练状态。
rng_state.pth 存储随机数生成器的状态，用于恢复随机种子。
scheduler.pt 存储学习率调度器的状态，用于恢复训练状态。
training_args.bin 存储训练参数，如学习率、批次大小等。
training_state.json 存储训练参数，如学习率、批次大小等。

项目结构如下

├── checkpoint	# lora微调保存的权重和配置
│   ├── adapter_config.json
│   ├── adapter_model.safetensors
│   └── README.md
├── lora_net_structure1.py # 参见上一篇文章
├── lora_net_structure2.py	# 针对每一层Linear设置Lora配置
├── lora_train.py # lora的训练文件
├── model	# 原始训练模型保存路径
│   ├── model.pth
│   └── model_state.pth
├── predict.py	# 推理文件
├── sample_input.py # 推理数据来源
├── simple_mlp.py	# 网络结构
└── train.py	# 原始模型训练文件

相关文章

LoRA微调基础知识点
LoRA微调模型结构可训参数和配置详解
LoRA训练推理部署流程
更新中…
[LoRA基于BERT模型微调实践]