跳到正文
原文
Hugging Face Blog·· 2023-04-05精选AI 评分61

StackLLaMA：用 RLHF 训练 LLaMA 的实操指南

StackLLaMA: A hands-on guide to train LLaMA with RLHF

AI 导读

Hugging Face 发布 StackLLaMA 模型及完整训练流程，用 SFT、奖励模型和 RLHF 三步训练 LLaMA 在 Stack Exchange 上答题，并开源了基于 TRL 库的全部训练管线。

推荐理由

原文完整披露了 RLHF 三步训练流程和踩坑细节，读者可据此复现或在单卡上跑起自己的对齐实验。

正文 · AI 翻译

像 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 发布），并且整个训练流程作为 HF Mirror 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 中描述的方法，为每个答案分配一个分数：

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).

对于奖励模型，我们总是需要每个问题有两个答案进行比较，稍后我们会看到。有些问题有几十个答案，因此会产生许多可能的答案对。我们每个问题最多采样十个答案对，以限制每个问题的数据点数量。最后，我们通过将 HTML 转换为 Markdown 来清理格式，使模型的输出更易读。您可以在这里找到数据集以及处理 notebook。

高效训练策略

即使训练最小的 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（float16 中为 13GB）和 openai/whisper-large（15GB GPU RAM）。要了解有关使用 peft 的更多信息，请参阅我们的 github 仓库 或之前的 博客文章（https://hf-mirror.com/blog/trl-peft），了解如何在消费级硬件上训练 20b 参数的模型。

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

chapter10_ddp.png

我们使用 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

监督微调

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

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

chapter10_preprocessing-clm.png

使用这种方法，训练效率要高得多，因为通过模型的每个 token 都参与了训练，而填充 token 通常会从损失中被掩蔽。如果你没有太多数据，更关心偶尔截断一些溢出上下文的 token，你也可以使用经典的数据加载器。

打包由 ConstantLengthDataset 处理，然后我们在用 Trainer 加载模型后使用 peft。首先，我们以 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 格式。请注意，你还需要从源码安装 🤗 Transformers，直到 v4.28 发布。

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

奖励建模与人类偏好

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

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

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

loss⁡(θ)=−E(x,yj,yk)∼D[log⁡(σ(rθ(x,yj)−rθ(x,yk)))] \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 r 是模型的评分，yj 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

我们使用 100,000 对候选样本中的一个子集，并在一个包含 50,000 对样本的留出集上进行评估。以 4 这一适中的训练批量大小，我们使用 LoRA peft 适配器、采用 Adam 优化器和 BF16 精度对 LLaMA 模型训练一个 epoch。我们的 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 循环了。它大致遵循三个步骤：

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

Untitled

在分词并传入模型之前，Query 和 Response 提示按如下方式套用模板：

Question: <Query>

Answer: <Response>

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

用 RL 训练语言模型的一个常见问题是，模型可能学会通过生成完全乱码来利用奖励模型，从而使奖励模型给出高奖励。为了平衡这一点，我们在奖励中加入了惩罚项：我们保留一个不训练的模型作为参考，并通过计算 KL 散度将新模型的生成结果与参考模型进行比较：

R⁡(x,y)=r⁡(x,y)−βKL⁡(x,y) \operatorname{R}(x, y)=\operatorname{r}(x, y)- \beta \operatorname{KL}(x, y)

其中 r r 是来自奖励模型的奖励，而 KL⁡(x,y) \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)

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

Per batch reward at each step during training. The model’s performance plateaus after around 1000 steps. 训练期间每一步的每批次奖励。模型性能在大约 1000 步后趋于平稳。

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

llama prompt

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

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

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

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

Wow this run must be great, look at that sweet, sweet, reward! 哇，这次运行一定很棒，看看那甜美的奖励！

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

KL总是正值，不是吗？

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

KLpen(x,y)=log⁡(πϕRL(y∣x)/πSFT(y∣x)) KL_{pen}(x,y) = \log \left(\pi_\phi^{\mathrm{RL}}(y \mid x) / \pi^{\mathrm{SFT}}(y \mid x)\right)

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

Negative KL

在生成响应时需要小心，我们建议先使用简单的采样策略，然后再考虑更复杂的生成方法。

持续存在的问题

仍有许多问题需要我们更深入地理解和解决。例如，偶尔会出现损失突增，这可能导致进一步的不稳定。

Loss spikes

随着我们识别并解决这些问题，我们会将更改trl提交到上游，确保社区能够受益。

结论

在这篇文章中，我们详细介绍了RLHF的整个训练周期，从准备带有人工标注的数据集开始，到使语言模型适应特定领域，训练奖励模型，最后使用强化学习训练模型。

通过使用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://hf-mirror.com/blog/stackllama },
    doi          = { 10.57967/hf/0513 },
    publisher    = { HF Mirror Blog }
}

致谢

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

来源：Hugging Face Blog · hf-mirror.com

#教程/实践#开源/仓库#数据/训练#Hugging Face#Meta