在一张 24 GB 的消费级显卡上用 RLHF 微调 20B LLMs

huggingface.co/blog/zh/trl…

咱们很快乐正式发布 trl 与 peft 的集成，使任何人都能够更轻松地运用强化学习进行大型言语模型 (LLM) 微调！在这篇文章中，咱们解说了为什么这是现有微调办法的有竞争力的替代计划。

请注意， peft 是一种通用东西，能够使用于许多 ML 用例，但它对 RLHF 特别风趣，因为这种办法特别需求内存！

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

介绍

LLMs & RLHF

LLM 结合 RLHF (人类反应强化学习) 似乎是构建非常强大的 AI 体系 (例如 ChatGPT) 的下一个首选办法。

运用 RLHF 练习言语模型一般包含以下三个进程:

1. 在特定领域或指令和人类演示语料库上微调预练习的 LLM;
2. 收集人类标示的数据集，练习一个奖赏模型;
3. 运用 RL (例如 PPO)，用此数据集和奖赏模型进一步微调进程 1 中的 LLM。

根底 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 中，你还能够在参阅模型和活泼模型之间运用共享层以防止整个副本。 模型解毒示例中展现了此功用的具体示例。

大规模练习

大规模练习是具有应战性的。第一个应战是在可用的 GPU 设备上拟合模型，及其优化器状况。 单个参数占用的 GPU 内存量取决于其“精度”(或更具体地说是 dtype)。 最常见的 dtype 是 float32 (32 位) 、 float16 和 bfloat16 (16 位)。 最近，“奇特的”精度支撑开箱即用的练习和推理 (具有特定条件和束缚)，例如 int8 (8 位)。 简而言之，要在 GPU 设备上加载一个模型，每十亿个参数在 float32 精度上需求 4GB，在 float16 上需求 2GB，在 int8 上需求 1GB。 假如你想了解关于这个话题的更多信息，请查看这篇研究深入的 文章。

假如您运用 AdamW 优化器，每个参数需求 8 个字节 (例如，假如您的模型有 1B 个参数，则模型的完整 AdamW 优化器将需求 8GB GPU 内存 来历)。

许多技能现已被采用以应对大规模练习上的应战。最了解的范式是管道并行、张量并行和数据并行。

经过数据并行性，同一模型并行托管在多台机器上，而且每个实例都被供给不同的数据批次。 这是最直接的并行策略，本质上是复制单 GPU 的状况，而且现已被 trl 支撑。 运用管道并行和张量并行，模型自身分布在机器上: 在管道并行中，模型按层拆分，而张量并行则跨 GPU 拆分张量操作 (例如矩阵乘法)。运用这些模型并行策略，你需求将模型权重分片到许多设备上，这需求你定义跨进程的激活和梯度的通信协议。 这完成起来并不简略，或许需求采用一些框架，例如 [Megatron-DeepSpeed](https://github.com/microsoft/Megatron-DeepSpeed) 或 [Nemo](https://github.com/NVIDIA/NeMo)。其他对扩展练习至关重要的东西也需求被强调，例如自习惯激活检查点和交融内核。 能够在 扩展阅览 找到有关并行范式的进一步阅览。

因而，咱们问自己下面一个问题: 仅用数据并行咱们能够走多远？咱们能否运用现有的东西在单个设备中习惯超大型练习进程 (包含活泼模型、参阅模型和优化器状况)？ 答案似乎是必定的。 首要因素是: 适配器和 8 位矩阵乘法！ 让咱们在以下部分中介绍这些主题:

8 位矩阵乘法

高效的 8 位矩阵乘法是论文 LLM.int8() 中初次引入的一种办法，旨在处理量化大规模模型时的性能下降问题。 所提出的办法将在线性层中使用的矩阵乘法分解为两个阶段: 在 float16 中将被履行的异常值躲藏状况部分和在 int8 中被履行的“非异常值”部分。

简而言之，假如运用 8 位矩阵乘法，则能够将全精度模型的巨细减小到 4 分之一 (因而，对于半精度模型，能够减小 2 分之一)。

低秩适配和 PEFT

在 2021 年，一篇叫 LoRA: Low-Rank Adaption of Large Language Models 的论文标明，能够经过冻住预练习权重，并创立查询和值层的注意力矩阵的低秩版原本对大型言语模型进行微调。这些低秩矩阵的参数远少于原始模型，因而能够运用更少的 GPU 内存进行微调。 作者证明，低阶适配器的微调取得了与微调完整预练习模型相当的成果。

这种技能答应运用一小部分内存来微调 LLM。 然而，也有一些缺陷。因为适配器层中的额定矩阵乘法，前向和反向传递的速度大约是原来的两倍。

什么是 PEFT？

Parameter-Efficient Fine-Tuning (PEFT) 是一个 Hugging Face 的库，它被创造出来以支撑在 LLM 上创立和微调适配器层。 peft 与 Accelerate 无缝集成，用于运用了 DeepSpeed 和 Big Model Inference 的大规模模型。

此库支撑许多先进的模型，而且有大量的比如，包含:

• 因果言语建模
• 条件生成
• 图画分类
• 8 位 int8 练习
• Dreambooth 模型的低秩适配
• 语义分割
• 序列分类
• 词符分类

该库仍在广泛和积极的开发中，许多即将推出的功用将在未来几个月内公布。

运用低质适配器微调 20B 参数量的模型

现在先决条件现已处理，让咱们一步步过一遍整个管道，并用图阐明如安在单个 24GB GPU 上运用上述东西运用 RL 微调 20B 参数量的 LLM！

第 1 步: 在 8 位精度下加载你的活泼模型

运用 transformers 削减 LLM 内存的“免费午饭”是运用 LLM.int8 中描述的办法，以 8 位精度加载模型。 这能够经过在调用 from_pretrained 办法时简略地增加标志 load_in_8bit=True 来履行 (你能够在 文档中 阅览更多相关信息)。

如前一节所述，核算加载模型所需的 GPU 内存量的“技巧”是依据“十亿个参数量”进行考虑。 因为一个字节需求 8 位，因而全精度模型 (32 位 = 4 字节) 每十亿个参数需求 4GB，半精度模型每十亿个参数需求 2GB，int8 模型每十亿个参数需求 1GB。

所以首要，咱们只加载 8 位的活泼模型。 让咱们看看第二步需求做什么！

第 2 步: 运用 peft 增加额定可练习的适配器

第二步是在模型中加载适配器并使这些适配器可练习。 这能够大幅削减活泼模型所需的可练习权重的数量。 此进程运用 peft 库，只需几行代码即可履行。 请注意，一旦适配器经过练习，您就能够轻松地将它们推送到 Hub 以供以后运用。

第 3 步: 运用相同的模型得到参阅和活泼 logits

因为适配器能够停用，咱们能够运用相同的模型来获取 PPO 的参阅和活泼的 logits 值，而无需创立两个相同的模型副本！ 这运用了 peft 库中的一个功用，即 disable_adapters 上下文管理器。

练习脚本概述

咱们现在将描述怎么运用 transformers 、 peft 和 trl 练习 20B 参数量的 gpt-neox 模型。 这个比如的最终目标是微调 LLM 以在内存受限的设置中生成积极的电影谈论。类似的进程能够使用于其他使命，例如对话模型。

全体来看有三个关键进程和练习脚本:

1. 脚本 1 – 在冻住的 8 位模型上微调低秩适配器，以便在 imdb 数据集上生成文本。
2. 脚本 2 – 将适配器层合并到根底模型的权重中并将它们存储在 Hub 上。
3. 脚本 3 – 对低等级适配器进行情感微调以创立正面评价。

咱们在 24GB NVIDIA 4090 GPU 上测试了这些进程。虽然能够在 24 GB GPU 上履行整个练习进程，但在 研究集群上的单个 A100 上无法进行完整的练习进程。

练习进程的第一步是对预练习模型进行微调。 一般这需求几个高端的 80GB A100 GPU，因而咱们选择练习低阶适配器。 咱们将其视为一种因果言语建模设置，并针从 imdb 数据集中练习了一个 epoch 的示例，该数据集具有电影谈论和指明积极还是消极情绪的标签。

为了运用现已适配了的模型并运用 RL 履行进一步微调，咱们首要需求组合自习惯权重，这是经过加载预练习模型和 16 位浮点精度的适配器而且加和权重矩阵 (使用恰当的缩放份额) 来完成的。

最后，咱们能够在冻住的、用 imdb 微调过的模型之上微调另一个低秩适配器。咱们运用一个 imdb 情感分类器 来为 RL 算法供给奖赏。

假如您想查看更多图表和文本生成，可在 链接处 获取此试验的完整权重和误差陈述。

定论

咱们在 trl 中完成了一项新功用，答使用户运用 peft 和 bitsandbytes 库以合理的成本运用 RLHF 微调大型言语模型。 咱们证明了能够在 24GB 消费级 GPU 上微调 gpt-neo-x (以 bfloat16 精度需求 40GB！)，咱们期望社区将广泛运用此集成来微调运用了 RLHF 的大型模型，并分享出色的工件。

咱们现已为接下来的进程确认了一些风趣的方向，以应战这种集成的极限:

• 这将如安在多 GPU 设置中扩展？ 咱们将首要探究这种集成将怎么依据 GPU 的数量进行扩展，是否能够开箱即用地使用数据并行，或许是否需求在任何相关库上采用一些新功用。
• 咱们能够运用哪些东西来提高练习速度？ 咱们观察到这种集成的首要缺陷是全体练习速度。 在未来，咱们将继续探究使练习更快的或许方向。

参阅

• 并行范式: hf.co/docs/transf…
• transformers 中的 8 位集成: hf.co/blog/hf-bit…
• LLM.int8 论文: arxiv.org/abs/2208.07…
• 梯度检查点解说: docs.aws.amazon.com/sagemaker/l…