Build a Large Language Model (From Scratch) 附录 E 使用 LoRA 进行参数高效微调

LoRA（低秩自适应，Low-Rank Adaptation）是应用最广泛的参数高效微调技术之一。以下讨论基于第 6 章中提供的垃圾消息分类微调示例。然而，LoRA 微调同样适用于第 7 章中讨论的监督指令微调。

E.1 LoRA 简介

LoRA 是一种通过仅调整模型权重参数的一小部分，使预训练模型更好地适应特定且通常较小的数据集的技术。"低秩"指的是将模型调整限制在总权重参数空间的较小维度子空间，从而有效捕获训练过程中对权重参数变化影响最大的方向。LoRA 方法之所以有用且广受欢迎，是因为它能够高效地对大模型进行特定任务的微调，显著降低了通常所需的计算成本和资源。

假设一个大型权重矩阵 WWW 与特定层相关联，LoRA 可以应用于大语言模型中的所有线性层。不过，为了便于说明，我们将重点关注单个层。

在训练深度神经网络时，反向传播过程中我们会学习一个矩阵 ΔW\Delta WΔW，其中包含了更新原始权重参数以最小化训练期间损失函数所需的信息。

说明： 接下来，我们会将"权重"一词用作"模型权重参数"的简写。

在常规训练和微调中，权重更新定义如下：

Wupdated=W+ΔW W_{\text{updated}} = W + \Delta W Wupdated=W+ΔW

LoRA 方法由 Hu 等人提出，它通过学习权重更新 ΔW\Delta WΔW 的近似值，提供了一种更有效的替代方法：

ΔW≈AB \Delta W \approx AB ΔW≈AB

在此方法中，AAA 和 BBB 是两个比 WWW 小得多的矩阵，ABABAB 表示 AAA 和 BBB 之间的矩阵乘法。

利用 LoRA，可以重新定义之前的权重更新公式：

Wupdated=W+AB W_{\text{updated}} = W + AB Wupdated=W+AB

图 E-1 并列展示了全量微调与 LoRA 的权重更新方法。

如果仔细观察，你可能会注意到图 E-1 中全量微调和 LoRA 的可视化表示与之前提出的公式略有不同。这种变化是由于矩阵乘法的分配律，它允许我们将原始权重与更新后的权重分开，而不是将它们组合在一起。

例如，在全景微调中，当以 xxx 作为输入数据时，可以将计算表示为如下形式：

x(W+ΔW)=xW+xΔW x(W + \Delta W) = xW + x\Delta W x(W+ΔW)=xW+xΔW

同样，可以将 LoRA 表示为以下公式：

x(W+AB)=xW+xAB x(W + AB) = xW + xAB x(W+AB)=xW+xAB

除了可以减少训练期间需要更新的权重数量，将 LoRA 权重矩阵与原始模型权重分开的能力使 LoRA 在实践中更加有用。实际上，这一特性允许预训练的模型权重保持不变，并且在使用模型时可以动态地应用 LoRA 矩阵。

保持 LoRA 的权重分离在实践中非常有用，因为它使得模型定制变得更加灵活，无须存储多个完整版本的大语言模型。这降低了存储需求并提高了可扩展性，因为在为每个特定客户或应用程序进行定制时，只需调整和保存较小的 LoRA 矩阵即可。

接下来，让我们看看如何使用 LoRA 对大语言模型进行微调以实现垃圾消息分类，就像第 6 章中的微调示例那样。

E.2 准备数据集

在将 LoRA 应用于垃圾消息分类示例之前，需要加载所使用的数据集和预训练模型。这里的代码与第 6 章中的数据准备部分相同。（可以直接打开并运行第 6 章的 Jupyter Notebook，然后在其中插入 E.4 节的 LoRA 代码，而无须重复代码。）

首先，下载数据集并将其保存为 CSV 文件，如代码清单 E-1 所示。

python 复制代码

# import urllib
import requests
from pathlib import Path
import pandas as pd
from previous_chapters import (
    download_and_unzip_spam_data,
    create_balanced_dataset,
    random_split
)

url = "https://archive.ics.uci.edu/static/public/228/sms+spam+collection.zip"
zip_path = "sms_spam_collection.zip"
extracted_path = "sms_spam_collection"
data_file_path = Path(extracted_path) / "SMSSpamCollection.tsv"


try:
    download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path)
except (requests.exceptions.RequestException, TimeoutError) as e:
    print(f"Primary URL failed: {e}. Trying backup URL...")
    url = "https://f001.backblazeb2.com/file/LLMs-from-scratch/sms%2Bspam%2Bcollection.zip"
    download_and_unzip_spam_data(url, zip_path, extracted_path, data_file_path)

df = pd.read_csv(data_file_path, sep="\t", header=None, names=["Label", "Text"])
balanced_df = create_balanced_dataset(df)
balanced_df["Label"] = balanced_df["Label"].map({"ham": 0, "spam": 1})

train_df, validation_df, test_df = random_split(balanced_df, 0.7, 0.1)
train_df.to_csv("train.csv", index=None)
validation_df.to_csv("validation.csv", index=None)
test_df.to_csv("test.csv", index=None)

接下来，创建 SpamDataset 实例，如代码清单 E-2 所示。

python 复制代码

import torch
import tiktoken
from previous_chapters import SpamDataset


tokenizer = tiktoken.get_encoding("gpt2")
train_dataset = SpamDataset("train.csv", max_length=None, tokenizer=tokenizer)
val_dataset = SpamDataset("validation.csv", max_length=train_dataset.max_length, tokenizer=tokenizer)
test_dataset = SpamDataset("test.csv", max_length=train_dataset.max_length, tokenizer=tokenizer)

创建 PyTorch 数据集对象后，实例化数据加载器，如代码清单 E-3 所示。

python 复制代码

from torch.utils.data import DataLoader

num_workers = 0
batch_size = 8

torch.manual_seed(123)

train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=num_workers,
    drop_last=True,
)

val_loader = DataLoader(
    dataset=val_dataset,
    batch_size=batch_size,
    num_workers=num_workers,
    drop_last=False,
)

test_loader = DataLoader(
    dataset=test_dataset,
    batch_size=batch_size,
    num_workers=num_workers,
    drop_last=False,
)

在验证步骤中，迭代数据加载器，并检查每个批次是否包含 8 个训练示例，且每个示例由 120 个词元组成：

python 复制代码

print("Train loader:")
for input_batch, target_batch in train_loader:
    pass

print("Input batch dimensions:", input_batch.shape)
print("Label batch dimensions", target_batch.shape)

复制代码

Train loader:
Input batch dimensions: torch.Size([8, 120])
Label batch dimensions torch.Size([8])

最后，打印每个数据集中的批次总数：

python 复制代码

print(f"{len(train_loader)} training batches")
print(f"{len(val_loader)} validation batches")
print(f"{len(test_loader)} test batches")

复制代码

130 training batches
19 validation batches
38 test batches

模型只取最后一个词元位置的输出 logits 来预测类别，损失是普通的交叉熵分类损失，也就是第六章的分类微调。根本用不到第七章的指令微调：「右移构造 target」和「屏蔽 padding 损失」这套机制。

E.3 初始化模型

我们复用第 6 章中的代码来加载和准备预训练的 GPT 模型。首先，下载模型权重，并将其加载到 GPTModel 类中，如代码清单 E-4 所示。

python 复制代码

from gpt_download import download_and_load_gpt2
from previous_chapters import GPTModel, load_weights_into_gpt
# Alternatively:
# from llms_from_scratch.ch04 import GPTModel
# from llms_from_scratch.ch05 import load_weights_into_gpt



CHOOSE_MODEL = "gpt2-small (124M)"
INPUT_PROMPT = "Every effort moves"

BASE_CONFIG = {
    "vocab_size": 50257,     # Vocabulary size
    "context_length": 1024,  # Context length
    "drop_rate": 0.0,        # Dropout rate
    "qkv_bias": True         # Query-key-value bias
}

model_configs = {
    "gpt2-small (124M)": {"emb_dim": 768, "n_layers": 12, "n_heads": 12},
    "gpt2-medium (355M)": {"emb_dim": 1024, "n_layers": 24, "n_heads": 16},
    "gpt2-large (774M)": {"emb_dim": 1280, "n_layers": 36, "n_heads": 20},
    "gpt2-xl (1558M)": {"emb_dim": 1600, "n_layers": 48, "n_heads": 25},
}

BASE_CONFIG.update(model_configs[CHOOSE_MODEL])

model_size = CHOOSE_MODEL.split(" ")[-1].lstrip("(").rstrip(")")
settings, params = download_and_load_gpt2(model_size=model_size, models_dir="gpt2")

model = GPTModel(BASE_CONFIG)
load_weights_into_gpt(model, params)
model.eval();

为了确保模型加载正确，仔细检查一下它是否可以生成连贯的文本：

python 复制代码

from previous_chapters import (
    generate_text_simple,
    text_to_token_ids,
    token_ids_to_text
)


text_1 = "Every effort moves you"

token_ids = generate_text_simple(
    model=model,
    idx=text_to_token_ids(text_1, tokenizer),
    max_new_tokens=15,
    context_size=BASE_CONFIG["context_length"]
)

print(token_ids_to_text(token_ids, tokenizer))

以下输出显示模型生成了连贯的文本，这表明模型权重已正确加载：

复制代码

Every effort moves you forward.

The first step is to understand the importance of your work

接下来，准备进行分类微调的模型，类似于第 6 章，我们会替换输出层：

python 复制代码

torch.manual_seed(123)

num_classes = 2
model.out_head = torch.nn.Linear(in_features=768, out_features=num_classes)

print("Device:", device)
model.to(device)

最后，计算未微调模型的初始分类准确率（预计约为 50%，这表明该模型无法可靠地区分垃圾消息和非垃圾消息）：

python 复制代码

from previous_chapters import calc_accuracy_loader
# Alternatively:
# from llms_from_scratch.ch06 import calc_accuracy_loader

torch.manual_seed(123)
train_accuracy = calc_accuracy_loader(train_loader, model, device, num_batches=10)
val_accuracy = calc_accuracy_loader(val_loader, model, device, num_batches=10)
test_accuracy = calc_accuracy_loader(test_loader, model, device, num_batches=10)

print(f"Training accuracy: {train_accuracy*100:.2f}%")
print(f"Validation accuracy: {val_accuracy*100:.2f}%")
print(f"Test accuracy: {test_accuracy*100:.2f}%")

复制代码

Training accuracy: 46.25%
Validation accuracy: 45.00%
Test accuracy: 48.75%

E.4 使用 LoRA 进行参数高效微调

接下来，我们将使用 LoRA 来调整或微调大语言模型。我们首先初始化一个 LoRA 层，它创建了矩阵 A 和 B，并设置了 alpha 缩放因子和 rank（r）。该层可以接受输入并计算相应的输出，如图 E-2 所示。

在代码中，该 LoRA 层可以按如下方式（参见代码清单 E-5）实现。

python 复制代码

class LoRALayer(torch.nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha):
        super().__init__()
        self.A = torch.nn.Parameter(torch.empty(in_dim, rank))
        # A 用 Kaiming 初始化（与标准权重一致），B 初始化为 0
        # 保证训练开始时 A @ B = 0，即初始不改变原模型行为
        torch.nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
        self.B = torch.nn.Parameter(torch.zeros(rank, out_dim))

        # rank、alpha 是超参数：训练和推理全程固定，不参与梯度更新
        # 缩放系数 s = alpha / rank，是 LoRA 增量的"音量旋钮"
        #
        # 为什么 s 同时影响前向和梯度？因为前向 output = s·(x·A·B)，s 是常数因子。
        # 对 B 求梯度时链式法则会把它原样提出来：
        #   ∂L/∂B = (∂L/∂output)·(∂output/∂B) = s·[(∂L/∂output)·(x·A)]
        # 方括号是"没缩放时的原梯度"，前面多乘一个 s（对 A 同理）。
        # 结论：s 放大前向影响的同时，等比例放大梯度和每步更新量。
        #
        # 设 rank=8，x@A@B 输出=2.0，原模型输出=10.0：
        #   alpha=16：s=2.0 → 增量=4.0，输出=14.0；梯度也×2.0，学得快
        #   alpha=8 ：s=1.0 → 增量=2.0，输出=12.0；梯度 ×1.0
        #   alpha=4 ：s=0.5 → 增量=1.0，输出=11.0；梯度也 ×0.5，学得慢（更温柔）
        # 即 alpha 越小，影响越弱、更新越慢，两者同向。作用与学习率重叠，常固定一个调另一个。
        self.alpha = alpha
        self.rank = rank

    def forward(self, x):
        # 除以 rank 做归一化（原书没有这步，非必须但更规范）：
        # A@B 是 rank 个秩1外积之和，rank 越大输出量级越大。设每项量级为 c：
        #   不缩放：rank=4 → 4c，rank=64 → 64c（差 16 倍，换 rank 就得重调学习率）
        #   /rank ：rank=4 → (alpha/4)*4c=alpha*c，rank=64 → (alpha/64)*64c=alpha*c
        # 归一化后量级恒为 alpha*c，换 rank 可复用同一学习率
        # （更精确的方差归一化可参考 rsLoRA 的 alpha/sqrt(rank)）
        x = (self.alpha / self.rank) * (x @ self.A @ self.B)
        return x

rank 控制着矩阵 A 和 B 的内部维度。本质上，该设置决定了 LoRA 引入的额外参数量，从而在模型的适应性和效率之间建立平衡。

另一个重要的设置是 alpha，它作为低秩自适应输出的缩放因子，主要决定了适应层的输出对原始层输出的影响程度。这可以被视为调节低秩适应对层输出影响的一种方式。到目前为止，我们实现的 LoRALayer 类使我们能够对层的输入进行转换。

在 LoRA 中，典型的目标是替换现有的线性层，从而允许权重更新直接应用于已有的预训练权重，如图 E-3 所示。

为了整合原始线性层的权重，现在创建一个 LinearWithLoRA 层。该层利用之前实现的 LoRALayer，旨在替换神经网络中现有的线性层，比如 GPTModel 中的自注意力模块或前馈模块，如代码清单 E-6 所示。

python 复制代码

class LinearWithLoRA(torch.nn.Module):
    def __init__(self, linear, rank, alpha):
        super().__init__()
        self.linear = linear
        self.lora = LoRALayer(
            linear.in_features, linear.out_features, rank, alpha
        )

    def forward(self, x):
        return self.linear(x) + self.lora(x)

上述代码将标准线性层与 LoRALayer 结合在了一起。forward 方法通过将原始线性层和 LoRA 层的结果相加来计算输出。

由于权重矩阵 B（LoRALayer 中的 self.B）被初始化为零值，因此矩阵 A 和 B 的乘积将产生零矩阵。这确保了乘法不会改变原始权重，因为加零不会对它们产生影响。

为了将 LoRA 应用到之前定义的 GPTModel 中，我们引入了 replace_linear_ with_lora 函数。该函数会将模型中所有现有的 Linear 层替换为新创建的 LinearWithLoRA 层。

python 复制代码

def replace_linear_with_lora(model, rank, alpha):
    for name, module in model.named_children():
        if isinstance(module, torch.nn.Linear):
            # Replace the Linear layer with LinearWithLoRA
            setattr(model, name, LinearWithLoRA(module, rank, alpha))
        else:
            # Recursively apply the same function to child modules
            replace_linear_with_lora(module, rank, alpha)

现在我们已经实现了所有必要的代码，将 GPTModel 中的 Linear 层替换为新开发的 LinearWithLoRA 层，以便进行参数高效微调。接下来，我们将把 LinearWithLoRA 应用到 GPTModel 的多头注意力模块、前馈模块以及输出层中的所有 Linear 层，如图 E-4 所示。

在使用 LinearWithLoRA 层升级之前，首先冻结原始模型的参数：

python 复制代码

total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total trainable parameters before: {total_params:,}")

for param in model.parameters():
    param.requires_grad = False

total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total trainable parameters after: {total_params:,}")

现在，可以看到在 1.24 亿个模型参数中，没有一个是可训练的：

复制代码

Total trainable parameters before: 124,441,346
Total trainable parameters after: 0

接下来，使用 replace_linear_with_lora 来替换 Linear 层：

python 复制代码

replace_linear_with_lora(model, rank=16, alpha=16)

total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Total trainable LoRA parameters: {total_params:,}")

在添加 LoRA 层后，可训练的模型参数量如下所示：【这些是新增的参数】

复制代码

Total trainable LoRA parameters: 2,666,528

如你所见，使用 LoRA 时，我们将可训练参数的数量减少到了原来的 1/50。将 rank 和 alpha 设置为 16 是一个不错的默认选择，但增加 rank 参数也很常见，这反过来会增加可训练参数的数量。通常选择将 alpha 设置为 rank 的一半、两倍或等于 rank 的值。

接下来，通过打印模型架构来验证各层是否已按预期修改：

python 复制代码

model.to(device)

print(model)

复制代码

GPTModel(
  (tok_emb): Embedding(50257, 768)
  (pos_emb): Embedding(1024, 768)
  (drop_emb): Dropout(p=0.0, inplace=False)
  (trf_blocks): Sequential(
    (0): TransformerBlock(
      (att): MultiHeadAttention(
        (W_query): LinearWithLoRA(
          (linear): Linear(in_features=768, out_features=768, bias=True)
          (lora): LoRALayer()
        )
        (W_key): LinearWithLoRA(
          (linear): Linear(in_features=768, out_features=768, bias=True)
          (lora): LoRALayer()
        )
        (W_value): LinearWithLoRA(
          (linear): Linear(in_features=768, out_features=768, bias=True)
          (lora): LoRALayer()
        )
        (out_proj): LinearWithLoRA(
          (linear): Linear(in_features=768, out_features=768, bias=True)
          (lora): LoRALayer()
        )
        (dropout): Dropout(p=0.0, inplace=False)
      )
...
    (linear): Linear(in_features=768, out_features=2, bias=True)
    (lora): LoRALayer()
  )
)
Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...

该模型现在包含新的 LinearWithLoRA 层，这些层由设置为不可训练的原始 Linear 层和新的 LoRA 层组成，我们将对后者进行微调。在开始微调模型之前，先计算一下初始分类准确率：

python 复制代码

torch.manual_seed(123)
train_accuracy = calc_accuracy_loader(train_loader, model, device, num_batches=10)
val_accuracy = calc_accuracy_loader(val_loader, model, device, num_batches=10)
test_accuracy = calc_accuracy_loader(test_loader, model, device, num_batches=10)

print(f"Training accuracy: {train_accuracy*100:.2f}%")
print(f"Validation accuracy: {val_accuracy*100:.2f}%")
print(f"Test accuracy: {test_accuracy*100:.2f}%")

复制代码

Training accuracy: 46.25%
Validation accuracy: 45.00%
Test accuracy: 48.75%

这些准确率值与第 6 章中的值相同。出现这一结果是因为 LoRA 矩阵 B 被初始化为零，因此矩阵 AB 的乘积产生了零矩阵。这确保了乘法不会改变原始权重，因为加零不会对它们产生影响。

现在，进入令人兴奋的部分------使用第 6 章中的训练函数来微调模型，如代码清单 E-7 所示。在 M3 MacBook Air 笔记本电脑上，训练大约需要 15 分钟；而在 V100 或 A100 GPU 上，训练时间不到半分钟。

python 复制代码

import time
from previous_chapters import train_classifier_simple
# Alternatively:
# from llms_from_scratch.ch06 import train_classifier_simple


start_time = time.time()

torch.manual_seed(123)

optimizer = torch.optim.AdamW(model.parameters(), lr=8e-4, weight_decay=0.1)

num_epochs = 5
train_losses, val_losses, train_accs, val_accs, examples_seen = train_classifier_simple(
    model, train_loader, val_loader, optimizer, device,
    num_epochs=num_epochs, eval_freq=50, eval_iter=5,
)

end_time = time.time()
execution_time_minutes = (end_time - start_time) / 60
print(f"Training completed in {execution_time_minutes:.2f} minutes.")

复制代码

Ep 1 (Step 000000): Train loss 3.820, Val loss 3.462
Ep 1 (Step 000050): Train loss 0.346, Val loss 0.325
Ep 1 (Step 000100): Train loss 0.063, Val loss 0.144
Training accuracy: 100.00% | Validation accuracy: 92.50%
Ep 2 (Step 000150): Train loss 0.054, Val loss 0.045
Ep 2 (Step 000200): Train loss 0.058, Val loss 0.122
Ep 2 (Step 000250): Train loss 0.041, Val loss 0.199
Training accuracy: 100.00% | Validation accuracy: 95.00%
Ep 3 (Step 000300): Train loss 0.020, Val loss 0.153
Ep 3 (Step 000350): Train loss 0.017, Val loss 0.186
Training accuracy: 100.00% | Validation accuracy: 95.00%
Ep 4 (Step 000400): Train loss 0.017, Val loss 0.099
Ep 4 (Step 000450): Train loss 0.001, Val loss 0.170
Ep 4 (Step 000500): Train loss 0.117, Val loss 0.222
Training accuracy: 97.50% | Validation accuracy: 92.50%
Ep 5 (Step 000550): Train loss 0.038, Val loss 0.235
Ep 5 (Step 000600): Train loss 0.019, Val loss 0.252
Training accuracy: 100.00% | Validation accuracy: 100.00%
Training completed in 2.16 minutes.

在这里，使用 LoRA 训练模型比不使用 LoRA 训练模型花费的时间更长（参见第 6 章），因为 LoRA 层在前向传播过程中引入了额外的计算。但是，对于较大的模型，反向传播的成本更高，此时使用 LoRA 的模型通常比不使用 LoRA 的模型训练速度更快。

如你所见，该模型经过了充分的训练，并且验证准确率非常高。下面让我们可视化损失曲线，以更好地观察训练是否收敛。

python 复制代码

from previous_chapters import plot_values
# Alternatively:
# from llms_from_scratch.ch06 import plot_values

epochs_tensor = torch.linspace(0, num_epochs, len(train_losses))
examples_seen_tensor = torch.linspace(0, examples_seen, len(train_losses))

plot_values(epochs_tensor, examples_seen_tensor, train_losses, val_losses, label="loss")

除了根据损失曲线评估模型，我们还计算了完整的训练集、验证集和测试集的准确率（在训练期间，我们通过设置 eval_iter=5 来估算来自 5 个批次的训练集和验证集的准确率）：

python 复制代码

train_accuracy = calc_accuracy_loader(train_loader, model, device)
val_accuracy = calc_accuracy_loader(val_loader, model, device)
test_accuracy = calc_accuracy_loader(test_loader, model, device)

print(f"Training accuracy: {train_accuracy*100:.2f}%")
print(f"Validation accuracy: {val_accuracy*100:.2f}%")
print(f"Test accuracy: {test_accuracy*100:.2f}%")

复制代码

Training accuracy: 99.81%
Validation accuracy: 97.99%
Test accuracy: 96.67%

这些结果表明，该模型在训练集、验证集和测试集上表现良好。训练集准确率达到100%，模型完美地学习了训练数据。然而，略低的验证集准确率和测试集准确率（分别为 96.64%和 98%）表明存在一定程度的过拟合，因为与训练集相比，该模型在未见过的数据上的泛化效果不佳。总体而言，考虑到我们只微调了相对较少的模型权重（LoRA 仅有 270 万个参数，而原先的模型参数量是 1.24 亿），这些结果令人印象非常深刻。