使用 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 的更多信息，请参阅文档。

以 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 之前微调模型没有什么特别之处 - 这只是我们在此应用的预训练中的因果语言建模目标。为了有效地使用数据，我们使用了一种称为打包的技术：我们不是在批次中为每个样本设置一个文本，然后填充到最长文本或模型的最大上下文中，而是连接许多文本，并在其间使用 EOS 标记，并剪切上下文大小的块以填充批次，而无需任何填充。

通过这种方法，训练效率更高，因为通过模型传递的每个标记也都经过训练，这与通常从损失中屏蔽的填充标记形成对比。如果您没有太多数据，并且更担心偶尔会截断一些溢出上下文的标记，您也可以使用经典的数据加载器。

打包由 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_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 格式。请注意，您还需要从源代码安装 🤗 Transformers，直到 v4.28 发布。

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

奖励建模和人类偏好

原则上，我们可以直接使用人类注释通过 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

我们利用 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,
)

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

基于人类反馈的强化学习

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

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

在被标记化并传递给模型之前，查询和回复提示的模板如下

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 的博客文章。