使用 QLoRA 微调 Llama2

在本教程中，我们将了解 QLoRA，它是 LoRA 的增强版，它以 4 位量化精度保持冻结的模型参数，从而减少内存使用量。我们将逐步介绍如何使用 QLoRA 在 torchtune 中使用不到 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 及其底层抽象的深入介绍，请参阅本教程的 QLoRA 在 torchtune 中的深入介绍 部分。

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

要以 QLoRA 风格量化 LoRALinear 层，只需将 quantize_base 标志作为 True 传递到 LoRALinear 中。此标志将导致基础模型权重被量化为 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"])

在幕后，这将把 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 的 lora_llama_2_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 次迭代记录峰值内存统计信息。让我们了解 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 并设置一个帐户。有关在 torchtune 中使用 W&B 的更多详细信息，请参阅我们的“记录到 Weights & Biases”配方。

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

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

深入探究：从 LoRA 构建 QLoRA

本深入探究部分从本教程的 使用 QLoRA 节省内存 部分开始，并深入探讨如何使用 NF4Tensor 进行量化，并在正向传递中适当地处理。

首先，我们将从一个普通的最小 LoRA 层开始，该层取自 LoRA 教程 并进行了增强以支持量化

import torch
from torch import nn
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: torch.Tensor) -> torch.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。