使用LoRA微调Llama2

本指南将介绍 LoRA，这是一种参数高效的微调技术， 并展示如何使用 torchtune 对 Llama2 模型进行 LoRA 微调。 如果您已经了解 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 将输入投影到一个更低的秩（在实践中通常为四或八）， 而 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的自注意力机制中，in_dim=out_dim=4096用于Q、K和V的投影。这意味着一个秩为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配方（详情请见此处），你只需要传递相关的检查点路径。加载模型权重和设置可训练参数将在配方中处理。

LoRA 微调配方在 torchtune

最后，我们可以将所有内容整合在一起，并使用 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 会应用于 Q 和 V 投影，秩为 8。 一些关于 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 的 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层的秩。 （我们将并行缩放alpha与LoRA秩，如上所述。）

为了比较实验结果，我们可以在 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 教程。