使用QLoRA微调Llama2

在本教程中，我们将学习 QLoRA，这是对 LoRA 的改进，在保持模型参数冻结的同时，以 4 位量化精度存储，从而减少内存使用。我们将介绍如何在 torchtune 中利用 QLoRA 在不到 10 GB 的内存下微调 Llama2-7b 模型。强烈建议您首先了解 torchtune 中的 LoRA 微调。

你将学到什么

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

• QLoRA 在 torchtune 中的概述

• 如何在 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 使用了 NF4Tensor 抽象，该抽象来自 torchao 库，用于构建论文中指定的 QLoRA 组件。torchao 是一个 PyTorch 原生库，允许您对模型进行量化和剪枝。

使用QLoRA节省内存

在本节中，我们将概述如何在torchtune中对一个LoRALinear层应用QLoRA。有关torchtune中QLoRA的详细信息以及底层抽象的深入探讨， 请参阅本教程中的torchtune中的QLoRA深入探讨部分。

QLoRA 的一个核心思想是区分计算数据类型（dtypes）与存储数据类型。具体而言，QLoRA 以 4 位精度（即存储 dtype）存储基础模型参数，并以原始更高精度（即计算 dtype）执行计算，通常采用 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

使用 QLoRA 在 torchtune

我们现在将介绍如何初始化支持 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'>

接下来，有一些细节对于保存检查点（即 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 上。这种卸载是为了避免内存峰值；如果我们不这样做，就必须在 GPU 上保留整个 bf16/fp32 格式的 state_dict 副本。

将所有内容整合在一起：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 的 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 并设置一个账户。有关在 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。