使用 QLoRA 微调 Llama2

在本教程中，我们将了解 QLoRA，这是 LoRA 的一项增强功能，它以 4 位量化精度保持冻结模型参数，从而减少内存使用量。我们将逐步了解如何在 torchtune 中利用 QLoRA 以不到 10 GB 的内存微调 Llama2-7b 模型。强烈建议首先了解 torchtune 中的 LoRA 微调。

您将学到什么

• QLoRA 如何在 LoRA 微调中节省内存

• torchtune 中 QLoRA 的概述

• 如何在 torchtune 中运行 QLoRA 微调

先决条件

• 熟悉 torchtune

• 确保 安装了 torchtune

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

• 熟悉 torchtune 中的 LoRA

什么是 QLoRA？

QLoRA 以 LoRA 为基础，可进一步节省内存。在 LoRA 中，模型参数可以被认为存在于两个分区中：适配器（添加到神经网络不同层的低秩矩阵）和基本模型参数（是原始模型一部分的参数）。在香草 LoRA 风格的训练中，这两个参数都以相同的精度（通常为 fp32 或 bf16）保存，因此计算的激活和中间梯度为 fp32/bf16。

QLoRA 进一步将基本模型参数量化为定制的 4 位 NormalFloat (NF4) 数据类型，从而将参数内存使用量减少 4-8 倍，同时在很大程度上保留模型精度。因此，绝大多数参数仅占用 4 位（与 bf16/fp32 数据类型的 16 位或 32 位相反）。此量化通过原始 QLoRA 论文 中突出显示的方法完成。适配器参数仍以原始精度保存，激活、梯度和优化器状态仍以更高的精度存在，以保持精度。

QLoRA 作者引入了两个关键抽象来减少内存使用并避免精度下降：定制的 4 位 NormalFloat 类型，以及将量化参数本身量化的双重量化方法，以节省更多内存。torchtune 使用 NF4Tensor 抽象，该抽象来自 torchao 库，以根据论文构建 QLoRA 组件。torchao 是一个 PyTorch 原生库，允许你量化和剪枝模型。

使用 QLoRA 节省内存

在本节中，我们将概述如何将 QLoRA 应用于 torchtune 中的 LoRALinear 层。如需深入了解 torchtune 中的 QLoRA 和底层抽象，请参阅本教程的 torchtune 中的 QLoRA 深入解析 部分。

QLoRA 的一个核心思想是计算和存储数据类型 (dtypes) 之间的区别。具体来说，QLoRA 以 4 位精度（即存储数据类型）存储基础模型参数，并在原始的更高精度（计算数据类型）中运行计算，通常是 fp32 或 bf16。作为第一步，QLoRA 需要将这些基础模型参数量化为 4 位精度并存储它们。

要以 QLoRA 样式量化 LoRALinear 层，只需将 quantize_base 标志作为 True 传递到 LoRALinear 中。此标志将导致基础模型权重被量化并由 NF4Tensor 数据类型支持。前向传递也将自动处理以使用 NF4Tensor 数据类型，具体来说，NF4 基础权重将被反量化为计算精度，激活将被计算，并且只有 4 位参数将被存储用于反向传递中的梯度计算，避免了存储更高精度计算数据类型会产生的额外内存使用。

下面是一个创建量化 LoRALinear 层的示例，与未量化的 LoRALinear 层进行比较。正如我们所看到的，量化层消耗的内存比未量化层少约 8 倍。

import torch
from torchtune.modules.peft import LoRALinear

torch.set_default_device("cuda")
qlora_linear = LoRALinear(512, 512, rank=8, alpha=0.1, quantize_base=True)
print(torch.cuda.memory_allocated())  # 177,152 bytes
del qlora_linear
torch.cuda.empty_cache()
lora_linear = LoRALinear(512, 512, rank=8, alpha=0.1, quantize_base=False)
print(torch.cuda.memory_allocated()) # 1,081,344 bytes

在 torchtune 中使用 QLoRA

我们现在将介绍如何初始化支持 QLoRA 的 Llama2-7b 模型以及有关 QLoRA 检查点的详细信息。

使用 torchtune，您可以使用类似于 LoRA 构建器（lora_llama_2_7b）的简单构建器将 QLoRA 应用于 Llama2 模型。下面是一个使用 QLoRA 初始化 Llama2-7b 模型的简单示例

from torchtune.models.llama2 import qlora_llama2_7b

qlora_model = qlora_llama2_7b(lora_attn_modules=["q_proj", "v_proj"])

在底层，这将把 LoRA 应用于所有注意力层中的 q_proj 和 v_proj 矩阵，并进一步将这些矩阵中的基础参数量化为 NF4 数据类型。请注意，基础模型参数的量化仅应用于配置为添加 LoRA 适配器的层。例如，在这种情况下，注意力层中的 k_proj 和 output_proj 没有应用 LoRA，因此它们的底层模型参数没有被量化。我们可以通过打印特定注意力层的底层模型参数数据类型来看到这一点

attn = qlora_model.layers[0].attn
print(type(attn.q_proj.weight))  # <class 'torchao.dtypes.nf4tensor.NF4Tensor'>
print(type(attn.k_proj.weight))  # <class 'torch.nn.parameter.Parameter'>

接下来，有一些细节对于 QLoRA 启用的模型的检查点（即 state_dict）至关重要。为了与 torchtune 的 检查点 很好的集成，我们需要将 NF4Tensors 转换回它们的原始精度（通常为 fp32/bf16）。这允许 QLoRA 训练的检查点与生态系统的其余部分（在 torchtune 及其他地方，例如训练后量化、评估、推理）很好地互操作。此转换过程还允许将 LoRA 适配器权重合并回基础模型，就像在典型的 LoRA 训练流程中所做的那样。

为实现此目的，当使用 torchtune 的 qlora_llama2_7b 构建器时，我们会自动注册一个钩子 reparametrize_as_dtype_state_dict_post_hook，它在对顶层模型调用 .state_dict() 之后运行。此钩子将 NF4Tensors 转换回它们的原始精度，同时也将这些转换后的张量卸载到 CPU。此卸载是为了避免内存峰值；如果不这样做，我们必须在 GPU 上保留 state_dict 的整个 bf16/fp32 副本。

将所有内容放在一起：QLoRA 微调

将所有内容放在一起，我们现在可以使用 torchtune 的 LoRA 配方 和 QLoRA 配置 来微调模型。

请确保你已按照 这些说明 下载 Llama2 权重和标记器。然后，你可以运行以下命令，在单个 GPU 上对 Llama2-7B 执行 QLoRA 微调。

tune run lora_finetune_single_device --config llama2/7B_qlora_single_device

注意

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

默认情况下，此运行应在模型初始化时和训练期间每 100 次迭代记录峰值内存统计信息。让我们了解 QLoRA 在 LoRA 训练之上实现的内存节省。LoRA 训练可以按如下方式运行

tune run lora_finetune_single_device --config llama2/7B_lora_single_device

您应该在模型初始化和训练期间看到打印的内存使用情况。以下是 LoRA 模型初始化的示例日志

Memory Stats after model init::
GPU peak memory allocation: 13.96 GB
GPU peak memory reserved: 13.98 GB
GPU peak memory active: 13.96 GB

下表比较了 QLoRA 在模型初始化和训练期间保留的内存与香草 LoRA 的内存。我们可以看到，在模型初始化期间，QLoRA 将峰值内存减少了约 35%，在模型训练期间减少了约 40%

微调方法	保留的峰值内存，模型初始化	保留的峰值内存，训练
LoRA	13.98 GB	15.57 GB
QLoRA	9.13 GB	9.29 GB

从日志中可以看到，开箱即用的训练性能相当慢，每秒慢于 1 次迭代

1|149|Loss: 0.9157477021217346:   1%|          | 149/25880 [02:08<6:14:19,  1.15it/s

为了加快速度，我们可以利用 torch.compile 编译我们的模型并运行编译结果。为了配合 QLoRA 训练，必须使用 PyTorch 的夜间构建。要将 PyTorch 更新到最新的夜间版本，请参阅 安装说明。更新后，您可以通过配置覆盖将编译标志指定为 True

tune run lora_finetune_single_device --config llama2/7B_qlora_single_device compile=True

从日志中，我们可以看到速度提高了约 200%（在训练稳定后经过数百次迭代）

1|228|Loss: 0.8158286809921265:   1%|          | 228/25880 [11:59<1:48:16,  3.95it/s

下面可以看到 QLoRA 和 LoRA 之间平滑损失曲线的比较。

注意

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

作为一项练习，您还可以尝试运行一些评估任务或手动检查已保存检查点输出的生成（可以在 output_dir 中找到）。

在最后一节中，我们将深入探讨如何从 LoRA 组件构建 QLoRA 组件。

深入探讨：从 LoRA 构建 QLoRA

此深入部分从本教程的 使用 QLoRA 节省内存 部分继续，并深入探讨如何使用 NF4Tensor 执行量化并在前向传递中适当地处理。

首先，我们将从 LoRA 教程 中获取一个香草极简 LoRA 层，并对其进行增强以支持量化

from torch import nn, Tensor
import torch.nn.functional as F
from torchao.dtypes.nf4tensor import linear_nf4, to_nf4

class LoRALinear(nn.Module):
  def __init__(
    self,
    in_dim: int,
    out_dim: int,
    rank: int,
    alpha: float,
    dropout: float,
    quantize_base: bool
  ):
    # These are the weights from the original pretrained model
    self.linear = nn.Linear(in_dim, out_dim, bias=False)
    self.linear_weight = self.linear.weight
    # Use torchao's to_nf4 API to quantize the base weight if needed.
    if quantize_base:
      self.linear_weight = to_nf4(self.linear_weight)
    # 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:
    # frozen_out would be the output of the original model
    if quantize_base:
      # Call into torchao's linear_nf4 to run linear forward pass w/quantized weight.
      frozen_out  = linear_nf4(x, self.weight)
    else:
      frozen_out = F.linear(x, self.weight)

    # 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

如上所述，torchtune 依赖于 torchao 来获取 QLoRA 所需的一些核心组件。其中包括 NF4Tensor，以及包括 to_nf4 和 linear_nf4 在内的有用实用程序。

在 LoRA 层之上的主要更改是使用 to_nf4 和 linear_nf4 API。

to_nf4 接受未量化的（bf16 或 fp32）张量，并生成权重的 NF4 表示。有关更多详细信息，请参阅 to_nf4 的 实现。当使用量化的基础模型权重运行时，linear_nf4 处理前向传递和自动微分。它使用传入的激活和未量化的权重计算前向传递作为常规 F.linear。量化权重保存在反向传播中，而不是未量化版本的权重，以避免由于在反向传播中存储更高精度的变量来计算梯度而导致的额外内存使用。有关更多详细信息，请参阅 linear_nf4。