作者: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 Notebook。

引言

大语言模型(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,这使得微调变得具有挑战性,并且对所有人来说都不那么容易获得。为了弥合这一差距,参数高效微调(PEFT)方法如今被社区广泛采用。

参数高效微调 (PEFT) 方法

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

它们可以通过其概念框架进行区分:该方法是微调现有参数的子集、引入新参数、引入可训练提示等?我们建议读者阅读下面分享的论文,其中详细比较了现有的 PEFT 方法。

Venn diagram

图片摘自论文:缩减规模以扩大规模:参数高效微调指南

对于这篇博文,我们将重点介绍用于大语言模型的低秩适应(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 矩阵输出的隐藏状态在同一 dtype。

压缩和量化大语言模型最近成为了一个热门话题,因为 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 的新硬件,应将其设置为最优的 BFloat16,以获得最佳性能

总而言之,通过结合这些量化过程的改进和慷慨使用 LoRA,我们将模型压缩了 90% 以上,并且在没有通常的量化性能下降的情况下保持了完整的模型性能,同时通过每层的 16 位 LoRA 适配器保留了完整的微调能力。

QLoRA 的实践应用

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

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

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

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

每个可训练参数总计 14 字节乘以 0.0029(因为 QLoRA 最终只有 0.29% 的可训练参数),这使得 QLoRA 训练设置成本约为 4.5GB 以适应,但在实践中需要约 7-10GB 以包含始终为半精度的中间隐藏状态(在下一节分享的 Google Colab 演示中,序列长度为 512 时为 7 GB,序列长度为 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. “训练语言模型遵循带有人类反馈的指令。” arXiv 预印本 arXiv:2203.02155 (2022)。

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

  • 数据打包 (Packing):不是批量中每个样本一个文本,然后填充到最长文本或模型的最大上下文长度,而是将大量文本用句子结束 (EOS) 符号连接起来,并截取上下文大小的块来填充批次,而无需任何填充。这种方法显著提高了训练效率,因为模型处理的每个 token 都为训练做出了贡献。

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 Notebook,您可以通过此链接查看。我们使用了上面各节分享的所有组件,并在 UltraChat 数据集上使用 QLoRA 微调了 llama-7b 模型。如下截图所示,当使用序列长度 1024 和批量大小 4 时,内存使用量非常低(约 10GB)。

Memory usage diagram