使用QLoRA微调Llama2

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

What you will learn

• QLoRA 如何比 LoRA 微调节省内存

• torchtune 中 QLoRA 的概述

• 如何在torchtune中运行QLoRA微调

Prerequisites

• 熟悉 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的一个核心思想是计算和存储数据类型（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

在 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的lora_llama_2_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次迭代记录峰值内存统计信息。让我们了解在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|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 表示的权重。有关更多详细信息，请参阅 implementation 的 to_nf4。 linear_nf4 在使用量化基础模型权重运行时处理前向传播和自动梯度。它像常规的 F.linear 一样计算前向传播，使用传入的激活和未量化的权重。为了节省内存，量化后的权重被保存用于反向传播，而不是未量化的权重版本，以避免在反向传播中存储更高精度的变量来计算梯度。有关更多详细信息，请参阅 linear_nf4。