使用 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 中，可以将模型参数视为存在于两个分区中：适配器，它们是 添加到神经网络不同层的低秩矩阵，以及 Base Model 参数，这些参数是 原始模型。在普通的 LoRA 式训练中，这两个参数都以相同的精度保持（通常为 fp32 或 bf16），并且 因此，计算的活化和中间梯度在 FP32/BF16 中。

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

QLoRA 作者介绍了两个关键抽象来减少内存使用并避免准确性下降：定制的 4 位 NormatFloat type 和双量化方法，该方法对量化参数本身进行量化以节省更多内存。Torchtune 使用 从 torchao 库中抽象出 NF4Tensor 来构建论文中指定的 QLoRA 组件。 torchao 是一个 PyTorch 原生库，允许您量化和修剪模型。

使用 QLoRA 节省内存

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

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

要量化LoRALinear层，只需将标志传入 intoquantize_baseTrueLoRALinear.此标志 将导致 Base Model 权重被量化并由 dtype 提供支持。前向传递也将被自动处理以使用 dtype， 具体来说，基本权重将被去量化为计算精度，将计算激活，并且仅存储 4 位参数以进行梯度计算 在向后传递中，避免了存储更高精度的计算 DTtype 所产生的额外内存使用。NF4TensorNF4TensorNF4

以下是创建量化层与未量化层的比较示例。正如我们所看到的，量化层消耗 内存比未量化的对应物少 ~8 倍。LoRALinearLoRALinear

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 builder （） 的简单构建器将 QLoRA 应用于 Llama2 模型。下面是一个简单的示例 在启用 QLoRA 的情况下初始化 Llama2-7b 模型：lora_llama_2_7b

from torchtune.models.llama2 import qlora_llama2_7b

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

在后台，这会将 LoRA 应用于所有注意力层中的 和 矩阵，并进一步量化基本参数 在这些矩阵中设置为 dtype.请注意，基础模型参数的量化仅适用于配置为具有 添加了 LoRA 适配器。例如，在这种情况下，注意力层没有应用 LoRA，因此它们的 基本模型参数未量化。我们可以通过打印特定注意力层的基本模型参数 dtypes 来看到这一点：q_projv_projNF4k_projoutput_proj

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 的模型进行检查点（即）来说，有几个细节是必不可少的。 为了与 torchtune 的 checkpointing 很好地集成，我们需要转换回它们的 原始精度（通常为 FP32/BF16）。这使得 QLoRA 训练的检查点能够与生态系统的其余部分很好地互作，在 torchtune 等（例如，训练后量化、评估、推理）。此转换过程还允许将 LoRA 适配器权重合并回基本模型 在典型的 LoRA 训练流程中。state_dictNF4Tensors

为了实现这一点，当使用 torchtune 的构建器时，我们会自动注册一个钩子 ， 在调用 top level model 之后运行。这个钩子会转换回它们原来的精度，同时也会卸载这些 将张量转换为 CPU。这种卸载是为了避免内存达到峰值;如果我们不这样做，我们将不得不在 GPU 上维护 的整个 bf16/fp32 副本。qlora_llama2_7breparametrize_as_dtype_state_dict_post_hook.state_dict()NF4Tensorsstate_dict

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

综上所述，我们现在可以使用 torchtune 的 LoRA 配方对模型进行微调。 使用 QLoRA 配置。

确保您首先按照这些说明下载了 Llama2 权重和分词器。 然后，您可以运行以下命令，在单个 GPU 上执行 Llama2-7B 的 QLoRA 微调。

tune run lora_finetune_single_device --config llama2/7B_qlora_single_device

注意

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

默认情况下，此运行应在模型初始化时和每 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吉字节	15.57 吉字节
QLoRA	9.13 吉字节	9.29 吉字节

从日志中可以看出，开箱即用的训练性能相当慢，每次迭代 1 次还慢 第二：

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

为了加快速度，我们可以利用它来编译我们的模型并运行编译后的结果。使用 QLoRA 训练，必须使用 PyTorch 的 nightly 版本。要将 PyTorch 更新到最新的 nightly， 请参阅安装说明。更新后， 您可以通过 config override 指定 compile 标志：torch.compileTrue

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 对 QLoRA 所需的一些核心组件依赖于 torchao。这包括 以及有用的实用程序，包括 和 。NF4Tensorto_nf4linear_nf4

LoRA 层之上的关键变化是 和 API 的使用。to_nf4linear_nf4

to_nf4接受未量化的（BF16 或 FP32）张量，并生成权重的表示形式。有关更多详细信息，请参阅 implementation of 。 在使用量化的基本模型权重运行时处理 forward pass 和 autograd。它将 forward pass 作为 incoming active 和 unquantized 权重的 regular 进行计算。量化的权重被保存为反向，而不是未量化的权重版本，以避免额外的 由于存储更高精度的变量来计算向后传递中的梯度而导致的内存使用。有关详细信息，请参阅 linear_nf4。NF4to_nf4linear_nf4F.linear