使用 TRL 训练 LLaMA 模型

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

高效的训练策略

即使训练最小的 LLaMA 模型也需要大量的内存。一些简单的计算：在 bf16 中，每个参数使用 2 个字节（在 fp32 中为 4 个字节），另外还使用 8 个字节，例如在 Adam 优化器中（有关更多信息，请参阅 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

监督微调

在开始训练奖励模型和使用强化学习调整模型之前，如果模型在我们感兴趣的领域已经表现良好，这将很有帮助。在我们的案例中，我们希望它能够回答问题，而对于其他用例，我们可能希望它能够遵循指令，在这种情况下，指令微调是一个好主意。实现这一点最简单的方法是继续使用语言建模目标对来自该领域或任务的文本训练语言模型。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

我们利用 10 万对候选者的子集，并在 5 万个保留集中进行评估。使用 4 的适度训练批次大小，我们使用 Adam 优化器和 BF16 精度对 Llama 模型进行单轮训练，使用 LoRA peft 适配器。我们的 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 的博客文章。