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

发布于2023年4月5日

在 GitHub 上更新

刘易斯·滕斯托尔 (Lewis Tunstall)

ChatGPT、GPT-4和Claude等模型是强大的语言模型，它们通过“人类反馈强化学习”（RLHF）方法进行了微调，以便更好地符合我们对它们的行为预期和使用方式。

在这篇博客文章中，我们将展示使用RLHF结合以下方法训练LlaMa模型以回答Stack Exchange问题的全部步骤：

监督式微调（SFT）
奖励/偏好建模（RM）
来自人类反馈的强化学习（RLHF）

摘自InstructGPT论文：Ouyang, Long, et al. "Training language models to follow instructions with human feedback." arXiv preprint arXiv:2203.02155 (2022)。

通过结合这些方法，我们发布了StackLLaMA模型。该模型可在🤗 Hub上获取（原始LLaMA模型请参见Meta的LLaMA发布），并且整个训练管道作为Hugging Face TRL库的一部分提供。为了让您了解该模型的能力，请尝试下面的演示！

LLaMA模型

进行RLHF时，从一个有能力的模型开始非常重要：RLHF步骤仅仅是一个微调步骤，旨在使模型与我们期望的交互方式和响应方式保持一致。因此，我们选择使用最近推出且性能优异的LLaMA模型。LLaMA模型是Meta AI开发的最新大型语言模型。它们的大小范围从7B到65B参数，并在1T到1.4T个token上进行了训练，这使得它们非常强大。我们使用7B模型作为所有后续步骤的基础！要访问该模型，请使用Meta AI的表单。

Stack Exchange数据集

收集人类反馈是一项复杂且昂贵的任务。为了在此示例中引导该过程，同时仍构建一个有用的模型，我们使用了StackExchange数据集。该数据集包含来自StackExchange平台（包括StackOverflow的代码和许多其他主题）的问题及其对应的答案。它对这种用例很有吸引力，因为答案附带了点赞数和已接受答案的标签。

我们遵循Askell et al. 2021中描述的方法，并为每个答案分配一个分数：

分数 = log2 (1 + 点赞数) 四舍五入到最近的整数，如果提问者接受了答案则加1（如果点赞数为负数，我们将其分数设为-1）。

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

高效训练策略

即使是训练最小的LLaMA模型也需要巨大的内存。简单计算一下：在bf16中，每个参数使用2字节（fp32中是4字节），此外Adam优化器还使用了8字节（更多信息请参阅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上微调大型模型成为可能。值得注意的演示是在Google Colab上微调facebook/opt-6.7b（13GB，float16）和openai/whisper-large（15GB GPU RAM）。要了解更多关于使用peft的信息，请参阅我们的GitHub仓库或关于在消费硬件上训练200亿参数模型的上一篇博客文章（https://huggingface.co/blog/trl-peft）。

现在我们可以将非常大的模型放入单个GPU中，但训练可能仍然非常缓慢。在这种情况下，最简单的策略是数据并行：我们将相同的训练设置复制到独立的GPU中，并向每个GPU传递不同的批次。通过这种方式，您可以并行化模型的前向/后向传播，并随着GPU数量的增加而扩展。

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

accelerate launch --multi_gpu --num_machines 1  --num_processes 8 my_accelerate_script.py
torchrun --nnodes 1  --nproc_per_node 8 my_torch_script.py

监督式微调

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

在进行RLHF之前对模型进行微调没有什么特别之处——我们只是在这里应用了预训练中的因果语言建模目标。为了高效地利用数据，我们使用了一种称为“打包”（packing）的技术：在批处理中，我们不是每个样本一个文本然后填充到最长文本或模型最大上下文，而是将许多文本用EOS标记连接起来，并切割上下文大小的块来填充批处理，无需任何填充。

通过这种方法，训练效率大大提高，因为模型中传递的每个token都经过训练，而填充token通常会从损失中屏蔽掉。如果您数据不多，并且更担心偶尔会截断一些超出上下文的token，您也可以使用传统的DataLoader。

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

# load model in 8bit
model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        load_in_8bit=True,
        device_map={"": Accelerator().local_process_index}
    )
model = prepare_model_for_int8_training(model)

# add LoRA to model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)

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

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

现在我们已经针对任务对模型进行了微调，我们可以训练一个奖励模型了。

奖励建模和人类偏好

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

一种替代直接反馈的有效技巧是在RL循环之前，根据人类标注训练一个奖励模型。奖励模型的目标是模仿人类如何评价文本。构建奖励模型有几种可能的策略：最直接的方法是预测标注（例如，评分或“好”/“坏”的二元值）。实际上，更好的方法是预测两个示例的排序，即奖励模型接收给定提示 $x$ 的两个候选 $(y_k, y_j)$ ，并预测人类标注者会认为哪个更好。

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

$\operatorname{loss}(\theta)=- E_{\left(x, y_j, y_k\right) \sim D}\left[\log \left(\sigma\left(r_\theta\left(x, y_j\right)-r_\theta\left(x, y_k\right)\right)\right)\right]$

其中 $r$ 是模型的得分， $y_j$ 是首选候选。

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

class RewardTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"],  attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss

我们利用了10万对候选子集，并在5万个保留集上进行评估。在批处理大小为4的情况下，我们使用LoRA peft适配器对LLaMA模型进行单次epoch训练，使用Adam优化器和BF16精度。我们的LoRA配置是：

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

训练通过Weights & Biases进行日志记录，在8块A100 GPU上使用🤗研究集群花费了几个小时，模型最终达到了67%的准确率。虽然这听起来分数不高，但这项任务本身也非常困难，即使对于人类标注者来说也是如此。

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

来自人类反馈的强化学习

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

从提示中生成响应
使用奖励模型对响应进行评分
使用评分运行强化学习策略优化步骤

查询和响应提示在分词并传递给模型之前，模板化如下：

Question: <Query>

Answer: <Response>

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

使用强化学习训练语言模型的一个常见问题是，模型可能通过生成完全无意义的内容来利用奖励模型，这会导致奖励模型分配高奖励。为了平衡这一点，我们对奖励进行惩罚：我们保留一个不进行训练的模型作为参考，通过计算KL散度来比较新模型生成的内容和参考模型生成的内容。

$\operatorname{R}(x, y)=\operatorname{r}(x, y)- \beta \operatorname{KL}(x, y)$

其中 $r$ 是来自奖励模型的奖励， $\operatorname{KL}(x,y)$ 是当前策略与参考模型之间的KL散度。

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

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    question_tensors = batch["input_ids"]
        
    # sample from the policy and generate responses
    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Compute sentiment score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    # Log stats to WandB
    ppo_trainer.log_stats(stats, batch, rewards)

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

训练过程中每批次的奖励。模型性能在大约1000步后趋于平稳。

那么，训练后的模型能做什么呢？我们来看看！

虽然我们目前不应该完全相信它在LLaMA方面的建议，但这个答案看起来很连贯，甚至提供了一个谷歌链接。接下来，我们来看看一些训练挑战。

挑战、不稳定性及变通方法

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

更高的奖励意味着更好的性能，对吗？

哇，这次运行一定很棒，看看这甜美的奖励！

通常在强化学习中，你希望获得最高的奖励。在RLHF中，我们使用一个不完美的奖励模型，而PPO算法一旦有机会就会利用这些不完美之处。这可能表现为奖励的突然增加，然而当我们查看策略生成的文本时，它们大多包含“```”字符串的重复，因为奖励模型发现包含代码块的Stack Exchange答案通常比没有代码块的答案排名更高。幸运的是，这个问题很少出现，而且通常KL惩罚应该可以抵消这种利用行为。

KL值总是正数，不是吗？

正如我们之前提到的，KL惩罚项用于使模型的输出保持接近基础策略。一般来说，KL散度衡量两个分布之间的距离，并且总是一个正值。然而，在trl中，我们使用KL的估计值，该估计值在期望上等于真实的KL散度。

$KL_{pen}(x,y) = \log \left(\pi_\phi^{\mathrm{RL}}(y \mid x) / \pi^{\mathrm{SFT}}(y \mid x)\right)$

显然，当从策略中采样到的token的概率低于SFT模型时，这将导致负的KL惩罚，但平均而言，它将是正的，否则您将无法正确地从策略中采样。然而，一些生成策略可以强制生成某些token，或者某些token可以被抑制。例如，在批量生成时，完成的序列会被填充，并且在设置最小长度时，EOS token会被抑制。模型可以为这些token分配非常高或低的概率，这会导致负KL值。由于PPO算法优化奖励，它会追逐这些负惩罚，从而导致不稳定性。

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

持续存在的问题

仍有许多问题需要我们更好地理解和解决。例如，损失有时会出现尖峰，这可能导致进一步的不稳定性。

随着我们识别和解决这些问题，我们将把更改上游到trl，以确保社区能够从中受益。

结论

在这篇文章中，我们介绍了RLHF的整个训练周期，从准备带有“人类标注”的数据集开始，接着将语言模型适应到特定领域，训练奖励模型，最后使用RL训练模型。

通过使用peft，任何人都可以使用单个GPU运行我们的示例！如果训练太慢，您可以使用数据并行，无需更改代码，只需添加更多GPU即可扩展训练。

对于实际应用，这只是第一步！一旦您训练好一个模型，您必须对其进行评估，并将其与其他模型进行比较，以了解其性能如何。这可以通过对不同模型版本的生成结果进行排名来完成，类似于我们构建奖励数据集的方式。

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

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

引用

@misc {beeching2023stackllama,
    author       = { Edward Beeching and
                     Younes Belkada and
                     Kashif Rasul and
                     Lewis Tunstall and
                     Leandro von Werra and
                     Nazneen Rajani and
                     Nathan Lambert
                   },
    title        = { StackLLaMA: An RL Fine-tuned LLaMA Model for Stack Exchange Question and Answering },
    year         = 2023,
    url          = { https://huggingface.co/blog/stackllama },
    doi          = { 10.57967/hf/0513 },
    publisher    = { Hugging Face Blog }
}

致谢

我们感谢Philipp Schmid分享了他精彩的流式文本生成演示，我们的演示就是基于此。我们还要感谢Omar Sanseviero和Louis Castricato对博客文章草稿提供了宝贵而详细的反馈。

更多博客文章

使用直接偏好优化方法对 LLM 进行偏好调优

由 2024年1月18日 • 69

使用 DPO 微调 Llama 2

由 2023年8月8日 • 60

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录发表评论