使用 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 位 NormatFloat 类型，以及对量化参数本身进行双重量化的方式，以节省更多内存。torchtune 使用来自 torchao 库 的 NF4Tensor 抽象概念来构建论文中指定的 QLoRA 组件。torchao 是一个 PyTorch 原生库，允许您量化和修剪模型。

使用 QLoRA 节省内存

在本节中，我们将概述如何在 torchtune 中将 QLoRA 应用于 LoRALinear 层。有关 torchtune 中 QLoRA 和底层抽象概念的深入细节，请参阅本教程的 torchtune 中的 QLoRA 深入探讨 部分。

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

要以 QLoRA 风格量化 torchtune 中的 LoRALinear 层，只需将 quantize_base 标志传递到 LoRALinear 中，并将其设置为 True。此标志将导致基础模型权重被量化，并由 NF4Tensor 数据类型支持。正向传递也将自动处理，以与 NF4Tensor 数据类型协同工作，具体来说，NF4 基础权重将被反量化为计算精度，将计算激活，并且在反向传递中只会存储 4 位参数以进行梯度计算，从而避免了存储更高精度计算 dtype 所产生的额外内存使用量。

以下是以量化 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"])

在后台，这将向所有注意力层的 q_proj 和 v_proj 矩阵应用 LoRA，并将这些矩阵中的基础参数进一步量化为 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。这种卸载是为了避免内存峰值；如果没有，我们将不得不保留 state_dict 在 GPU 上的完整 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 次迭代记录内存峰值统计信息。让我们了解一下在 LoRA 训练之上启用 QLoRA 的内存节省。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 的 nightly 版本。有关将 PyTorch 更新到最新的 nightly 版本，请参阅 安装说明。更新后，您可以通过配置覆盖指定编译标志为 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 表示。有关更多详细信息，请参阅 实现。 linear_nf4 在使用量化的基础模型权重运行时处理正向传递和自动微分。它将正向传递计算为使用传入激活和未量化权重的常规 F.linear。量化权重保存在反向传递中，而不是权重的未量化版本，以避免由于在反向传递中存储更高精度变量来计算梯度而导致的额外内存使用。有关更多详细信息，请参阅 linear_nf4。