在 24GB 消费级 GPU 上用 RLHF 微调 200 亿参数大语言模型

我们很高兴正式发布 trl 与 peft 的集成，让任何人都能更轻松地使用强化学习微调大语言模型 (LLM)！在这篇文章中，我们将解释为什么这是一种能与现有微调方法相媲美的替代方案。

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

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

引言

大语言模型与 RLHF

大语言模型 (LLM) 与人类反馈强化学习 (RLHF) 相结合，似乎是构建如 ChatGPT 这般强大 AI 系统的下一个首选方法。

使用 RLHF 训练语言模型通常包括以下三个步骤

1- 在特定领域或指令语料库以及人类演示上微调预训练的 LLM

2- 收集一个由人类标注的数据集并训练一个奖励模型

3- 使用奖励模型和这个数据集，通过强化学习（例如 PPO）进一步微调步骤 1 中的 LLM

基础 LLM 的选择在这里至关重要。在撰写本文时，可以“开箱即用”于许多任务的“最佳”开源 LLM 是经过指令微调的 LLM。著名的模型有：BLOOMZ、Flan-T5、Flan-UL2，以及 OPT-IML。这些模型的缺点是它们的尺寸。要获得一个不错的模型，至少需要使用 100 亿以上参数规模的模型，这在全精度下需要高达 40GB 的 GPU 显存，仅仅是为了在单个 GPU 设备上加载模型，还完全没有进行任何训练！

什么是 TRL？

trl 库旨在让强化学习步骤变得更简单、更灵活，以便任何人都能在自己的自定义数据集和训练设置上使用强化学习微调他们的语言模型。在众多应用中，你可以使用这个算法来微调模型以生成正面的电影评论、进行可控生成或使模型毒性降低。

使用 trl，你可以在分布式环境或单个设备上运行最流行的深度强化学习算法之一，PPO！我们利用 Hugging Face 生态系统中的 accelerate 库来实现这一点，以便任何用户都可以将实验扩展到相当可观的规模。

使用强化学习微调语言模型大致遵循下面详述的协议。这需要原始模型的 2 个副本；为避免活动模型与其原始行为/分布偏离太远，你需要在每个优化步骤计算参考模型的 logits。这给优化过程增加了一个硬性约束，因为每个 GPU 设备上至少需要两个模型副本。如果模型尺寸增大，将整个设置装入单个 GPU 变得越来越棘手。

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

大规模训练

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

如果你使用 AdamW 优化器，每个参数需要 8 字节（例如，如果你的模型有 10 亿个参数，模型的完整 AdamW 优化器将需要 8GB GPU 显存 - 来源）。

已经采用了许多技术来应对这些大规模挑战。最熟悉的范式是流水线并行、张量并行和数据并行。

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

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

8-bit 矩阵乘法

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

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

低秩自适应与 PEFT

2021年，一篇名为《LoRA: 大语言模型的低秩自适应》的论文表明，可以通过冻结预训练权重并创建注意力矩阵中查询和值层的低秩版本来对大语言模型进行微调。这些低秩矩阵的参数比原始模型少得多，从而可以用少得多的 GPU 显存进行微调。作者们证明，微调低秩适配器取得了与微调整个预训练模型相当的结果。

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

什么是 PEFT？

参数高效微调（PEFT）是一个 Hugging Face 库，旨在支持在 LLM 上创建和微调适配器层。peft 与 🤗 Accelerate 无缝集成，利用 DeepSpeed 和大模型推理技术处理大规模模型。

该库支持许多最先进的模型，并有一系列广泛的示例，包括

• 因果语言建模
• 条件生成
• 图像分类
• 8-bit int8 训练
• Dreambooth 模型的低秩自适应
• 语义分割
• 序列分类
• 词元分类

该库仍在广泛和积极的开发中，未来几个月将宣布许多新功能。

使用低秩适配器微调 200 亿参数模型

现在，先决条件都已介绍完毕，让我们一步一步地过一遍整个流程，并通过图表解释如何使用上面提到的工具，在单个 24GB GPU 上用强化学习微调一个 200 亿参数的 LLM！

步骤 1：以 8-bit 精度加载你的活动模型

使用 transformers 实现 LLM 内存减少的“免费午餐”方法是，使用 LLM.int8 中描述的方法以 8-bit 精度加载你的模型。这可以通过在调用 from_pretrained 方法时简单地添加标志 load_in_8bit=True 来实现（你可以在这里阅读更多相关信息）。

如前一节所述，一个计算加载模型所需 GPU 显存量的“技巧”是以“十亿参数”为单位来思考。由于一个字节需要 8 比特，因此全精度模型（32bit = 4bytes）每十亿参数需要 4GB，半精度模型每十亿参数需要 2GB，而 int8 模型每十亿参数需要 1GB。

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

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

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

步骤 3：使用同一模型获取参考和活动 logits

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

训练脚本概述：

我们现在将描述我们如何使用 transformers、peft 和 trl 训练一个 200 亿参数的 gpt-neox 模型。这个例子的最终目标是在内存受限的情况下微调一个 LLM 来生成正面的电影评论。类似的步骤可以应用于其他任务，如对话模型。

总共有三个关键步骤和训练脚本

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

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

训练过程的第一步是在预训练模型上进行微调。通常这需要多个高端的 80GB A100 GPU，所以我们选择训练一个低秩适配器。我们将其视为因果语言建模设置，并在 imdb 数据集的示例上训练一个 epoch，该数据集包含电影评论和表明它们是正面还是负面情绪的标签。

为了使用适配过的模型并用强化学习进行进一步微调，我们首先需要合并适配的权重，这是通过以 16-bit 浮点数加载预训练模型和适配器，并与权重矩阵相加（应用了适当的缩放）来实现的。

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

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

结论

我们在 trl 中实现了一项新功能，允许用户通过利用 peft 和 bitsandbytes 库，以合理的成本使用 RLHF 微调大语言模型。我们证明了在 24GB 消费级 GPU 上微调 gpt-neo-x（在 bfloat16 下为 40GB！）是可能的，我们期待社区能广泛使用这一集成来微调更大的模型，利用 RLHF 并分享优秀的成果。

我们已经确定了下一步推动该集成极限的一些有趣方向

• 这在多 GPU 设置中将如何扩展？我们主要将探索这种集成在 GPU 数量方面的扩展情况，是否可以开箱即用地应用数据并行，或者是否需要在任何相关库中采用新功能。
• 我们可以利用哪些工具来提高训练速度？我们观察到，这种集成的主要缺点是整体训练速度。未来我们希望能探索可能的方向，使训练速度更快。

参考文献

• 并行范式： https://huggingface.co/docs/transformers/v4.17.0/en/parallelism
• transformers 中的 8-bit 集成： https://huggingface.co/blog/hf-bitsandbytes-integration
• LLM.int8 论文： https://arxiv.org/abs/2208.07339
• 梯度检查点解释： https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html