使用 LoRA 微调 Llama2¶

本指南将教您了解 LoRA，一种参数高效的微调技术，并向您展示如何使用 torchtune 通过 LoRA 微调 Llama2 模型。如果您已经了解 LoRA 是什么，并想直接在 torchtune 中运行您自己的 LoRA 微调，您可以跳转到 torchtune 中的 LoRA 微调配方。

您将学到什么

LoRA 是什么以及它如何在微调期间节省内存
torchtune 中 LoRA 组件的概述
如何使用 torchtune 运行 LoRA 微调
如何尝试不同的 LoRA 配置

先决条件

熟悉 torchtune
确保安装 torchtune
确保您已下载 Llama2-7B 模型权重

什么是 LoRA？¶

LoRA 是一种基于适配器的参数高效微调方法，它将可训练的低秩分解矩阵添加到神经网络的不同层，然后冻结网络的其余参数。LoRA 最常应用于 Transformer 模型，在这种情况下，通常将低秩矩阵添加到每个 Transformer 层自注意力机制中的某些线性投影中。

注意

如果您不熟悉，请查看这些参考资料，了解秩的定义和低秩近似的讨论。

通过使用 LoRA 进行微调（而不是微调所有模型参数），您可以期望看到内存节省，因为具有梯度的参数数量大大减少。当使用带有动量的优化器（如 AdamW）时，您可以期望从优化器状态中看到进一步的内存节省。

注意

LoRA 内存节省主要来自梯度和优化器状态，因此如果您的模型的峰值内存出现在其 forward() 方法中，则 LoRA 可能不会减少峰值内存。

LoRA 是如何工作的？¶

LoRA 用低秩近似替换权重更新矩阵。通常，任意 nn.Linear(in_dim,out_dim) 层的权重更新的秩可能高达 min(in_dim,out_dim)。LoRA（以及其他相关论文，如 Aghajanyan et al.）假设 LLM 微调期间这些更新的内在维度实际上可以低得多。为了利用此属性，LoRA 微调将冻结原始模型，然后从低秩投影添加可训练的权重更新。更明确地说，LoRA 训练两个矩阵 A 和 B。A 将输入向下投影到更小的秩（在实践中通常为 4 或 8），而 B 向上投影回原始线性层输出的维度。

下图给出了完整微调（左侧）与 LoRA 微调（右侧）的单个权重更新步骤的简化表示。LoRA 矩阵 A 和 B 用作蓝色部分中全秩权重更新的近似值。

尽管 LoRA 在模型 forward() 中引入了一些额外的参数，但只有 A 和 B 矩阵是可训练的。这意味着，对于秩为 r 的 LoRA 分解，我们需要存储的梯度数量从 in_dim*out_dim 减少到 r*(in_dim+out_dim)。（请记住，通常 r 远小于 in_dim 和 out_dim。）

例如，在 7B Llama2 的自注意力机制中，Q、K 和 V 投影的 in_dim=out_dim=4096。这意味着秩为 r=8 的 LoRA 分解将给定投影的可训练参数数量从 \(4096 * 4096 \approx 15M\) 减少到 \(8 * 8192 \approx 65K\)，减少了 99% 以上。

让我们看一下 LoRA 在原生 PyTorch 中的最小实现。

import torch
from torch import nn

class LoRALinear(nn.Module):
  def __init__(
    self,
    in_dim: int,
    out_dim: int,
    rank: int,
    alpha: float,
    dropout: float
  ):
    # These are the weights from the original pretrained model
    self.linear = nn.Linear(in_dim, out_dim, bias=False)

    # These are the new LoRA params. In general rank << in_dim, out_dim
    self.lora_a = nn.Linear(in_dim, rank, bias=False)
    self.lora_b = nn.Linear(rank, out_dim, bias=False)

    # Rank and alpha are commonly-tuned hyperparameters
    self.rank = rank
    self.alpha = alpha

    # Most implementations also include some dropout
    self.dropout = nn.Dropout(p=dropout)

    # The original params are frozen, and only LoRA params are trainable.
    self.linear.weight.requires_grad = False
    self.lora_a.weight.requires_grad = True
    self.lora_b.weight.requires_grad = True

  def forward(self, x: torch.Tensor) -> torch.Tensor:
    # This would be the output of the original model
    frozen_out = self.linear(x)

    # lora_a projects inputs down to the much smaller self.rank,
    # then lora_b projects back up to the output dimension
    lora_out = self.lora_b(self.lora_a(self.dropout(x)))

    # Finally, scale by the alpha parameter (normalized by rank)
    # and add to the original model's outputs
    return frozen_out + (self.alpha / self.rank) * lora_out

还有一些关于初始化的其他细节，我们在此省略，但如果您想了解更多信息，可以查看我们在 LoRALinear 中的实现。现在我们了解了 LoRA 的工作原理，让我们看看如何将其应用于我们最喜欢的模型。

将 LoRA 应用于 Llama2 模型¶

使用 torchtune，我们可以轻松地将 LoRA 应用于具有各种不同配置的 Llama2。让我们看看如何在 torchtune 中构建带有和不带有 LoRA 的 Llama2 模型。

from torchtune.models.llama2 import llama2_7b, lora_llama2_7b

# Build Llama2 without any LoRA layers
base_model = llama2_7b()

# The default settings for lora_llama2_7b will match those for llama2_7b
# We just need to define which layers we want LoRA applied to.
# Within each self-attention, we can choose from ["q_proj", "k_proj", "v_proj", and "output_proj"].
# We can also set apply_lora_to_mlp=True or apply_lora_to_output=True to apply LoRA to other linear
# layers outside of the self-attention.
lora_model = lora_llama2_7b(lora_attn_modules=["q_proj", "v_proj"])

注意

仅调用 lora_llama_2_7b 不会处理可训练参数的定义。有关如何执行此操作，请参阅下方。

让我们更仔细地检查一下这些模型。

# Print the first layer's self-attention in the usual Llama2 model
>>> print(base_model.layers[0].attn)
MultiHeadAttention(
  (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (output_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (pos_embeddings): RotaryPositionalEmbeddings()
)

# Print the same for Llama2 with LoRA weights
>>> print(lora_model.layers[0].attn)
MultiHeadAttention(
  (q_proj): LoRALinear(
    (dropout): Dropout(p=0.0, inplace=False)
    (lora_a): Linear(in_features=4096, out_features=8, bias=False)
    (lora_b): Linear(in_features=8, out_features=4096, bias=False)
  )
  (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (v_proj): LoRALinear(
    (dropout): Dropout(p=0.0, inplace=False)
    (lora_a): Linear(in_features=4096, out_features=8, bias=False)
    (lora_b): Linear(in_features=8, out_features=4096, bias=False)
  )
  (output_proj): Linear(in_features=4096, out_features=4096, bias=False)
  (pos_embeddings): RotaryPositionalEmbeddings()
)

请注意，正如预期的那样，我们的 LoRA 模型的层在 Q 和 V 投影中包含额外的权重。此外，检查 lora_model 和 base_model 的类型会显示它们都是相同的 TransformerDecoder 的实例。（请随意自行验证。）

为什么这很重要？torchtune 使您可以轻松地直接从我们的 Llama2 模型加载 LoRA 的检查点，而无需任何包装器或自定义检查点转换逻辑。

# Assuming that base_model already has the pretrained Llama2 weights,
# this will directly load them into your LoRA model without any conversion necessary.
lora_model.load_state_dict(base_model.state_dict(), strict=False)

注意

每当使用 strict=False 加载权重时，您都应验证加载的 state_dict 中任何缺失或额外的键是否符合预期。torchtune 的 LoRA 配方通过 validate_missing_and_unexpected_for_lora() 默认执行此操作。

一旦我们加载了基础模型权重，我们还希望仅将 LoRA 参数设置为可训练。

from torchtune.modules.peft.peft_utils import get_adapter_params, set_trainable_params

# Fetch all params from the model that are associated with LoRA.
lora_params = get_adapter_params(lora_model)

# Set requires_grad=True on lora_params, and requires_grad=False on all others.
set_trainable_params(lora_model, lora_params)

# Print the total number of parameters
total_params = sum([p.numel() for p in lora_model.parameters()])
trainable_params = sum([p.numel() for p in lora_model.parameters() if p.requires_grad])
print(
  f"""
  {total_params} total params,
  {trainable_params}" trainable params,
  {(100.0 * trainable_params / total_params):.2f}% of all params are trainable.
  """
)

6742609920 total params,
4194304 trainable params,
0.06% of all params are trainable.

注意

如果您直接使用 LoRA 配方（如此处详述），您只需传递相关的检查点路径。加载模型权重和设置可训练参数将在配方中处理。

torchtune 中的 LoRA 微调配方¶

最后，我们可以将所有内容放在一起，并使用 torchtune 的 LoRA 配方微调模型。确保您首先按照这些说明下载 Llama2 权重和分词器。然后，您可以运行以下命令，使用两个 GPU（每个 GPU 的 VRAM 至少为 16GB）执行 Llama2-7B 的 LoRA 微调

tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora

注意

确保指向您的 Llama2 权重和分词器的位置。这可以通过添加 checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path 或直接修改 7B_lora.yaml 文件来完成。请参阅我们的“关于配置的一切”配方，了解有关如何轻松克隆和修改 torchtune 配置的更多详细信息。

注意

您可以根据 (a) 您可用的 GPU 数量和 (b) 您的硬件的内存限制来修改 nproc_per_node 的值。

前面的命令将使用 torchtune 的出厂设置运行 LoRA 微调，但我们可能希望进行一些实验。让我们仔细看一下 lora_finetune_distributed 配置的一些内容。

# Model Arguments
model:
  _component_: lora_llama2_7b
  lora_attn_modules: ['q_proj', 'v_proj']
  lora_rank: 8
  lora_alpha: 16
...

我们看到默认是将 LoRA 应用于秩为 8 的 Q 和 V 投影。一些关于 LoRA 的实验发现，将 LoRA 应用于自注意力机制中的所有线性层，并将秩增加到 16 或 32 可能是有益的。请注意，这可能会增加我们的最大内存，但只要我们保持 rank<<embed_dim，影响应该是相对较小的。

让我们运行这个实验。我们还可以增加 alpha（通常，将 alpha 和 rank 一起缩放是一个好的做法）。

tune run --nnodes 1 --nproc_per_node 2 lora_finetune_distributed --config llama2/7B_lora \
lora_attn_modules=['q_proj','k_proj','v_proj','output_proj'] \
lora_rank=32 lora_alpha=64 output_dir=./lora_experiment_1

下图可以看到此运行与我们的基线在前 500 步上的（平滑）损失曲线的比较。

../_images/lora_experiment_loss_curves.png

注意

上图是使用 W&B 生成的。您可以使用 torchtune 的 WandBLogger 生成类似的损失曲线，但您需要单独安装 W&B 并设置帐户。有关在 torchtune 中使用 W&B 的更多详细信息，请参阅我们的“记录到 Weights & Biases”配方。

使用 LoRA 权衡内存和模型性能¶

在前面的示例中，我们在两个设备上运行了 LoRA。但是考虑到 LoRA 的低内存占用，我们可以使用支持 bfloat16 浮点格式的大多数商用 GPU 在单个设备上运行微调。这可以通过以下命令完成

tune run lora_finetune_single_device --config llama2/7B_lora_single_device

在单个设备上，我们可能需要更加注意我们的峰值内存。让我们运行一些实验，看看我们在微调期间的峰值内存。我们将沿着两个轴进行实验：首先，哪些模型层应用了 LoRA，其次，每个 LoRA 层的秩。（我们将与 LoRA 秩并行缩放 alpha，如上所述。）

为了比较我们的实验结果，我们可以评估我们在 truthfulqa_mc2 上的模型，这是来自语言模型的 TruthfulQA 基准测试的任务。有关如何使用 torchtune 的 EleutherAI 评估工具集成运行此任务和其他评估任务的更多详细信息，请参阅我们的端到端工作流程教程。

以前，我们仅为每个自注意力模块中的线性层启用 LoRA，但实际上我们可以将 LoRA 应用于其他线性层：MLP 层和我们模型的最终输出投影。请注意，对于 Llama-2-7B，最终输出投影映射到词汇表维度（32000 而不是其他线性层中的 4096），因此为此层启用 LoRA 将比其他层更多地增加我们的峰值内存。我们可以对我们的配置进行以下更改

# Model Arguments
model:
  _component_: lora_llama2_7b
  lora_attn_modules: ['q_proj', 'k_proj', 'v_proj', 'output_proj']
  apply_lora_to_mlp: True
  apply_lora_to_output: True
...

注意

以下所有微调运行都使用 llama2/7B_lora_single_device 配置，该配置的默认批量大小为 2。修改批量大小（或其他超参数，例如优化器）将影响峰值内存和最终评估结果。

LoRA 层	秩	Alpha	峰值内存	准确率 (truthfulqa_mc2)
仅 Q 和 V	8	16	15.57 GB	0.475
所有层	8	16	15.87 GB	0.508
仅 Q 和 V	64	128	15.86 GB	0.504
所有层	64	128	17.04 GB	0.514

我们可以看到，我们的基线设置给出了最低的峰值内存，但我们的评估性能相对较低。通过为所有线性层启用 LoRA 并将秩增加到 64，我们看到在此任务上的准确率绝对提高了近 4%，但我们的峰值内存也增加了约 1.4GB。这些只是一些简单的实验；我们鼓励您运行自己的微调，以找到适合您特定设置的权衡。

此外，如果您想进一步降低模型的峰值内存（并且仍然可能获得相似的模型质量结果），您可以查看我们的 QLoRA 教程。