由尤尼斯·贝尔卡达、马克·孙、蒂图斯·冯·科勒、索拉布·曼格鲁卡尔、本杰明·博桑、利桑德·德布特、史蒂文·刘共同撰写

我们展示了如何使用 LoRA 和 PyTorch 以及 Hugging Face 生态系统中的工具,在一个典型的消费级 GPU(NVIDIA T4 16GB)上微调一个 7B 参数模型,并使用完整的可复现 Google Colab 笔记本。

引言

大型语言模型(LLMs)在工业应用中表现出令人印象深刻的性能。通常,开发者寻求针对特定用例和应用程序定制这些LLMs,以优化其性能。然而,LLMs的设计本身就是大型的,需要大量的 GPU 来进行微调。

让我们通过尝试在一个免费层级的 Google Colab 实例(1x NVIDIA T4 16GB)上微调 Llama 模型来关注一个具体的例子。Llama-2 7B 有 70 亿个参数,如果以全精度加载,总大小为 28GB。鉴于我们的 GPU 内存限制(16GB),模型甚至无法加载,更不用说在 GPU 上训练了。这种内存需求可以通过降低一半来满足,性能下降可以忽略不计。您可以在这里了解更多关于在半精度和混合精度下运行模型进行训练的信息。

什么使得我们的 Llama 微调如此昂贵?

在使用 Adam 优化器进行全量微调,并采用半精度模型和混合精度模式的情况下,我们需要为每个参数分配:

  • 2 个字节的权重
  • 2 个字节的梯度
  • Adam 优化器状态占用 4 + 8 字节

每个可训练参数占用 16 字节,总计 112GB(不包括中间隐藏状态)。鉴于目前最大的 GPU 可拥有高达 80GB 的 GPU VRAM,这使得微调变得具有挑战性,并且对每个人来说都难以接触。为了弥合这一差距,社区广泛采用了参数高效微调(PEFT)方法。

参数高效微调(PEFT)方法

PEFT 方法旨在大幅减少模型的可训练参数数量,同时保持与完整微调相同的性能。

可以通过其概念框架来区分:该方法是否微调现有参数的子集,引入新参数,引入可训练提示等?我们建议读者阅读以下论文,该论文广泛比较了现有的 PEFT 方法。

Venn diagram

图片来自论文:《缩小规模以扩大规模:参数高效微调指南》

对于这篇博客文章,我们将重点关注低秩适配(LoRA)大型语言模型,因为它是社区中最广泛采用的 PEFT 方法之一。

使用🤗 PEFT 进行大型语言模型的低秩适配(LoRA)

胡等人从微软团队提出的 LoRA 方法于 2021 年发布,其工作原理是在一个基础模型(我们将表示为 base model)中附加额外的可训练参数。

为了使微调更高效,LoRA 将一个大的权重矩阵分解为两个较小的低秩矩阵(称为更新矩阵)。这些新矩阵可以训练以适应新的数据,同时保持整体变化数量较低。原始权重矩阵保持冻结,不接受任何进一步的调整。为了产生最终结果,原始权重和适应后的权重被结合。

这种方法具有几个优点:

  • LoRA 通过大幅减少可训练参数的数量,使微调更加高效。
  • 原始预训练权重保持冻结状态,这意味着您可以在其基础上构建多个轻量级且便携的 LoRA 模型,用于各种下游任务。
  • LoRA 与其他许多参数高效方法正交,可以与其中许多方法结合使用。
  • 使用 LoRA 微调的模型性能与全量微调模型性能相当。
  • 当适配器权重与基础模型合并时,LoRA 不会增加任何推理延迟。

原则上,LoRA 可以应用于神经网络中权重矩阵的任何子集以减少可训练参数的数量。然而,为了简化并进一步提高参数效率,在 Transformer 模型中 LoRA 通常仅应用于注意力块。LoRA 模型中的可训练参数数量取决于低秩更新矩阵的大小,这主要取决于秩 r 和原始权重矩阵的形状。

Animated diagram that show how LoRA works in practice

动画图展示了 LoRA 在实际中的应用方式 - 原始内容适配自 LoRA 原始论文中的图 1

下面是一个代码片段,展示了如何使用 Hugging Face PEFT 库训练 LoRA 模型:

code snippet showing how to train LoRA model using  Hugging Face PEFT library

基础模型可以是任何 dtype 利用 SOTA LLM 量化,并以 4 位精度加载基础模型

根据 LoRA 公式,只要基础模型的隐藏状态与 LoRA 矩阵的输出隐藏状态属于同一数据类型(‘dtype’),就可以在任何数据类型(‘dtype’)中对基础模型进行压缩。

随着 SOTA 模型变得越来越大,压缩和量化大型语言模型最近已成为一个令人兴奋的话题,因为它们越来越难以服务于最终用户。社区中许多人提出了各种方法,以最小化性能下降来有效地压缩LLMs。

这就是 bitsandbytes 库的作用所在。它的目的是让 Tim Dettmers 的研究成果,这位在量化以及深度学习硬件加速器使用方面的领先学术专家,对公众更加易于接触。

QLoRA: bitsandbytes 对 AI 民主化的一大核心贡献

LLMs 的量化主要关注推理阶段的量化,但 QLoRA(量化模型权重 + 低秩适配器)论文展示了在大型模型规模中使用冻结量化权重进行反向传播的突破性实用价值。

使用 QLoRA,我们在所有规模和模型上实现了与 16 位微调相当的性能,同时将微调内存占用减少了超过 90%,从而使得在消费级硬件上对 SOTA 模型进行微调成为可能。

在这种方法中,LoRA 对于微调和纠正微小的残差量化误差都至关重要。由于量化模型的尺寸显著减小,因此可以在每个网络层大量放置低秩适配器,而这些适配器加起来仍然只占原始模型权重内存占用的 0.2%。通过这样的 LoRA 使用,我们实现了与 16 位全模型微调等效的性能。

System diagram

除了大量使用 LoRA,为了实现 4 位模型的超高保真微调,QLoRA 还使用了 3 个额外的算法技巧:

  1. 4 位正常浮点数(NF4)量化,一种利用模型权重正态分布特性的自定义数据类型,将等量的权重(每个块)分配到每个量化箱中,从而提高信息密度。
  2. 双重量化,量化量化常数(进一步节省空间)。
  3. 分页优化器,防止梯度检查点期间的内存峰值导致内存不足错误。

一个有趣的方面是在 GPU 缓存中对 4 位权重进行去量化,矩阵乘法以 16 位浮点运算进行。换句话说,我们使用低精度存储数据类型(在我们的情况下是 4 位,但原则上可以互换)和一个正常精度的计算数据类型。这很重要,因为后者默认为 32 位,以适应硬件兼容性和数值稳定性,但对于支持它的较新硬件,应将其设置为最优的 BFloat16 以实现最佳性能。

总结来说,通过结合对量化过程的优化和 LoRA 的广泛使用,我们使模型压缩超过 90%,同时保留了完整的模型性能,没有出现通常的量化降级,并且还保留了使用 16 位 LoRA 适配器在每一层进行微调的全部能力。

实际应用中的 QLoRA

这些 SOTA 量化方法打包在 bitsandbytes 库中,并方便地与 HuggingFace 🤗 Transformers 集成。例如,要使用LLM.int8 和 QLoRA 算法,只需将 load_in_8bitload_in_4bit 传递给 from_pretrained 方法。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "facebook/opt-125m"
# For LLM.int8()
# model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True)

# For QLoRA
model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True)

您可以在文档的此特定部分中了解更多有关量化功能的信息:https://huggingface.co/docs/transformers/main_classes/quantization

当使用 QLoRA 与 Adam 优化器结合 4 位基模型和混合精度模式时,我们需要为每个参数分配:

  • 0.5 字节用于权重
  • 2 字节用于梯度
  • 4 + 8 字节用于 Adam 优化器状态

每个可训练参数总共 14 字节,乘以 0.0029,因为最终我们只有 0.29%的可训练参数使用 QLoRA,这使得 QLoRA 的训练设置成本大约为 4.5GB,但在实际中需要包括中间隐藏状态,这些状态始终以半精度存储(对于序列长度为 512 时为 7GB,对于序列长度为 1024 时为 10GB),在下一节中分享的 Google Colab 演示中需要大约 7-10GB。

下面是使用 Hugging Face PEFT 训练 QLoRA 模型的代码片段:

code snippet showing how to train QLoRA model using Hugging Face PEFT

使用 TRL 进行LLM训练

ChatGPT、GPT-4 和 Claude 等模型是强大的语言模型,它们通过一种称为人类反馈强化学习(RLHF)的方法进行了微调,以更好地符合我们的期望和期望的使用方式。微调过程分为 3 个步骤:

  • 监督微调(SFT)
  • 奖励/偏好建模(RM)
  • 基于人类反馈的强化学习(RLHF)

Process diagram

来自 InstructGPT 论文:Ouyang, Long, et al. “通过人类反馈训练语言模型以遵循指令。” arXiv 预印本 arXiv:2203.02155(2022 年)。

这里,我们只关注监督微调步骤。我们使用与预训练类似的过程在新数据集上训练模型。目标是预测下一个标记(因果语言建模)。可以应用多种技术来提高训练效率:

  • 打包:在批次中每个样本只有一个文本并填充到最长的文本或模型的最大上下文之前,我们使用 End-Of-Sentence(EOS)标记将许多文本连接起来,并切割上下文大小的块来填充批次而不进行任何填充。这种方法显著提高了训练效率,因为每个由模型处理的标记都对训练有贡献。

Sample diagram

  • 仅在完成时进行训练:我们希望模型能够理解提示并生成答案。与其在(提示+答案)的整个输入上训练模型,不如仅对完成部分进行训练,这样训练会更高效。

您可以使用 SFTTrainer 使用这些技术进行监督微调:

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=training_arguments,
    train_dataset=train_dataset,
    dataset_text_field="text",
    max_seq_length=1024,
    packing=True,
)

由于 SFTTrainer 后端由🤗 accelerate 提供支持,您可以通过一行代码轻松地将训练适应您的硬件配置!

例如,如果您有 2 个 GPU,您可以使用以下命令进行分布式数据并行训练:

accelerate launch --num_processes=2 training_llama_script.py

将所有部件组合在一起

我们制作了一个完整的可复现的 Google Colab 笔记本,您可以通过此链接进行检查。我们使用上述部分中共享的所有组件,并在 UltraChat 数据集上使用 QLoRA 微调了 llama-7b 模型。如图下所示,当使用序列长度为 1024 和批大小为 4 时,内存使用率保持非常低(大约 10GB)。

Memory usage diagram