使用 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 等人）假设在 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% 以上。

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

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_state_dict_for_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 和等级一起缩放是一个好习惯）。

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 教程。