StackLLaMA：使用 RLHF 训练 LLaMA 的实践指南

LLaMA模型

在进行 RLHF 时，重要的是从一个有能力的模型开始：RLHF 这一步只是一个微调步骤，可以使模型与我们想要与之交互的方式以及我们期望它的响应方式保持一致。因此，我们选择使用最近推出的高性能LLaMA 模型。LLaMA 模型是 Meta AI 开发的最新大型语言模型。它们的大小从 7B 到 65B 不等，并接受了 1T 到 1.4T 令牌的训练，使它们非常有能力。我们使用 7B 模型作为以下所有步骤的基础！要访问该模型，请使用Meta AI 中的表单。

堆栈交换数据集

收集人类反馈是一项复杂而昂贵的工作。为了在构建有用模型的同时引导此示例的过程，我们使用了StackExchange 数据集。该数据集包括来自 StackExchange 平台的问题及其相应的答案（包括针对代码和许多其他主题的 StackOverflow）。这个用例很有吸引力，因为答案与赞成票的数量和已接受答案的标签一起出现。

我们遵循Askell 等人中描述的方法。2021并为每个答案分配一个分数：

score = log2 (1 + upvotes) rounded to the nearest integer, plus 1 if the questioner accepted the answer (we assign a score of −1 if the number of upvotes is negative).

对于奖励模型，我们将始终需要每个问题有两个答案进行比较，我们稍后会看到。有些问题有几十个答案，导致许多可能的对。我们对每个问题最多采样 10 个答案对，以限制每个问题的数据点数量。最后，我们通过将 HTML 转换为 Markdown 来清理格式，使模型的输出更具可读性。您可以在此处找到数据集和处理笔记本。

高效的训练策略

即使训练最小的 LLaMA 模型也需要大量内存。一些快速数学运算：在 bf16 中，每个参数使用 2 个字节（在 fp32 中使用 4 个字节）以及使用的 8 个字节，例如，在 Adam 优化器中（有关更多信息，请参阅 Transformers 中的性能文档）。因此，一个 7B 参数模型将(2+8)*7B=70GB仅用于适应内存，并且在计算注意力分数等中间值时可能需要更多。因此，即使像这样，您也无法在单个 80GB A100 上训练模型。您可以使用一些技巧，例如更高效的半精度训练优化器，将更多内存挤入内存，但迟早会用完。

另一种选择是使用参数高效微调 (PEFT) 技术，例如peft可以对以 8 位加载的模型执行低秩自适应 (LoRA) 的库。

线性层的低秩自适应：在冻结层（蓝色）旁边添加额外参数（橙色），并将生成的编码隐藏状态与冻结层的隐藏状态一起添加。

以 8 位加载模型会大大减少内存占用，因为每个参数只需要一个字节用于权重（例如，7B LlaMa 的内存为 7GB）。LoRA 不是直接训练原始权重，而是在某些特定层（通常是注意力层）之上添加小的适配器层；因此，可训练参数的数量大大减少。

在这种情况下，经验法则是为每十亿个参数分配 ~1.2-1.4GB（取决于批量大小和序列长度）以适应整个微调设置。正如上面附带的博客文章中详述的那样，这可以以低成本微调更大的模型（在 NVIDIA A100 80GB 上高达 50-60B 比例模型）。

这些技术可以在消费设备和 Google Colab 上微调大型模型。值得注意的演示是微调facebook/opt-6.7b（13GB in float16）和openai/whisper-largeGoogle Colab（15GB GPU RAM）。要了解有关使用 的更多信息peft，请参阅我们的github 存储库或之前关于在消费类硬件上训练 20b 参数模型的博文( https://huggingface.co/blog/trl-peft ))。

现在我们可以将非常大的模型放入单个 GPU 中，但训练速度可能仍然很慢。这种情况下最简单的策略是数据并行：我们将相同的训练设置复制到不同的 GPU 中，并将不同的批次传递给每个 GPU。有了这个，您可以并行化模型的前向/后向传递，并根据 GPU 的数量进行扩展。

我们使用transformers.Traineror accelerate，它们都支持数据并行性而无需任何代码更改，只需在使用torchrunor调用脚本时传递参数即可accelerate launch。下面分别用accelerate和在单台机器上运行带有 8 个 GPU 的训练脚本torchrun。

监督微调

在我们开始训练奖励模型并使用 RL 调整我们的模型之前，如果该模型在我们感兴趣的领域已经很好，它会有所帮助。在我们的例子中，我们希望它回答问题，而对于其他用例，我们可能希望它遵循说明，在这种情况下，指令调整是个好主意。实现这一目标的最简单方法是继续使用来自领域或任务的文本的语言建模目标来训练语言模型。StackExchange 数据集非常庞大（超过 1000 万条指令），因此我们可以轻松地在其中的一个子集上训练语言模型。

在进行 RLHF 之前微调模型没有什么特别之处——它只是我们在此处应用的预训练的因果语言建模目标。为了有效地使用数据，我们使用了一种称为打包的技术：我们不是在批处理中为每个样本使用一个文本，然后填充到模型的最长文本或最大上下文，而是将大量文本与 EOS 标记连接在一起之间和剪切上下文大小的块以填充批处理而没有任何填充。

使用这种方法，训练效率更高，因为通过模型传递的每个标记也经过训练，这与通常掩盖损失的填充标记不同。如果您没有太多数据并且更关心偶尔切断一些溢出上下文的标记，您也可以使用经典数据加载器。

包装由 处理，然后我们可以在使用 加载模型后ConstantLengthDataset使用。首先，我们在 int8 中加载模型，准备训练，然后添加 LoRA 适配器。Trainerpeft

我们使用因果语言建模目标训练模型几千步并保存模型。由于我们将根据不同的目标再次调整模型，因此我们将适配器权重与原始模型权重合并。

免责声明：由于 LLaMA 的许可，我们仅发布此适配器权重和以下部分中的模型检查点。您可以通过填写 Meta AI 的表单申请访问基础模型的权重，然后通过运行此脚本将它们转换为 🤗 Transformers 格式。请注意，在发布之前，您还需要从源代码安装 🤗 Transformers v4.28。

现在我们已经为任务微调了模型，我们准备好训练奖励模型了。

奖励建模和人类偏好

原则上，我们可以直接使用 RLHF 和人工注释来微调模型。然而，这将需要我们在每次优化迭代后将一些样本发送给人类进行评级。由于收敛所需的训练样本数量以及人类阅读和注释器速度的固有延迟，这是昂贵且缓慢的。

一个代替直接反馈效果很好的技巧是在 RL 循环之前收集的人类注释上训练奖励模型。奖励模型的目标是模仿人类如何评价文本。建立奖励模型有几种可能的策略：最直接的方法是预测注释（例如评分或“好”/“坏”的二进制值）。在实践中，更好的方法是预测两个例子的排名，其中奖励模型有两个候选者(是�,是�)( yk的,是j的)对于给定的提示$X$并且必须预测哪一个会被人类注释者评为更高。

这可以转化为以下损失函数：

使用 StackExchange 数据集，我们可以根据分数推断用户更喜欢这两个答案中的哪一个。有了这些信息和上面定义的损失，我们就可以transformers.Trainer通过添加自定义损失函数来修改。

我们利用 100,000 对候选人的一个子集，并评估 50,000 对保留的集合。在适度的训练批量大小为 4 的情况下，我们使用具有 BF16 精度的 Adam 优化器使用 LoRApeft适配器针对单个时期训练 LLaMA 模型。我们的 LoRA 配置是：

训练通过Weights & Biases记录，并使用 🤗 研究集群在 8-A100 GPU 上花费了几个小时，模型最终达到了67% 的准确率。虽然这听起来分数很低，但任务也非常艰巨，即使对于人工注释者也是如此。

如下一节所述，生成的适配器可以合并到冻结模型中并保存以供下游进一步使用。

从人类反馈中强化学习

有了经过微调的语言模型和奖励模型，我们现在可以运行 RL 循环了。大致分为三个步骤：

1. 根据提示生成响应
2. 使用奖励模型对响应进行评分
3. 使用评级运行强化学习策略优化步骤

查询和响应提示在被标记化并传递给模型之前按如下方式模板化：

SFT、RM 和 RLHF 阶段使用相同的模板。

使用 RL 训练语言模型的一个常见问题是，该模型可以通过生成完整的乱码来学习利用奖励模型，这会导致奖励模型分配高奖励。为了平衡这一点，我们在奖励中增加了一个惩罚：我们保留了一个我们没有训练的模型的参考，并通过计算 KL-divergence 将新模型的生成与参考模型进行比较：

再一次，我们利用peft内存高效训练，这在 RLHF 上下文中提供了额外的优势。在这里，参考模型和策略共享相同的基础，即 SFT 模型，我们以 8 位加载并在训练期间冻结。我们使用 PPO 专门优化策略的 LoRA 权重，同时共享基本模型的权重。

我们使用 🤗 研究集群在 3x8 A100-80GB GPU 上训练了 20 小时，但您也可以更快地获得不错的结果（例如，在 8 个 A100 GPU 上训练约 20 小时后）。训练运行的所有训练统计数据都可以在Weights & Biases上找到。

训练期间每一步的每批奖励。该模型的性能在大约 1000 步后达到稳定状态。

那么模型经过训练后能做什么呢？我们来看一下！

虽然我们不应该相信它关于 LLaMA 问题的建议，但答案看起来连贯，甚至提供了一个谷歌链接。让我们来看看接下来的一些训练挑战。

挑战、不稳定性和解决方法

用 RL 训练 LLM 并不总是一帆风顺的。我们今天演示的模型是许多实验、失败的运行和超参数扫描的结果。即使那样，该模型也远非完美。在这里，我们将分享我们在制作这个例子的过程中遇到的一些观察和头痛。

更高的奖励意味着更好的表现，对吧？

哇，这次跑步一定很棒，看看那甜美的奖励！

一般来说，在 RL 中，你希望获得最高的奖励。在 RLHF 中，我们使用了一个不完美的奖励模型，如果有机会，PPO 算法将利用这些不完美。这可能表现为奖励的突然增加，但是当我们从策略中查看文本生成时，它们主要包含字符串 ``` 的重复，因为奖励模型发现包含代码块的堆栈交换答案通常排名高于没有它的。幸运的是，这个问题很少被观察到，一般来说，KL 惩罚应该可以抵消这种攻击。

KL 总是一个正值，不是吗？

正如我们之前提到的，使用 KL 惩罚项是为了推动模型的输出保持接近基本策略的输出。通常，KL 散度衡量两个分布之间的距离，并且始终为正数。然而，在trl我们使用 KL 的估计值时，期望值等于实际 KL 散度。

显然，当从概率低于 SFT 模型的策略中抽取令牌时，这将导致负 KL 惩罚，但平均而言它将是正的，否则您将无法从策略中正确抽样。然而，一些生成策略可以强制生成一些令牌或者可以抑制一些令牌。例如，当批量生成完成的序列时，会填充序列，当设置最小长度时，EOS 令牌会被抑制。该模型可以为那些导致负 KL 的标记分配非常高或低的概率。当 PPO 算法针对奖励进行优化时，它会追逐这些负面惩罚，从而导致不稳定。

生成响应时需要小心，我们建议在求助于更复杂的生成方法之前始终先使用简单的采样策略。

持续的问题

还有一些问题需要我们更好地理解和解决。例如，损失偶尔会出现峰值，这可能会导致进一步的不稳定性。

当我们确定并解决这些问题时，我们将在上游进行更改trl，以确保社区能够受益。

结论

在这篇文章中，我们经历了 RLHF 的整个训练周期，从准备带有人工注释的数据集开始，使语言模型适应领域，训练奖励模型，最后用 RL 训练模型。

通过使用peft，任何人都可以在单个 GPU 上运行我们的示例！如果训练速度太慢，您可以在不更改代码的情况下使用数据并行性，并通过添加更多 GPU 来扩展训练。

对于真正的用例，这只是第一步！拥有经过训练的模型后，您必须对其进行评估并将其与其他模型进行比较，以了解它的好坏。这可以通过对不同模型版本的世代进行排名来完成，类似于我们构建奖励数据集的方式。

添加评估步骤后，乐趣就开始了：您可以开始迭代数据集和模型训练设置，看看是否有改进模型的方法。您可以将其他数据集添加到组合中，或者对现有数据集应用更好的过滤器。另一方面，您可以为奖励模型尝试不同的模型大小和架构，或者训练更长时间。

我们正在积极改进 TRL，使 RLHF 中涉及的所有步骤更容易获得，并很高兴看到人们用它构建的东西！如果您有兴趣做出贡献，请查看GitHub 上的问题。

引用