在24GB的消费级GPU上使用强化学习高阶反馈对20B LLM进行微调
在24GB GPU上使用强化学习高阶反馈微调20B LLM
我们很高兴正式发布与 peft 的集成，使得使用增强学习进行大型语言模型 (LLM) 微调更加容易接近任何人！在本文中，我们将解释为什么这是现有微调方法的一个有竞争力的替代方案。
请注意 peft 是一个通用工具，可以应用于许多机器学习用例，但对于 RLHF 来说特别有趣，因为该方法对内存的需求特别高！
如果您想直接深入了解代码，请查看 TRL 文档页面上的示例脚本。
介绍
LLMs 和 RLHF
LLMs 结合 RLHF（用人类反馈进行增强学习）似乎是构建非常强大的人工智能系统（如 ChatGPT）的下一个方法。
使用 RLHF 对语言模型进行训练通常包括以下三个步骤：
1- 在特定领域或指令和人类演示的语料库上对预训练的 LLM 进行微调
2- 收集人类注释数据集并训练奖励模型
3- 使用奖励模型和该数据集使用 RL（例如 PPO）进一步微调步骤 1 中的 LLM
在这里选择基础 LLM 是相当关键的。在撰写本文时，“最佳”的开源 LLM，可以直接用于许多任务的是指令微调的 LLM。值得注意的模型有：BLOOMZ、Flan-T5、Flan-UL2 和 OPT-IML。这些模型的缺点是它们的尺寸。为了获得一个体面的模型，您至少需要使用 10B+ 规模的模型，这将需要高达 40GB 的 GPU 内存以完整精度在单个 GPU 设备上适应模型，而不进行任何训练！
TRL 是什么？
trl 库旨在使 RL 步骤更加简单和灵活，以便任何人都可以使用 RL 在自定义数据集和训练设置上微调他们的语言模型。除了许多其他应用外，您可以使用此算法来微调模型以生成正面的电影评论、进行受控生成或使模型更少有毒。
使用 trl，您可以以分布式方式或在单个设备上运行最流行的深度 RL 算法之一 PPO！我们利用 Hugging Face 生态系统中的 accelerate 来实现这一点，以便任何用户都可以将实验扩展到有趣的规模。
使用 RL 对语言模型进行微调大致遵循下面详细的协议。这需要拥有原始模型的 2 个副本；为了避免活动模型在优化过程中偏离其原始行为/分布，您需要在每个优化步骤中计算参考模型的 logits。这给优化过程增加了一个硬约束，因为您始终需要至少两个模型的副本每个 GPU 设备。如果模型的大小增加，将变得越来越难以将设置适应单个 GPU。
在 trl 中，您还可以在参考模型和活动模型之间使用共享层以避免整个副本。这个特性的一个具体示例在解毒示例中展示。
规模训练
规模训练可能具有挑战性。第一个挑战是将模型及其优化器状态适应可用的 GPU 设备上。单个参数占用的 GPU 内存量取决于其“精度”（更具体地说是 dtype）。最常见的 dtype 是 float32（32 位）、float16 和 bfloat16（16 位）。最近，“奇特”的精度在训练和推断中得到了开箱即用的支持（在某些条件和约束下），例如 int8（8 位）。简而言之，为了在 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 训练至关重要的工具，例如自适应激活检查点和融合内核。关于并行范式的更多阅读资料可以在这里找到。
因此，我们提出了以下问题：仅使用数据并行能够走多远？我们能否使用现有工具将超大规模的训练过程（包括活跃模型、参考模型和优化器状态）适应到单个设备上？答案似乎是肯定的。主要的因素是：适配器和8位矩阵乘法！让我们在以下部分详细介绍这些主题：
8 位矩阵乘法
高效的8位矩阵乘法是一种在 LLM.int8() 论文中首次引入的方法，旨在解决量化大规模模型时性能下降的问题。该方法将在线性层中在底层应用的矩阵乘法分为两个阶段：将要在 float16 下执行的异常隐藏状态部分和在 int8 下执行的“非异常”部分。
简而言之，如果使用8位矩阵乘法，可以将全精度模型的大小减小4倍（因此对于半精度模型减小2倍）。
低秩适应和 PEFT
2021年，一篇名为 LoRA: Low-Rank Adaption of Large Language Models 的论文证明了可以通过冻结预训练权重并创建查询和值层注意力矩阵的低秩版本来进行大语言模型的微调。这些低秩矩阵的参数远少于原始模型，可以使用更少的 GPU 内存进行微调。作者证明了低秩适配器的微调结果与完整预训练模型的微调结果相当。
这种技术允许使用较少的内存要求对 LLM 进行微调。然而，也存在一些缺点。由于适配器层中存在额外的矩阵乘法，前向传播和反向传播大约会慢两倍。
PEFT 是什么？
参数高效微调（PEFT）是 Hugging Face 的一个库，用于支持在 LLM 上创建和微调适配器层。peft 与 🤗 Accelerate 无缝集成，可以利用 DeepSpeed 和大模型推理来支持大规模模型。
该库支持许多最先进的模型，并具有广泛的示例集，包括：
- 因果语言建模
- 条件生成
- 图像分类
- 8位 int8 训练
- Dreambooth 模型的低秩适应
- 语义分割
- 序列分类
- 标记分类
该库仍在广泛而积极地开发中，未来几个月将发布许多即将推出的功能。
使用低秩适配器微调 20B 参数模型
现在，我们已经搞清楚了前提条件，让我们逐步介绍整个流程，并通过图示解释如何在单个 24GB GPU 上使用上述工具微调一个具有 20B 参数的 LLM！
步骤 1：以8位精度加载活跃模型
使用 transformers 的一种“免费午餐”内存减少方法是使用 LLM.int8 中描述的方法以8位精度加载模型。只需要在调用 from_pretrained 方法时添加标志 load_in_8bit=True 即可（您可以在这里阅读更多信息）。
如前一节所述，计算需要加载模型所需的GPU内存的“黑科技”是根据“参数数量的十亿级”来思考。对于全精度模型（32位 = 4字节），每十亿个参数需要4GB内存，对于半精度模型每十亿个参数需要2GB内存，对于int8模型每十亿个参数需要1GB内存。
因此，首先让我们只加载8位的活跃模型。让我们看看第二步需要做什么！
第二步：使用peft添加额外的可训练的适配器
第二步是在模型中加载适配器并使这些适配器可训练。这将极大地减少活跃模型所需的可训练权重数量。这一步利用了peft库，可以用几行代码完成。请注意，一旦适配器被训练，您可以轻松将它们推送到Hub以便以后使用。
第三步：使用相同的模型获取参考和活跃的logits
由于适配器可以被停用，我们可以使用相同的模型获取PPO的参考和活跃的logits，而无需创建两个相同模型的副本！这利用了peft库中的一个特性，即disable_adapters上下文管理器。
训练脚本概述：
我们现在将描述如何使用transformers、peft和trl训练一个20B参数的gpt-neox模型。此示例的最终目标是在内存受限的环境中微调LLM以生成正面电影评论。类似的步骤也可以应用于其他任务，如对话模型。
总体而言，有三个关键步骤和训练脚本：
- 脚本 – 在imdb数据集上对冻结的8位模型进行文本生成的低秩适配器微调。
- 脚本 – 将适配器层合并到基础模型的权重中，并将其存储在hub上。
- 脚本 – 对低秩适配器进行情感微调以创建正面评论。
我们在一块24GB的NVIDIA 4090 GPU上测试了这些步骤。虽然完全可以在24GB的GPU上执行整个训练过程，但完整的训练运行是在单个A100上进行的，该A100位于🤗研究集群上。
训练过程中的第一步是对预训练模型进行微调。通常，这需要几个高端的80GB A100 GPU，所以我们选择了训练一个低秩适配器。我们将其视为因果语言建模设置，并对imdb数据集中的示例进行了一个时期的训练，该数据集包含电影评论和指示其情感是正面还是负面的标签。
为了将适应后的模型与RL一起进行进一步的微调，我们首先需要合并适应后的权重，这通过以16位浮点数加载预训练模型和适配器，并进行带有权重矩阵的摘要来实现（应用适当的缩放）。
最后，我们可以在冻结的imdb微调模型之上微调另一个低秩适配器。我们使用imdb情感分类器为RL算法提供奖励。
完整的Weights and Biases报告可以在此处找到，如果您想查看更多的图表和文本生成。
结论
我们在trl中实现了一项新功能，允许用户使用peft和bitsandbytes库以合理的成本微调大型语言模型。我们演示了在24GB的消费级GPU上对40GB的gpt-neo-x（使用bfloat16！）进行微调是可能的，我们预计这种集成将被广泛应用于使用RLHF微调更大模型并共享优秀成果的社区。
我们已经确定了下一步推动这种集成极限的一些有趣方向
- 在多GPU环境下，这个模型在扩展性方面会如何表现？我们主要探讨这个集成在GPU数量方面的扩展性，是否可以直接应用数据并行处理，或者是否需要在任何相关库中采用一些新的特性。
- 我们可以利用哪些工具来提高训练速度？我们观察到这个集成的主要缺点是整体训练速度较慢。将来，我们将乐于探索训练加速的可能方向。
参考资料
- 并行处理范式：https://huggingface.co/docs/transformers/v4.17.0/en/parallelism
transformers中的8位整合：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
