跳到正文
原文
Hugging Face Blog·· 2023-03-09精选AI 评分71

trl 与 peft 集成，可在 24GB 消费级 GPU 上用 RLHF 微调 20B 大模型

Fine-tuning 20B LLMs with RLHF on a 24GB consumer GPU

AI 导读

Hugging Face 官方发布 trl 与 peft 的集成，让使用 RLHF 的大语言模型微调更易上手。

推荐理由

原文给出 trl 与 peft 集成后在 24GB 消费级 GPU 上微调 20B 参数模型的完整路径，便于读者评估自身显存条件下的可行性。

正文 · AI 翻译

我们非常激动地正式发布 trl 与 peft 的集成，让任何人都能更轻松地使用强化学习对大型语言模型（LLM）进行微调！在这篇文章中，我们将解释为什么这是现有微调方法的一个有竞争力的替代方案。

请注意，peft 是一个通用工具，可以应用于许多机器学习用例，但它在 RLHF 中特别有趣，因为这种方法特别消耗内存！

如果你想直接深入研究代码，请直接查看 TRL 文档页面 上的示例脚本。

引言

LLM 与 RLHF

LLM 与 RLHF（基于人类反馈的强化学习）相结合，似乎是构建像 ChatGPT 这样非常强大的 AI 系统的下一个首选方法。

使用 RLHF 训练语言模型通常涉及以下三个步骤：

1- 在特定领域或指令与人类示范语料上微调预训练的 LLM

2- 收集人工标注数据集并训练奖励模型

3- 使用 RL（例如 PPO）结合奖励模型和该数据集，进一步微调步骤 1 中的 LLM

openai_diagram
ChatGPT 训练流程概览，从数据收集到 RL 部分。来源：OpenAI 的 ChatGPT 博客文章

基础 LLM 的选择在这里相当关键。在撰写本文时，可以“开箱即用”用于许多任务的“最佳”开源 LLM 是指令微调过的 LLM。值得注意的模型有：BLOOMZ、Flan-T5、Flan-UL2 和 OPT-IML。这些模型的缺点是它们的规模。要获得一个像样的模型，你至少需要处理 10B+ 规模的模型，仅为了在单个 GPU 设备上以全精度容纳模型而不进行任何训练，就需要高达 40GB 的 GPU 内存！

什么是 TRL？

trl 库旨在让 RL 步骤变得更加简单和灵活，以便任何人都可以在自定义数据集和训练设置上使用 RL 微调他们的 LM。在许多其他应用中，你可以使用该算法微调模型以生成积极的电影评论、进行可控生成或使模型毒性更低。

使用 trl，你可以以分布式方式或在单个设备上运行最流行的深度 RL 算法之一 PPO！我们利用 HF Mirror 生态系统中的 accelerate 来实现这一点，以便任何用户都可以将实验扩展到有趣的规模。

使用 RL 微调语言模型大致遵循下面详述的协议。这需要拥有原始模型的两份副本；为了避免活跃模型偏离其原始行为/分布太多，你需要在每个优化步骤计算参考模型的 logits。这为优化过程增加了一个硬性约束，因为每个 GPU 设备至少需要两份模型副本。如果模型规模增大，在单个 GPU 上容纳整个设置就变得越来越棘手。

trl_diagram
TRL 中 PPO 训练设置概览。

在 trl 中，你还可以使用参考模型和活跃模型之间的共享层来避免完整的副本。该功能的一个具体示例在去毒化示例中展示。

大规模训练

大规模训练可能会充满挑战。第一个挑战是将模型及其优化器状态装进可用的 GPU 设备。单个参数占用的 GPU 内存量取决于其“精度”（更具体地说是 dtype）。最常见的 dtype 是 float32（32 位）、float16 和 bfloat16（16 位）。最近，一些“奇特”精度已开箱即用地支持训练和推理（带有一定条件和约束），例如 int8（8 位）。简而言之，要在 GPU 设备上加载模型，每十亿参数在 float32 精度下需要 4GB，在 float16 下需要 2GB，在 int8 下需要 1GB。如果你想了解更多关于这个主题的内容，可以看看这篇更深入的博文：https://hf-mirror.com/blog/hf-bitsandbytes-integration。

如果使用 AdamW 优化器，每个参数需要 8 字节（例如，如果你的模型有 1B 参数，该模型的完整 AdamW 优化器将需要 8GB GPU 内存——来源）。

为了在规模化场景下应对这些挑战，人们采用了许多技术。最常见的范式是流水线并行、张量并行和数据并行。

model-parallelism
图片来源：这篇博文

在数据并行中，同一个模型被并行托管在多台机器上，每个实例接收不同的数据批次。这是最直接的并行策略，本质上复现了单 GPU 的情况，并且已经由 trl 支持。在流水线并行和张量并行中，模型本身被分布到多台机器上：在流水线并行中，模型按层切分，而张量并行则将张量操作（例如矩阵乘法）拆分到多个 GPU 上。使用这些模型并行策略时，你需要将模型权重分片到许多设备上，这要求你定义跨进程的激活值和梯度的通信协议。这实现起来并不简单，可能需要采用一些框架，例如 Megatron-DeepSpeed 或 Nemo。同样重要的是要强调其他对扩展 LLM 训练至关重要的工具，例如自适应激活检查点和融合内核。关于并行范式的进一步阅读可以在这里找到。

因此，我们问了自己以下问题：仅靠数据并行，我们能走多远？我们能否使用现有工具，将超大规模训练过程（包括活动模型、参考模型和优化器状态）放进单个设备中？答案似乎是肯定的。主要要素是：适配器和 8 位矩阵乘法！让我们在以下章节中介绍这些主题：

8 位矩阵乘法

高效 8 位矩阵乘法是一种最早在论文 LLM.int8() 中引入的方法，旨在解决量化大规模模型时的性能退化问题。所提出的方法将 Linear 层底层执行的矩阵乘法分解为两个阶段：将在 float16 中执行的部分为离群隐藏状态部分，以及在 int8 中执行的部分为“非离群”部分。

8bit-matmul
高效的 8 位矩阵乘法是一种首次在论文 LLM.int8() 中引入的方法，旨在解决量化大规模模型时的性能退化问题。所提出的方法将 Linear 层底层应用的矩阵乘法分解为两个阶段：以 float16 执行的离群隐藏状态部分，以及以 int8 执行的“非离群”部分。

简而言之，如果使用 8 位矩阵乘法，你可以将全精度模型的大小减少 4 倍（因此，对于半精度模型则减少 2 倍）。

低秩适配与 PEFT

2021 年，一篇名为 LoRA: Low-Rank Adaption of Large Language Models 的论文表明，可以通过冻结预训练权重，并为查询层和值层的注意力矩阵创建低秩版本来对大型语言模型进行微调。这些低秩矩阵的参数远少于原始模型，从而能够以少得多的 GPU 内存进行微调。作者证明，低秩适配器的微调取得了与微调完整预训练模型相当的结果。

lora-gif
输出激活的原始（冻结）预训练权重（左）由一个由权重矩阵 A 和 B 组成的低秩适配器（右）增强。

这项技术允许以一小部分内存需求对 LLM 进行微调。然而，也有一些缺点。由于适配器层中额外的矩阵乘法，前向和后向传递大约慢两倍。

什么是 PEFT？

Parameter-Efficient Fine-Tuning (PEFT) 是一个 HF Mirror 库，旨在支持在 LLM 上创建和微调适配器层。peft 与 🤗 Accelerate 无缝集成，可利用 DeepSpeed 和 Big Model Inference 处理大规模模型。

该库支持许多最先进的模型，并拥有大量示例，包括：

  • 因果语言建模
  • 条件生成
  • 图像分类
  • 8 位 int8 训练
  • Dreambooth 模型的低秩适配
  • 语义分割
  • 序列分类
  • 标记分类

该库仍处于广泛而活跃的开发中，未来几个月内将公布许多即将推出的功能。

使用低秩适配器微调 20B 参数模型

现在前置条件已经讲完，让我们逐步走过整个流程，并用图示解释如何在一张 24GB GPU 上使用上述工具，通过 RL 微调一个 20B 参数的 LLM！

第 1 步：以 8 位精度加载你的活动模型

step1
以 8 位精度加载模型相比全精度模型最多可节省 4 倍内存

使用 transformers 对 LLM 进行“免费午餐”式内存缩减的方法，是使用 LLM.int8 中描述的方法以 8 位精度加载模型。只需在调用 from_pretrained 方法时添加标志 load_in_8bit=True 即可实现（你可以在此处了解更多）。

如前一节所述，估算加载模型所需 GPU 内存的一个“窍门”是按“数十亿参数”来思考。由于一个字节需要 8 位，对于全精度模型（32 位 = 4 字节），每十亿参数需要 4GB；对于半精度模型，每十亿参数需要 2GB；对于 int8 模型，每十亿参数需要 1GB。

那么首先，我们只需以 8 位加载活动模型。让我们看看第二步需要做什么！

第 2 步：使用 peft 添加额外的可训练适配器

step2
你可以轻松地在冻结的 8 位模型上添加适配器，通过训练一小部分参数来降低优化器状态的内存需求

第二步是在模型内部加载适配器，并让这些适配器可训练。这能大幅减少活动模型所需训练权重的数量。此步骤利用 peft 库，只需几行代码即可完成。注意，一旦适配器训练完成，你可以轻松地将它们推送到 Hub 以供以后使用。

第 3 步：使用同一模型获得参考和活动的 logits

step3
你可以使用 peft API 轻松地禁用和启用适配器。

由于适配器可以被停用，我们可以使用同一个模型为 PPO 获取参考和活动的 logits，而无需创建同一模型的两个副本！这利用了 peft 库中的一个特性，即 disable_adapters 上下文管理器。

训练脚本概述：

现在我们将描述如何使用 transformers、peft 和 trl 训练一个 20B 参数的 gpt-neox 模型。本示例的最终目标是在内存受限环境中微调 LLM 以生成正面电影评论。类似的步骤也可应用于其他任务，比如对话模型。

总体上，有三个关键步骤和训练脚本：

  1. 脚本 - 在冻结的 8 位模型上微调低秩适配器，用于在 imdb 数据集上进行文本生成。
  2. 脚本 - 将适配器层合并到基础模型的权重中，并将其存储在 hub 上。
  3. 脚本 - 对低秩适配器进行情感微调，以生成正面评论。

我们在 24GB 的 NVIDIA 4090 GPU 上测试了这些步骤。尽管整个训练过程可以在 24GB GPU 上完成，但完整的训练运行是在 🤗 研究集群中的单个 A100 上进行的。

训练过程的第一步是在预训练模型上进行微调。通常这需要多块高端的 80GB A100 GPU，因此我们选择训练一个低秩适配器。我们将其视为因果语言建模任务，并在 imdb 数据集上的示例中训练了一个周期。该数据集包含电影评论以及表示正面或负面情感的标签。

loss-20b
gpt-neox-20b 模型在 imdb 数据集上进行一个周期训练期间的训练损失

为了拿适配后的模型进行进一步的 RL 微调，我们首先需要合并适配后的权重，这是通过以 16 位浮点加载预训练模型和适配器，并对权重矩阵求和（应用适当的缩放）来实现的。

最后，我们可以在冻结的 imdb 微调模型之上再微调另一个低秩适配器。我们使用一个 imdb 情感分类器为 RL 算法提供奖励。

reward-20b
对 peft 适配的 20B 参数模型进行 RL 微调以生成正面电影评论时的奖励均值。

如果你想查看更多的图表和文本生成结果，完整的 Weights and Biases 报告可在此实验中查看 这里。

结论

我们已在 trl 中实现了一项新功能，允许用户借助 peft 和 bitsandbytes 库，以合理的成本使用 RLHF 对大型语言模型进行微调。我们证明了在 24GB 消费级 GPU 上微调 gpt-neo-x（在 bfloat16 中为 40GB！）是可行的，并且我们预计社区将广泛使用这一集成，利用 RLHF 微调更大的模型并分享出色的成果。

我们已确定了一些有趣的方向，作为下一步来突破这一集成的极限

  • 这将如何在多 GPU 环境下扩展？ 我们将主要探索这一集成如何随 GPU 数量进行扩展，是否可以直接应用数据并行，还是需要在任何相关库中采用一些新功能。
  • 我们可以利用哪些工具来提高训练速度？ 我们观察到，这一集成的主要缺点是整体训练速度。未来我们热衷于探索可能的方向，使训练速度大幅提升。

参考文献

来源：Hugging Face Blog · hf-mirror.com

#产品更新#数据/训练#开源生态#Hugging Face