Demos/12. 应用 LoRA 到大模型的简单示例（PEFT）.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "347c39eb-0974-4880-b6dd-0cbfc33d7145",
   "metadata": {},
   "source": [
    "# 简介\n",
    "\n",
    "> 指导文章：[14. PEFT：在大模型中快速应用 LoRA](https://github.com/Hoper-J/AI-Guide-and-Demos-zh_CN/blob/master/Guide/14.%20PEFT：在大模型中快速应用%20LoRA.md)\n",
    "\n",
    "在线链接：[Kaggle](https://www.kaggle.com/code/aidemos/12-lora-peft) | [Colab](https://colab.research.google.com/drive/1-gWfn9xslSq6WlYDS9cinnyDEhBhjte4?usp=sharing)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07c339a9-af32-494f-88f0-1a5859da483d",
   "metadata": {},
   "source": [
    "## 安装必要的库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64c27f94-265c-49f3-a624-7f089f426577",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install transformers peft"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2006b231-0d05-444a-bdf8-4dcfed09f1c0",
   "metadata": {},
   "source": [
    "## 加载预训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8899fbab-45d1-43be-b198-403903d36900",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "# 加载预训练的 GPT-2 模型和分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained('gpt2')\n",
    "model = AutoModelForCausalLM.from_pretrained('gpt2')\n",
    "\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9c18030-1a96-4ba9-a157-bc4bf9293352",
   "metadata": {},
   "source": [
    "## 使用 PEFT 应用 LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd0f0a4-f692-47c1-b153-3eb5f38aded4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "\n",
    "# 配置 LoRA\n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,  # 任务类型：因果语言模型\n",
    "    inference_mode=False,          # 推理模式关闭，以进行训练\n",
    "    r=8,                           # 低秩值 r\n",
    "    lora_alpha=32,                 # LoRA 的缩放因子\n",
    "    lora_dropout=0.1,              # Dropout 概率\n",
    ")\n",
    "\n",
    "# 将 LoRA 应用到模型中\n",
    "model = get_peft_model(model, lora_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2658099e-7627-4950-b9ce-31167a9f6b3c",
   "metadata": {},
   "source": [
    "## 查看当前模型架构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b200673-a8e1-42bf-ad23-a6944ee3f5c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1a89d20-f03d-4938-94c8-42c79dae4fc8",
   "metadata": {},
   "source": [
    "## 查看增加的参数量\n",
    "\n",
    "应用 LoRA 后，我们一般都希望了解模型参数量的变化，它的计算其实很简单，\n",
    "\n",
    "### 理论计算\n",
    "\n",
    "对于每个应用了 LoRA 的层，增加的参数量为：\n",
    "\n",
    "$$\n",
    "\\text{增加的参数量} = r \\times (\\text{输入维度} + \\text{输出维度})\n",
    "$$\n",
    "\n",
    "- **`r`**：LoRA 的低秩值。\n",
    "- **输入维度**：层的输入特征数。\n",
    "- **输出维度**：层的输出特征数。\n",
    "\n",
    "### 使用 PEFT 查看参数\n",
    "\n",
    "`peft` 提供了查看模型参数的便捷方法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db8ddcf3-1ea0-4bfc-989e-890608e65e69",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 查看 LoRA 模块\n",
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6f2a584-f2f3-49c3-a010-0dbe3962fe41",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_trainable_parameters(model):\n",
    "    trainable_params = 0\n",
    "    all_params = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        num_params = param.numel()\n",
    "        all_params += num_params\n",
    "        if param.requires_grad:\n",
    "            trainable_params += num_params\n",
    "    print(f\"可训练参数量: {trainable_params}\")\n",
    "    print(f\"总参数量: {all_params}\")\n",
    "    print(f\"可训练参数占比: {100 * trainable_params / all_params:.2f}%\")\n",
    "    \n",
    "print_trainable_parameters(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91fe2425-7bf6-4fb3-be8e-a7f66257828e",
   "metadata": {},
   "source": [
    "## 准备数据并进行微调\n",
    "\n",
    "假设你已经有了训练数据集 `train_dataset`，下面是一个简单的样例代码。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6289cbbe-db23-4858-b58f-7b937e93afaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "# 定义训练参数\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',         # 模型保存和日志输出的目录路径\n",
    "    num_train_epochs=3,             # 训练的总轮数（epochs）\n",
    "    per_device_train_batch_size=16, # 每个设备（如GPU或CPU）上的训练批次大小，16表示每次输入模型的数据数量\n",
    "    learning_rate=5e-5,             # 学习率\n",
    "    logging_steps=10,               # 每隔多少步（steps）进行一次日志记录\n",
    "    save_steps=100,                 # 每隔多少步保存模型\n",
    ")\n",
    "\n",
    "# 创建 Trainer\n",
    "trainer = Trainer(\n",
    "    model=model,                    # 训练的模型对象，需要事先加载好\n",
    "    args=training_args,             # 上面定义的训练参数配置\n",
    "    train_dataset=train_dataset,    # 需要对应替换成已经处理过的dataset\n",
    ")\n",
    "\n",
    "# 开始训练\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c70f7e3-a20b-4e01-8931-cfcc1e655cb1",
   "metadata": {},
   "source": [
    "### 保存和加载 LoRA 微调的模型\n",
    "\n",
    "训练完成后，你可以保存或者加载 LoRA 微调的参数，下面是个简单的示例。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f14e02e8-d3ab-457a-b8cc-0eb8ec715bb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存 LoRA 参数\n",
    "model.save_pretrained('./lora_model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d9771b8-de83-440c-9fa7-b1f1506ed64f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载原始模型\n",
    "base_model = AutoModelForCausalLM.from_pretrained(\"gpt2\")\n",
    "\n",
    "# 加载 LoRA 参数\n",
    "from peft import PeftModel\n",
    "\n",
    "model = PeftModel.from_pretrained(base_model, './lora_model')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c89196ee-0362-4a4d-ac4e-2afda20d8edb",
   "metadata": {},
   "source": [
    "### 合并 LoRA 权重并卸载 PEFT 包装\n",
    "\n",
    "在完成微调后，可以使用 `merge_and_unload()` 将 LoRA 的权重合并回原始模型。这在部署和推理阶段非常有用，因为这样可以：\n",
    "\n",
    "- **减少依赖**：合并后，模型成为标准的 `transformers` 模型，不再需要 `peft` 库。\n",
    "- **提高推理效率**：减少了额外的计算开销，推理速度可能会有所提升。\n",
    "- **简化模型保存和加载**：不需要分别保存基础模型和 LoRA 参数。\n",
    "\n",
    "运行下面的代码："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bd8e84e-f821-4e38-a5aa-66a37ae53866",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 对比合并前后的模型\n",
    "print(\"合并前的模型结构：\")\n",
    "print(model)\n",
    "\n",
    "# 合并并卸载 LoRA 权重\n",
    "model = model.merge_and_unload()\n",
    "\n",
    "print(\"合并后的模型结构：\")\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23163a34-8a1d-42d6-8c33-dc5a8b1ce163",
   "metadata": {},
   "source": [
    "你应该注意到，合并后模型的 LoRA 层将被去除。\n",
    "\n",
    "现在，你可以像保存普通模型一样保存："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d90a5c-d65e-4b40-82ce-b75f6ef32037",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存合并后的模型\n",
    "model.save_pretrained('./merged_model')\n",
    "tokenizer.save_pretrained('./merged_model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebba1ffd-4faf-4999-bc6c-56b34b8f9a53",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "# 加载合并后的模型\n",
    "tokenizer = AutoTokenizer.from_pretrained('./merged_model')\n",
    "model = AutoModelForCausalLM.from_pretrained('./merged_model')\n",
    "\n",
    "# 进行推理\n",
    "inputs = tokenizer(\"Hello, World！\", return_tensors=\"pt\")\n",
    "outputs = model.generate(**inputs)\n",
    "print(tokenizer.decode(outputs[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "025fd4c3-788c-4996-a7a2-723c823a3956",
   "metadata": {},
   "source": [
    "**注意**：\n",
    "\n",
    "- **不可逆操作**：合并操作是不可逆的。如果你之后还需要进一步微调 LoRA 参数，需确保在合并前备份模型。\n",
    "- **无需 PEFT 库**：合并后的模型不再包含 LoRA 适配器（Adapter）的信息，因此在加载时无需使用 `PeftModel`。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2426a1-d564-4d34-b4b2-5ff14970063e",
   "metadata": {},
   "source": [
    "## 可能的错误及解决方案（TypeError: Expected state_dict to be dict-like...）\n",
    "\n",
    "在使用 `PEFT` 和 `LoRA` 进行模型微调和保存加载时，可能会遇到如下错误：\n",
    "\n",
    "```\n",
    "TypeError: Expected state_dict to be dict-like, got <class 'peft.peft_model.PeftModel'>.\n",
    "```\n",
    "\n",
    "### 错误原因\n",
    "\n",
    "一般是因为混合使用不同的保存和加载方式，这个错误不局限于 `PeftModel`，问题出在你用 `torch.save(model)` 保存整个模型却用 `load_state_dict()` 去加载，注意模型加载和保存的一致性。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e464c18f-0ec5-46a0-b01e-4d4c0d336977",
   "metadata": {},
   "source": [
    "### 错误重现\n",
    "\n",
    "下面我们来复现它，看是不是和你的操作一致（这里以 `PeftModel` 举例）：\n",
    "\n",
    "#### 1. 错误地保存整个 `PeftModel` 对象而不是其 `state_dict`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c2278f2-460e-4986-84b0-0fb3f1488261",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "from peft import get_peft_model, LoraConfig, TaskType\n",
    "\n",
    "# 加载预训练模型和分词器\n",
    "tokenizer = AutoTokenizer.from_pretrained('gpt2')\n",
    "model = AutoModelForCausalLM.from_pretrained('gpt2')\n",
    "\n",
    "# 配置 LoRA\n",
    "lora_config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    inference_mode=False,\n",
    "    r=8,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.1,\n",
    ")\n",
    "\n",
    "# 应用 LoRA\n",
    "model = get_peft_model(model, lora_config)\n",
    "\n",
    "# 错误地保存整个 PeftModel 对象\n",
    "torch.save(model, './model')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f55b963-d0b0-4ae0-b19a-b6962706edad",
   "metadata": {},
   "source": [
    "#### 2. 加载时传入 `PeftModel` 对象"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa3e70a7-ef3d-4e44-8b37-9d746f81ba66",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化模型\n",
    "model = AutoModelForCausalLM.from_pretrained(\"gpt2\")\n",
    "\n",
    "# 错误地加载模型，期望接收 state_dict 但实际加载了整个模型对象\n",
    "model.load_state_dict(torch.load('./model'))  # 这里会报错"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14b37526-74ec-4828-87e9-173dfabe3150",
   "metadata": {},
   "source": [
    "### 解决方法\n",
    "\n",
    "确保你保存和加载的对象是一致的：\n",
    "\n",
    "- `torch.save(model, '...')` 对应于 `torch.load(model, '...')`。\n",
    "- `torch.save(model.state_dict(), '...')` 对应于 `model.load_state_dict(torch.load('...'))`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225f0c8f-071a-416d-89c9-b01ebd3dca44",
   "metadata": {},
   "source": [
    "## 一个非常刁钻的错误：使用 get_peft_model() 后再导入 LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d28d412-3042-4ee4-bff4-105c860b2f40",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "from copy import deepcopy\n",
    "from peft import get_peft_model, LoraConfig, PeftModel\n",
    "\n",
    "# 固定随机数种子，确保结果可复现\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# 定义一个简单的线性模型\n",
    "class LinearModel(nn.Module):\n",
    "    def __init__(self, input_size, output_size):\n",
    "        super(LinearModel, self).__init__()\n",
    "        self.linear = nn.Linear(input_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n",
    "\n",
    "# 实例化线性模型\n",
    "model = LinearModel(input_size=10, output_size=1)\n",
    "\n",
    "# 在应用 LoRA 之前深拷贝原始模型，确保后续公平比较\n",
    "original_model = deepcopy(model)\n",
    "\n",
    "# 配置 LoRA 参数\n",
    "config = LoraConfig(\n",
    "    inference_mode=False,\n",
    "    r=4,\n",
    "    lora_alpha=16,\n",
    "    target_modules=['linear'],\n",
    ")\n",
    "\n",
    "# 将 LoRA 应用到模型中\n",
    "lora_model = get_peft_model(model, config)\n",
    "\n",
    "# 定义一个简单的损失函数和优化器\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = Adam(lora_model.parameters(), lr=1e-3)\n",
    "\n",
    "# 生成一些模拟的训练数据\n",
    "input_data = torch.randn(100, 10)  # 100 个样本，每个样本有 10 个特征\n",
    "target_data = torch.randn(100, 1)  # 对应的目标值\n",
    "\n",
    "# 训练一个回合\n",
    "lora_model.train()\n",
    "for epoch in range(1):  # 训练 1 个回合\n",
    "    optimizer.zero_grad()\n",
    "    outputs = lora_model(input_data)\n",
    "    loss = criterion(outputs, target_data)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "# 训练后保存 LoRA 权重\n",
    "lora_model.save_pretrained('linear_lora_model')\n",
    "\n",
    "# 方法 1：先使用 get_peft_model，再加载 LoRA 权重\n",
    "model1 = PeftModel.from_pretrained(get_peft_model(deepcopy(original_model), config), 'linear_lora_model')\n",
    "\n",
    "# 方法 2：直接加载 LoRA 权重\n",
    "model2 = PeftModel.from_pretrained(deepcopy(original_model), 'linear_lora_model')\n",
    "\n",
    "# 生成相同的输入数据以进行输出比较\n",
    "test_input = torch.randn(1, 10)\n",
    "\n",
    "# 比较四个模型的输出（原始模型，LoRA，方法1，方法2）\n",
    "def compare_model_outputs(input_data):\n",
    "    # 原始模型\n",
    "    original_output = original_model(input_data)\n",
    "    print(\"原始模型输出:\", original_output.detach().numpy())\n",
    "\n",
    "    # 训练后的 LoRA 模型\n",
    "    lora_output = lora_model(input_data)\n",
    "    print(\"训练后的 LoRA 模型输出:\", lora_output.detach().numpy())\n",
    "\n",
    "    # 方法 1：先使用 get_peft_model，再加载 LoRA\n",
    "    output1 = model1(input_data)\n",
    "    print(\"方法 1（先使用 get_peft_model，再加载 LoRA）输出:\", output1.detach().numpy())\n",
    "\n",
    "    # 方法 2：直接加载 LoRA\n",
    "    output2 = model2(input_data)\n",
    "    print(\"方法 2（直接加载 LoRA）输出:\", output2.detach().numpy())\n",
    "\n",
    "    if torch.allclose(original_output, output1):\n",
    "        print(\"\\n原始模型和方法 1 输出相同。\")\n",
    "    if torch.allclose(lora_output, output2):\n",
    "        print(\"训练后的 LoRA 模型和方法 2 输出相同。\\n\")\n",
    "\n",
    "# 比较两个模型的参数\n",
    "def compare_params(m1, m2):\n",
    "    for (n1, p1), (n2, p2) in zip(m1.named_parameters(), m2.named_parameters()):\n",
    "        if n1 != n2 or not torch.allclose(p1, p2):\n",
    "            print(f\"参数不匹配: \\n{n1}\\n{n2}\")\n",
    "            return False\n",
    "    return True\n",
    "\n",
    "# 比较四个模型的输出\n",
    "compare_model_outputs(test_input)\n",
    "\n",
    "# 检查方法 1 和方法 2 的参数是否一致\n",
    "if compare_params(model1, model2):\n",
    "    print(\"方法 1 和方法 2 的 LoRA 模型参数一致！\")\n",
    "else:\n",
    "    print(\"方法 1 和方法 2 的 LoRA 模型参数不一致！\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6617e10-78fd-403c-9f69-d034a6a1f086",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}