使用 RLHF 训练 LLaMA 的实践指南：StackLLaMA-Toy模板网

这篇具有很好参考价值的文章主要介绍了使用 RLHF 训练 LLaMA 的实践指南：StackLLaMA。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

由于LLaMA没有使用RLHF，后来有一个初创公司 Nebuly AI使用LangChain agent生成的数据集对LLaMA模型使用了RLHF进行学习，得到了ChatLLaMA模型，详情请参考：Meta开源的LLaMA性能真如论文所述吗？如果增加RLHF，效果会提升吗？，其实RLHF未必是必须的，主要是高质量的标注数据获取成本比较高，RLHF是一个trade-off。

StackLLaMA模型介绍

今天分享的StackLLaMA是按照InstructGPT论文的方法获得的，它的目的是，在算法流程上和ChatGPT类似，大致流程如下：

监督微调 (SFT)
奖励/偏好建模 (RM)
从人类反馈中强化学习 (RLHF)

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

主要区别在于：

基础模型：ChatGPT使用的是GPT3.5，StackLLaMA使用的是LLaMA；

SFT阶段：StackLLaMA使用的是StackExchange 数据集

（https://huggingface.co/datasets/HuggingFaceH4/stack-exchange-preferences），而ChatGPT的简单数据没有公开；

StackLLaMA的主要共贡献

- StackLLaMA模型开源了，并且在Huggingface Hub上可以使用，地址：https://huggingface.co/trl-lib/llama-7b-se-rl-peft；
- 集成到Hugging Face TRL库，为广大朋友提供了基础库使用，地址：https://huggingface.co/docs/trl/index；
- 开源了监督训练数据集StackExchange，地址：https://huggingface.co/datasets/HuggingFaceH4/stack-exchange-preferences；
- 开源了数据集和处理笔记本https://huggingface.co/datasets/lvwerra/stack-exchange-paired；
- 介绍了训练过程的细节以及解决方案；
- RLHF库：https://github.com/lvwerra/trl

Stack Exchange数据集

该数据集包括来自 StackExchange 平台的问题及其相应的答案（包括针对代码和许多其他主题的 StackOverflow）

根据论文《https://arxiv.org/abs/2112.00861》介绍，给每个答案一个score，公式如下：

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

对于奖励模型，始终需要每个问题有两个答案进行比较。对每个问题最多采样 10 个答案对，以限制每个问题的数据点数量。最后，将 HTML 转换为 Markdown 来清理格式，使模型的输出更具可读性。

高效的训练策略

即使训练最小的 LLaMA 模型也需要大量内存，比如bf16，每个参数需要2个字节来存储，fp32需要4个字节，更多可以参考：https://huggingface.co/docs/transformers/perf_train_gpu_one#optimizer。对于一个 7B 参数的模型，只是参数就需要(2+8)*7B=70GB的内存，实际存储还包括计算注意力分数等中间值，可能需要更多内存。因此，即使是7B的模型也无法在单个 80GB A100 上训练模型。

使用参数高效微调 (PEFT) 技术，比如使用https://github.com/huggingface/peft库来实现，这种技术可以加载8-bit模型执行低秩自适应 (LoRA)，如下图所示：

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

线性层的低秩自适应：在冻结层（蓝色）旁边添加额外参数（橙色），并将生成的编码隐藏状态与冻结层的隐藏状态一起添加。

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

在这种情况下，根据batch大小和序列长度不同每十亿个参数大概需要1.2-1.4GB内存，这可以以低成本微调更大的模型（在 NVIDIA A100 80GB 上高达 50-60B 比例模型）。

这些技术可以在消费级设备和 Google Colab 上微调大型模型。比如：facebook/opt-6.7b（13GB in float16）和openai/whisper-largeGoogle Colab（15GB GPU RAM）。更多参考：https://github.com/huggingface/peft和https://huggingface.co/blog/trl-peft

将非常大的模型放入单个 GPU 中，如果训练速度仍然很慢，可以使用数据并行技术，如下图所示：

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

数据并行可以使用transformers.Trainer和accelerate，无需任何代码更改，只需在使用torchrunor调用脚本时传递参数即可accelerate launch。下面分别用accelerate和在单台机器上运行带有 8 个 GPU 的训练脚本torchrun。

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

监督微调SFT

首先需要对监督数据做一些处理，传统方式通常是需要保证每个batch中序列长度是一样（采用填充或者截断），与传统的方式不同，GPT模型监督数据是把多个sentence通过EOS标记拼接到一起来使用，如下图所示：

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

使用peft加载模型之后就可以使用Trainer训练模型了。具体是：首先

首先导入int8模型，然后加入到训练准备，最后再添加LoRA adapters，代码如下所示：

# load model in 8bitmodel = 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 modellora_config = LoraConfig(    r=16,    lora_alpha=32,    lora_dropout=0.05,    bias="none",    task_type="CAUSAL_LM",)
model = get_peft_model(model, config)

Note：最终预测的时候需要把LoRA adapters的模型参与与LLaMA模型参数加起来使用。

通过运行脚本（https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py）将它们转换为 🤗 Transformers 格式。

奖励建模和人类偏好

奖励模型的目标是模仿人类如何评价文本。建立奖励模型有几种可能的策略：最直接的方法是预测注释（例如评分或“好”/“坏”的二进制值）。在实践中，更好的方法是预测两个例子的排名，奖励模型必须预测哪一个会被人类标注人员评分更高。

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

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

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

训练集100000个样本的子集，评估集50000个样本，batch大小为4，采用LoRA adapter（bf16混合精度的Adam优化器）fine-tuning LLaMA模型。LoRA 配置如下：

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

在 8个A100 GPU 训练了几个小时，训练的Weights & Biases记录地址：https://wandb.ai/krasul/huggingface/runs/wmd8rvq6?workspace=user-krasul

模型最终达到了67% 的准确率。虽然这听起来分数很低，但任务也非常艰巨，即使对于人工标注者也是如此。

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

从人类反馈中强化学习RLHF

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

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

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

Query and Response模板如下：

Question: <Query>
Answer: <Response>

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

使用 RL 训练语言模型的一个常见问题是，奖励模型为了得到高的reward，通常会生成一些完整的乱码序列。为了解决这个问题，在reward中增加了一个KL散度的惩罚，公式如下：

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

训练的时候，首先导入8-bit的SFT模型并且冻结参数，然后使用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 小时后）。训练过程每个step的reward变化如下图所示：

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

模型的性能在大约 1000 步后达到稳定状态。

StackLLaMA模型效果

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

答案看起来连贯，甚至提供了一个谷歌链接。让我们来看看接下来的一些训练挑战。

StackLLaMA模型训练的一些坑以及解决方案

更高的reward意味着更好的表现吗？

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

一般来说，在 RL 中希望获得最高的reward，但是在 RLHF 中，使用了一个不完美的奖励模型，PPO 算法将利用这些不完美，这可能表现为奖励的突然增加，但是当我们从策略中查看文本生成时，它们主要包含字符串 ``` 的重复，因为奖励模型发现包含代码块的stack exchange答案reward分数是最高的。幸运的是，这个问题很少被观察到，一般来说，KL 惩罚应该可以抵消这种攻击。

KL 总是一个正值，不是吗？

通常，KL 散度衡量两个分布之间的距离，并且始终为正数。然而，在trl库使用 KL 的估计值时，希望期望分布与实际分布尽可能接近。

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

显然，当从概率低于 SFT 模型的策略中采样token时，这将导致负 KL 惩罚，但平均而言它将是正的，否则您将无法从策略中正确抽样。然而，一些生成策略可以强制生成一些token或者可以抑制一些token。例如，当批量生成完成的序列时，会填充序列，当设置最小长度时，EOS 令牌会被抑制。该模型可以为那些导致负 KL 的token分配非常高或低的概率。当 PPO 算法针对奖励进行优化时，它会追逐这些负面惩罚，从而导致不稳定，如下图所示：

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