将 TRL 用于 LLaMA 模型

我们已经开始推出在 trl 中使用 Meta LLaMA 模型的示例（原始 LLaMA 模型请参阅 Meta 的 LLaMA 发布）。

高效的训练策略

即使是训练最小的 LLaMA 模型也需要巨大的内存。简单计算一下：在 bf16 精度下，每个参数使用 2 字节（在 fp32 中为 4 字节），此外，像 Adam 优化器还会额外使用 8 字节（更多信息请参阅 Transformers 中的性能文档）。因此，一个 7B 参数的模型仅仅为了加载到内存就需要 (2+8)*7B=70GB，并且在计算注意力分数等中间值时可能需要更多内存。所以，你甚至无法在一块 80GB 的 A100 上进行这样的训练。你可以使用一些技巧，比如更高效的优化器或半精度训练，来挤占更多内存，但迟早会耗尽。

另一种选择是使用参数高效微调（PEFT）技术，例如 peft 库，它可以对以 8 位加载的模型执行低秩适应（LoRA）。更多关于 peft + trl 的信息，请参阅 Peft 集成文档。

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

在这种情况下，一个经验法则是为每十亿参数分配约 1.2-1.4GB 内存（取决于批量大小和序列长度），以容纳整个微调设置。这使得在低成本的情况下微调更大的模型（在 NVIDIA A100 80GB 上可达 50-60B 规模的模型）成为可能。

现在我们可以将非常大的模型装入单个 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

监督式微调

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

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

通过这种方法，训练效率更高，因为每个通过模型的标记都参与了训练，而填充标记通常会从损失计算中被掩码掉。如果你的数据不多，更担心偶尔会截断一些超出上下文的标记，你也可以使用经典的数据加载器。

# 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_kbit_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, config)

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

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

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

奖励建模与人类偏好

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

一个效果很好的技巧是，在 RL 循环开始前，在收集到的人工标注数据上训练一个奖励模型，而不是直接使用反馈。奖励模型的目标是模仿人类如何评价一段文本。建立奖励模型有几种可能的策略：最直接的方法是预测标注（例如，评分或“好”/“坏”的二元值）。实际上，更好的做法是预测两个示例的排名，其中奖励模型会看到给定提示 x 的两个候选答案 (y_k, 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 适配器，以 BF16 精度和 Adam 优化器训练 Llama 模型一个周期。我们的 LoRA 配置如下：

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

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

从人类反馈中进行强化学习

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

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

查询和响应提示在被分词并传递给模型之前，会按照以下模板进行格式化：

Question: <Query>

Answer: <Response>

在 SFT、RM 和 RLHF 阶段都使用了相同的模板。我们再次利用 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 to 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)

有关其余细节和评估，请参阅我们关于 StackLLaMA 的博客文章。

< > 在 GitHub 上更新