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

科技资讯 投稿 28700℃ 0 评论

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

trl 与 peft 的集成,使任何人都可以更轻松地使用强化学习进行大型语言模型 (LLM 微调!在这篇文章中,我们解释了为什么这是现有微调方法的有竞争力的替代方案。

peft 是一种通用工具,可以应用于许多 ML 用例,但它对 RLHF 特别有趣,因为这种方法特别需要内存!

介绍

LLMs & RLHF

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

    在特定领域或指令和人类示范语料库上微调预训练的 LLM;
  1. 收集人类标注的数据集,训练一个奖励模型;
  2. 使用 RL (例如 PPO,用此数据集和奖励模型进一步微调步骤 1 中的 LLM。

ChatGPT 的训练协议概述,从数据收集到 RL 部分。 资料来源: OpenAI 的 ChatGPT 博文

什么是 TRL?

trl 库的目的是使 RL 的步骤更容易和灵活,让每个人可以在他们自己的数据集和训练设置上用 RL 微调 LM。在许多其他应用程序中,你可以使用此算法微调模型以生成 正面电影评论、进行 受控生成 或 降低模型的毒性。

trl 你可以在分布式管理器或者单个设备上运行最受欢迎的深度强化学习算法之一: PPO。我们利用 Hugging Face 生态系统中的 accelerate 来实现这一点,这样任何用户都可以将实验扩大到一个有趣的规模。

TRL 中 PPO 训练设置概述

trl 中,你还可以在参考模型和活跃模型之间使用共享层以避免整个副本。 模型解毒示例中展示了此功能的具体示例。

大规模训练

dtype。 最常见的 dtype 是 float32 (32 位 、 float16 和 bfloat16 (16 位。 最近,“奇异的”精度支持开箱即用的训练和推理 (具有特定条件和约束,例如 int8 (8 位。 简而言之,要在 GPU 设备上加载一个模型,每十亿个参数在 float32 精度上需要 4GB,在 float16 上需要 2GB,在 int8 上需要 1GB。 如果你想了解关于这个话题的更多信息,请查看这篇研究深入的文章: https://huggingface.co/blog/hf-bitsandbytes-integration。

许多技术已经被采用以应对大规模训练上的挑战。最熟悉的范式是管道并行、张量并行和数据并行。

图片来自 这篇博文

trl 支持。 使用管道并行和张量并行,模型本身分布在机器上: 在管道并行中,模型按层拆分,而张量并行则跨 GPU 拆分张量操作 (例如矩阵乘法。使用这些模型并行策略,你需要将模型权重分片到许多设备上,这需要你定义跨进程的激活和梯度的通信协议。 这实现起来并不简单,可能需要采用一些框架,例如 Megatron-DeepSpeed 或 Nemo。其他对扩展训练至关重要的工具也需要被强调,例如自适应激活检查点和融合内核。 可以在 扩展阅读 找到有关并行范式的进一步阅读。

8 位矩阵乘法

高效的 8 位矩阵乘法是论文 LLM.int8( 中首次引入的一种方法,旨在解决量化大规模模型时的性能下降问题。 所提出的方法将在线性层中应用的矩阵乘法分解为两个阶段: 在 float16 中将被执行的异常值隐藏状态部分和在 int8 中被执行的“非异常值”部分。

简而言之,如果使用 8 位矩阵乘法,则可以将全精度模型的大小减小到 4 分之一 (因此,对于半精度模型,可以减小 2 分之一。

低秩适配和 PEFT

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

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

什么是 PEFT?

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

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

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

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

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

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

transformers 减少 LLM 内存的“免费午餐”是使用 LLM.int8 中描述的方法,以 8 位精度加载模型。 这可以通过在调用 from_pretrained 方法时简单地添加标志 load_in_8bit=True 来执行 (你可以在 文档中 阅读更多相关信息。

所以首先,我们只加载 8 位的活跃模型。 让我们看看第二步需要做什么!

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

您可以轻松地在冻结的 8 位模型上添加适配器,从而通过训练一小部分参数来减少优化器状态的内存需求

peft 库,只需几行代码即可执行。 请注意,一旦适配器经过训练,您就可以轻松地将它们推送到 Hub 以供以后使用。

第 3 步: 使用同样的模型得到参考和活跃 logits

你可以方便地使用 peft 关闭和使能适配器

peft 库中的一个功能,即 disable_adapters 上下文管理器。

训练脚本概述

transformers 、 peft 和 trl 训练 20B 参数量的 gpt-neox 模型。 这个例子的最终目标是微调 LLM 以在内存受限的设置中生成积极的电影评论。类似的步骤可以应用于其他任务,例如对话模型。

    脚本 1 - 在冻结的 8 位模型上微调低秩适配器,以便在 imdb 数据集上生成文本。

  1. 脚本 2 - 将适配器层合并到基础模型的权重中并将它们存储在 Hub 上。

  2. 脚本 3 - 对低等级适配器进行情感微调以创建正面评价。

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

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

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

RL 微调 peft 适配过的 20B 参数量的模型以生成积极影评时的奖励均值。

结论

我们在 trl 中实现了一项新功能,允许用户利用 peft 和 bitsandbytes 库以合理的成本使用 RLHF 微调大型语言模型。 我们证明了可以在 24GB 消费级 GPU 上微调 gpt-neo-x (以 bfloat16 精度需要 40GB!,我们期望社区将广泛使用此集成来微调利用了 RLHF 的大型模型,并分享出色的工件。

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

参考

    并行范式: https://hf.co/docs/transformers/v4.17.0/en/parallelism
  • transformers 中的 8 位集成: https://hf.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

英文原文: https://hf.co/blog/trl-peft

译者: AIboy1993 (李旭东

审校、排版: zhongdongy (阿东

编程笔记 » 在一张 24 GB 的消费级显卡上用 RLHF 微调 20B LLMs

赞同 (144) or 分享 (0)
游客 发表我的评论   换个身份
取消评论

表情
(0)个小伙伴在吐槽