作者:Younes Belkada, Marc Sun, Titus von Köller, Sourab Mangrulkar, Benjamin Bossan, Lysandre Debut, Steven Liu

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

简介

大型语言模型 (LLM) 在工业应用中展现出令人印象深刻的能力。通常,开发者寻求针对特定用例和应用定制这些 LLM,以对其进行微调,从而获得更好的性能。然而,LLM 在设计上是大型的,需要大量的 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

图片来自论文:Scaling Down to Scale Up: A Guide to Parameter-Efficient Fine-Tuning

在本博客文章中,我们将重点介绍大型语言模型的低秩自适应 (LoRA),因为它是社区中最常用的 PEFT 方法之一。

使用 🤗 PEFT 的大型语言模型低秩自适应 (LoRA)

Microsoft 团队的 Hu 等人在 2021 年提出了 LoRA 方法,其工作原理是将额外的可训练参数附加到模型中(我们将其表示为基础模型)。

为了提高微调效率,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 公式,基础模型可以以任何数据类型 (‘dtype’) 压缩,只要来自基础模型的隐藏状态与来自 LoRA 矩阵的输出隐藏状态的数据类型相同即可。

压缩和量化大型语言模型最近已成为一个令人兴奋的话题,因为 SOTA 模型变得越来越大,并且更难以服务和供最终用户使用。社区中的许多人提出了各种有效压缩 LLM 的方法,同时将性能降级降至最低。

这就是 bitsandbytes 库的用武之地。其目的是使 Tim Dettmers(一位量化和深度学习硬件加速器使用方面的领先学术专家)的尖端研究能够为公众所用。

QLoRA:bitsandbytes 对 AI 大众化的核心贡献之一

LLM 的量化主要侧重于推理量化,但 QLoRA(量化模型权重 + 低秩适配器)论文展示了通过冻结的量化权重进行反向传播在大型模型尺度上的突破性效用。

借助 QLoRA,我们在所有尺度和模型上都匹配了 16 位微调性能,同时将微调内存占用减少了 90% 以上,从而允许在消费级硬件上微调 SOTA 模型。

在这种方法中,LoRA 对于微调校正最小的残余量化误差都至关重要。由于量化模型的大小显着减小,因此可以在每个网络层慷慨地放置低秩适配器,这些适配器加在一起仍然只占原始模型权重内存占用的 0.2%。通过如此使用 LoRA,我们实现了已证明与 16 位全模型微调相当的性能。

System diagram

除了慷慨地使用 LoRA 之外,为了实现 4 位模型的高保真微调,QLoRA 还使用了 3 个进一步的算法技巧:

  1. 4 位 NormalFloat (NF4) 量化,一种自定义数据类型,利用模型权重正态分布的属性,并将相等数量的权重(每个块)分配给每个量化 bin,从而提高信息密度。
  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://hugging-face.cn/docs/transformers/main_classes/quantization

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

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

由于我们最终只使用 QLoRA 获得了 0.29% 的可训练参数,因此每个可训练参数总共 14 字节乘以 0.0029,这使得 QLoRA 训练设置成本约为 4.5GB 以容纳,但实际上需要约 7-10GB 以包含中间隐藏状态,这些状态始终是半精度的(在下一节中共享的 Google Colab 演示中,序列长度为 512 时为 7GB,序列长度为 1024 时为 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. “Training language models to follow instructions with human feedback.” arXiv preprint arXiv:2203.02155 (2022).

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

  • 打包:我们不采用批次中每个样本一个文本,然后填充到最长文本或模型的最大上下文的方式,而是将大量文本与句子结尾 (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 笔记本,您可以通过此链接查看。我们使用了上面各节中共享的所有组件,并使用 QLoRA 在 UltraChat 数据集上微调了一个 llama-7b 模型。正如从下面的屏幕截图中可以观察到的那样,当使用序列长度为 1024 和批次大小为 4 时,内存使用率保持非常低(约为 10GB)。

Memory usage diagram