我们很高兴正式发布 trl
与 peft
的集成,使任何人都可以更轻松地使用强化学习进行大型语言模型 (LLM) 微调!在这篇文章中,我们解释了为什么这是现有微调方法的有竞争力的替代方案。
请注意, peft
是一种通用工具,可以应用于许多 ML 用例,但它对 RLHF 特别有趣,因为这种方法特别需要内存!
如果你想直接深入研究代码,请直接在 TRL 的文档页面 直接查看示例脚本。
介绍
LLMs & RLHF
LLM 结合 RLHF (人类反馈强化学习) 似乎是构建非常强大的 AI 系统 (例如 ChatGPT) 的下一个首选方法。
使用 RLHF 训练语言模型通常包括以下三个步骤:
- 在特定领域或指令和人类示范语料库上微调预训练的 LLM;
- 收集人类标注的数据集,训练一个奖励模型;
- 使用 RL (例如 PPO),用此数据集和奖励模型进一步微调步骤 1 中的 LLM。
ChatGPT 的训练协议概述,从数据收集到 RL 部分。 资料来源: OpenAI 的 ChatGPT 博文
基础 LLM 的选择在这里是至关重要的。在撰写本文时,可以“开箱即用”地用于许多任务的“最佳”开源 LLM 是指令微调 LLMs。著名的模型有: BLOOMZ Flan-T5、Flan-UL2 和 OPT-IML。这些模型的缺点是它们的尺寸。要获得一个像样的模型,你至少需要玩 10B+ 级别的模型,在全精度情况下这将需要高达 40GB GPU 内存,只是为了将模型装在单个 GPU 设备上而不进行任何训练!
什么是 TRL?
trl
库的目的是使 RL 的步骤更容易和灵活,让每个人可以在他们自己的数据集和训练设置上用 RL 微调 LM。在许多其他应用程序中,你可以使用此算法微调模型以生成 正面电影评论、进行 受控生成 或 降低模型的毒性。
使用 trl
你可以在分布式管理器或者单个设备上运行最受欢迎的深度强化学习算法之一: PPO。我们利用 Hugging Face 生态系统中的 accelerate
来实现这一点,这样任何用户都可以将实验扩大到一个有趣的规模。
使用 RL 微调语言模型大致遵循下面详述的协议。这需要有 2 个原始模型的副本; 为避免活跃模型与其原始行为/分布偏离太多,你需要在每个优化步骤中计算参考模型的 logits 。这对优化过程增加了硬约束,因为你始终需要每个 GPU 设备至少有两个模型副本。如果模型的尺寸变大,在单个 GPU 上安装设置会变得越来越棘手。
TRL 中 PPO 训练设置概述
在 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 个字节 (例如,如果您的模型有 1B 个参数,则模型的完整 AdamW 优化器将需要 8GB GPU 内存 来源)。
许多技术已经被采用以应对大规模训练上的挑战。最熟悉的范式是管道并行、张量并行和数据并行。
图片来自 这篇博文
通过数据并行性,同一模型并行托管在多台机器上,并且每个实例都被提供不同的数据批次。 这是最直接的并行策略,本质上是复制单 GPU 的情况,并且已经被 trl
支持。 使用管道并行和张量并行,模型本身分布在机器上: 在管道并行中,模型按层拆分,而张量并行则跨 GPU 拆分张量操作 (例如矩阵乘法)。使用这些模型并行策略,你需要将模型权重分片到许多设备上,这需要你定义跨进程的激活和梯度的通信协议。 这实现起来并不简单,可能需要采用一些框架,例如 Megatron-DeepSpeed
或 Nemo
。其他对扩展训练至关重要的工具也需要被强调,例如自适应激活检查点和融合内核。 可以在 扩展阅读 找到有关并行范式的进一步阅读。
因此,我们问自己下面一个问题: 仅用数据并行我们可以走多远?我们能否使用现有的工具在单个设备中适应超大型训练过程 (包括活跃模型、参考模型和优化器状态)? 答案似乎是肯定的。 主要因素是: 适配器和 8 位矩阵乘法! 让我们在以下部分中介绍这些主题:
8 位矩阵乘法
高效的 8 位矩阵乘法是论文 LLM.int8() 中首次引入的一种方法,旨在解决量化大规模模型时的性能下降问题。 所提出的方法将在线性层中应用的矩阵乘法分解为两个阶段: 在 float16 中将被执行的异常值隐藏状态部分和在 int8 中被执行的“非异常值”部分。
高效的 8 位矩阵乘法是论文 LLM.int8() 中首次引入的一种方法,旨在解决量化大规模模型时的性能下降问题。 所提出的方法将在线性层中应用的矩阵乘法分解为两个阶段: 在 float16 中被执行的异常值隐藏状态部分和在 int8 中被执行的“非异常值”部分。
简而言之,如果使用 8 位矩阵乘法,则可以将全精度模型的大小减小到 4 分之一 (因此,对于半精度模型,可以减小 2 分之一)。
低秩适配和 PEFT
在 2021 年,一篇叫 LoRA: Low-Rank Adaption of Large Language Models 的论文表明,可以通过冻结预训练权重,并创建查询和值层的注意力矩阵的低秩版本来对大型语言模型进行微调。这些低秩矩阵的参数远少于原始模型,因此可以使用更少的 GPU 内存进行微调。 作者证明,低阶适配器的微调取得了与微调完整预训练模型相当的结果。
原始 (冻结的) 预训练权重 (左) 的输出激活由一个由权重矩阵 A 和 B 组成的低秩适配器 (右) 增强。
这种技术允许使用一小部分内存来微调 LLM。 然而,也有一些缺点。由于适配器层中的额外矩阵乘法,前向和反向传递的速度大约是原来的两倍。
什么是 PEFT?
Parameter-Efficient Fine-Tuning (PEFT) 是一个 Hugging Face 的库,它被创造出来以支持在 LLM 上创建和微调适配器层。 peft
与