使用 LoRA 微调 Llama2

本指南将向您介绍 LoRA，这是一种参数高效的微调技术。 并向您展示如何使用 torchtune 通过 LoRA 微调 Llama2 模型。 如果您已经知道 LoRA 是什么，并希望直接开始运行 您自己的 LoRA finetune 中，您可以跳转到 torchtune 中的 LoRA finetuning 配方。

您将学到什么

• 什么是 LoRA 以及它如何在微调期间节省内存

• torchtune 中的 LoRA 组件概述

• 如何使用 torchtune 运行 LoRA 微调

• 如何试验不同的 LoRA 配置

先决条件

• 熟悉 torchtune

• 确保安装 torchtune

• 确保您已下载 Llama2-7B 模型权重

什么是 LoRA？

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

注意

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

通过使用 LoRA 进行微调（而不是微调所有模型参数）， 由于 具有梯度的参数数。当使用具有 momentum 的优化器时， 像 AdamW 一样， 您可以预期从 Optimizer 状态看到进一步的内存节省。

注意

LoRA 内存节省主要来自梯度和优化器状态、 因此，如果模型的有效方法中具有峰值内存，则 LoRA 可能不会降低峰值内存。forward()

LoRA 的工作原理是什么？

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

下图给出了完整微调中单个权重更新步骤的简化表示 （左侧）与使用 LoRA 的权重更新步骤（右侧）进行比较。LoRA 矩阵 和 用作蓝色完整等级权重更新的近似值。AB

尽管 LoRA 在模型中引入了一些额外的参数，但只有 和 矩阵是可训练的。 这意味着，使用 rank LoRA 分解时，我们需要存储的梯度数量会减少 从 到 。（请记住，通常 比 和 小得多。forward()ABrin_dim*out_dimr*(in_dim+out_dim)rin_dimout_dim

例如，在 7B Llama2 的自注意力中，对于 Q、K、 和 V 投影。这意味着 rank 的 LoRA 分解将减少可训练的数量 给定投影的参数 from to ， a 减少 99% 以上。in_dim=out_dim=4096r=8

我们来看看原生 PyTorch 中 LoRA 的最小实现。

from torch import nn, Tensor

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: Tensor) -> 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

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

将 LoRA 应用于 Llama2 模型

使用 torchtune，我们可以轻松地将 LoRA 应用于具有各种不同配置的 Llama2。 我们来看看如何在有和没有 LoRA 的情况下在 torchtune 中构建 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)
CausalSelfAttention(
  (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)
CausalSelfAttention(
  (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_modelbase_model

为什么这很重要？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)

注意

每当使用 加载权重时，您都应该验证 中的任何缺失或额外的键 加载的符合预期。torchtune 的 LoRA 配方默认通过例如 .strict=Falsestate_dict

加载基本模型权重后，我们还希望仅将 LoRA 参数设置为 trainable。

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 配方（如此处所述），则只需将 relevant checkpoint 路径。加载模型权重和设置可训练参数将得到照顾 的。

torchtune 中的 LoRA 微调配方

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

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

注意

确保指向您的 Llama2 权重和分词器的位置。这是可以做到的 通过添加或直接修改文件。有关如何轻松克隆和修改 torchtune 配置的更多详细信息，请参阅我们的 All About Configs。checkpointer.checkpoint_files=[my_model_checkpoint_path] tokenizer_checkpoint=my_tokenizer_checkpoint_path7B_lora.yaml

注意

您可以根据 （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 个步骤的基线之间的（平滑）损失曲线的比较。

注意

以上数字是使用 W&B 生成的。您可以使用 torchtune 生成类似的损失曲线，但您需要安装 W&B 并单独设置一个帐户。

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

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

tune run lora_finetune_single_device --config llama2/7B_lora_single_device

在单个设备上，我们可能需要更多地了解我们的峰值内存。让我们运行一些实验 以查看 Finetune 期间的峰值内存。我们将沿着两个轴进行实验： 首先，应用了 LoRA 的模型层，其次，每个 LoRA 层的排名。（我们将扩大规模 alpha 与 LoRA 等级平行，如上所述。

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

以前，我们只为每个自我注意模块中的线性层启用了 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。修改 batch size（或其他超参数，例如优化器）将影响峰值内存 以及最终评估结果。

LoRA 图层	排	阿尔法	峰值内存	精度 （truthfulqa_mc2）
仅 Q 和 V	8	16	15.57 吉字节	0.475
所有图层	8	16	15.87 吉字节	0.508
仅 Q 和 V	64	128	15.86 吉字节	0.504
所有图层	64	128	17.04 吉字节	0.514

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

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