社区供稿 | 基于 LoRA 的 RLHF: 记一次不太成功但有趣的百川大模型调教经历

• 项目地址：https://github.com/beyondguo/LLM-Tuning

众所周知，整个 RLHF (基于人类反馈的强化学习) 分为这么三步：

• SFT (Supervised Fine-Tuning): 有监督的微调，使用正常的 instruction following 或者对话的样本，来训练模型的基础对话、听从 prompt 的能力；
• RM (Reward Modeling): 基于人类的偏好和标注，来训练一个能模拟人偏好的打分模型；
• RL (Reinforcement Learning): 在前面的 SFT 模型的基础上，借助 RM 提供反馈，来不断通过 PPO 的强化学习框架来调整模型的行为。

为了节省训练资源，快速了解整个 RLHF 的过程，我这里每一步的训练，都采用 LoRA 微调的方式：使用 LoRA 进行 SFT，使用 LoRA 训练 Reward Model，以及使用 LoRA 来进行强化学习 PPO 过程。

下面使用 baichuan-7B 作为基座模型，来实现整个 RLHF 过程。为什么选择 baichuan-7B 呢（而不是 ChatGLM 等模型呢）？因为 baichuan-7B 是一个纯纯的基座模型，本身没有对话能力，因此很适合检验我们训练的效果到底好不好；另一方面，这是一个很强大的基座模型，尤其在中文上，因此“调教”的潜力很大，相比 BLOOM 等模型可能更能训练出效果。

注：我这里的训练，并不是为了得到一个综合性能多么好的 Chat 模型（这些事儿留给专业机构团队来做），而是走通整个 RLHF 的过程，了解其中可能的技术难点，感受强化学习和有监督学习的区别，从而获得大模型“调教”的一手经验。

首先，我要感谢 Huggingface 团队开源的 Stack-LLaMA（https://huggingface.co/blog/stackllama）以及相关教程，我这里整个步骤，都是参考这个教程进行的，然后针对 baichuan 做了定制化的适配，并全部改成基于 LoRA 来训练。

好，不废话了，咱们开始吧！

1. SFT：训练一个拥有基础对话能力的模型

这里，请大家直接回顾一下使用 HC3 数据集来让 baichuan-7B 有对话能力 这个教程，因为我这里会直接复用这里的所有训练过程和最终的模型。

这里的 SFT 数据，我使用的是 HC3(Human-ChatGPT Comparison Corpus) 数据集：

• HC3 项目：https://github.com/Hello-SimpleAI/chatgpt-comparison-detection

基于 LoRA 训练之后的大概效果是这样的：

• 输入: "你是谁开发的啊"
• 原始 baichuan-7B: "我就是你，我是你自己。(自性)"
• ChatBaichun-HC3: "我是一个计算机程序，由一个人或一群人编写。我的目的是帮助人们解决问题和回答问题。"

理清关系：

• 我们称 baichuan-7B 为 base-model ；
• 训练时，我们需要加载一个 language model head，即成为 base-model-ForCausalLM ；
• 经过 SFT 的训练之后，我们得到对应的 LoRA 模型记为 sft-lora-adapter ；
• 当我们要使用 SFT 的模型时，就是把 sft-lora-adapter 加载到 base-model-ForCausalLM 上，加载之后，我们记为 sft-model。

2. RM：训练一个奖励模型

奖励模型，就是一个打分模型，能够判断模型对于同一个 prompt 的不同输出，哪个更好，哪个更差。

具体地，我们需要一批人类标注的对不同回答的排序数据，然后基于这样的排序，构造得分，或者更简单一点——构造标签，然后训练一个 regression 模型来打分。

中文的开源高质量排序数据几乎没有，而我又没法请一批人来真正打标一批数据，所以我决定直接使用一个现成的英文 reward 数据集，然后利用 ChatGPT 都翻译成中文，来将就用一用：

• 原始英文 reward 数据集：https://huggingface.co/datasets/yitingxie/rlhf-reward-datasets
• 翻译成中文的 reward 数据集：https://huggingface.co/datasets/beyond/rlhf-reward-single-round-trans_chinese

这个数据集中，每个prompt都对应一个"chosen"字段和一个"rejected"字段，分别代表一个更好的回答和一个更差的回答。

大致思路：

要训练打分模型，我们可以直接使用 sigmoid 二分类的方式，给 base-model 添加一个 num_labels=1 的 classification head，即 base-model-ForSequenceClassification。这相当于添加了一个 output-dim=1 的 linear 层，模型的输出就一个一维的 logits，这个 logits 经过 sigmoid 函数标准化之后就可以作为分数了。模型的输入实际上是 prompt+answer。

然后，我们的 chosen 回答和 rejected 回答都是成对输入的，我们希望前者得到的分数比后者的分数更高，二者的差别越大越好。因此，我们可以设计一个特殊的损失函数：

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

其中 j 代表的是包含 chosen 回答的输入，k 代表的是包含 rejected 回答的输入，模型分别对二者计算，得到各自的 logits，二者的差值，经过logsigmoid处理，得到损失函数。用这个损失，来对模型进行优化。

当然，由于我们设计了特殊的双输入输出结构，我们还需要做一些其他的修改，比如自定义 compute_metrics 函数来监控训练过程中的准确率，设计特殊的 data collator 来满足这种数据格式，这些内容都在 reward_modeling.py 中实现并进行了详细地注释。

下面，我们可以使用这样的命令，来训练 reward model：

# train_rm.sh
CUDA_VISIBLE_DEVICES=0,1,2,3 python reward_modeling.py \
    --model_name baichuan-inc/baichuan-7B \
    --lora_target_models W_pack \
    --num_train_epochs 2 \
    --eval_steps 200 \
    --save_steps 50 \
    --per_device_train_batch_size 10 \
    --per_device_eval_batch_size 16 \
    --train_subset -1 \
    --eval_subset -1 \
    --max_length 1000

下面是训练过程中的 wandb 记录：

100%|███████████████████████████████████████████████████████████| 3972/3972 [7:06:28<00:00,  6.44s/it]

wandb: 
wandb: Run history:
wandb:                  eval/accuracy ▁▄▅▆▆▆▆▇▇▇▇▇▇▇▇▇▇███████████████████████
wandb:                      eval/loss █▆▅▄▄▃▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁
wandb:                    train/epoch ▁▁▁▂▂▂▂▂▂▃▃▃▃▃▄▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb:              train/global_step ▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
wandb:            train/learning_rate ███▇▇▇▇▇▇▆▆▆▆▆▆▅▅▅▅▅▄▄▄▄▄▄▃▃▃▃▃▂▂▂▂▂▂▁▁▁
wandb:                     train/loss █▇▆▃▅▇▃▅▅▆▆▄▅▆▃▅▃▂▄▂▂▃▄▃▃▅▁▂▄▅▃▅▆▅▃▄▄▅▃▄

wandb: Run summary:
wandb:                  eval/accuracy 0.71892
wandb:                      eval/loss 0.54986
wandb:                   eval/runtime 281.5312
wandb:        eval/samples_per_second 17.742
wandb:          eval/steps_per_second 1.112
wandb:                    train/epoch 2.0
wandb:              train/global_step 3972
wandb:            train/learning_rate 0.0
wandb:                     train/loss 0.5684
wandb:               train/total_flos 0.0
wandb:               train/train_loss 0.56024

可见最终我们的 reward model 在验证集上的准确率达到了 71%，属于一个还不错的结果了。

理清关系：

• 依然使用 baichuan-7B 作为 base-model ；
• 训练时，我们需要加载一个 1-label classificaition model head，即成为 base-model-ForSequenceClassification ；
• 经过 RM 的训练之后，我们得到对应的 LoRA 模型记为 reward-lora-adapter ；
• 当我们要使用 RM 的模型时，就是把 reward-lora-adapter 加载到 base-model-ForSequenceClassification 上，加载之后，我们记为 reward-model。

3. RL：基于 PPO 的强化学习过程

好了，现在咱们已经拥有了一个 SFT 版本的模型和一个 RM 模型，是时候让 AI helps AI 了！（RM helps SFT model）

大概思路：首先看看在强化学习这里，我们需要哪些模型：

• sft-model，它已经具备了基本的 chat 能力，现在我们希望借助 PPO 方法来进一步提高它，我们记在 PPO 过程中不断提高的这个模型为 ppo-model；
• reward-model，它可以帮我们对 ppo-model 的不同输出进行打分，从而对 PPO 过程提供指导；
• reference model ref-model，它是一个正则化的作用，希望经过优化后的 ppo-model 不要跟 ref-model 差别太大。我们可以直接使用 sft-model 作为 refl-model。

那 ppo-model 模型结构上跟 sft-model 一模一样吗？也不是，还需要一个小小的修改——添加一个 value head, 通过 value head，模型生成的每个 token 都对应一个 value，这代表当前每个 token 的某个分值；

然后，我们还要计算每个 token 理想的分值，这就要用到前面的 reward-model 给出的一个 score，以及 ref-model 跟 ppo-model 对比时，产生的 kl 损失。具体地，ppo-model 生成的句子会对应一个 score，这个可以理解为“生成这样的句子，应得到多少的奖励”；另外，每个 token 都会有一个 ppo-model 对应的概率分布和 ref-model 对应的概率分布，这二者可以计算得到一个损失，这个可以理解为“生成这样的token，要接受多少的惩罚”，这个惩罚，就是对前面奖励的一个正则项。把 score 加到 kl损失序列的最后一个有效token的位置上，就构成了正则之后的 reward。

有了这个每个 token 的 reward，就可以拿来跟带有 value head 的 ppo-model 得到的每个 token 的 value 进行对比然后计算 loss 了，然后就基于这个 loss 来对模型参数进行优化。

关于这里面详细的计算细节，建议读者详细阅读：trl 包中 ppo_trainer.py 中的 compute_rewards, train_minibatch 等函数，应该有更清楚的理解。

上面解释了，为什么我们需要加 value head，为什么需要 ref model，以及这些不同模块之间是如何交互的。

最后，我们再次理清一下关系：

• 在 base-model 上加载 sft-lora-adapter 得到 sft-model；
• 把 sft-model 拷贝一份，作为 ref-model；
• 给 sft-model 添加一个 value head，得到 ppo-model，这也是我们要强化学习训练的对象；
• 强化学习阶段，我们复用之前的 sft-lora-adapter，就是依然只训练 adapter 那部分的参数，相当于对 sft 的 lora adapter 进一步优化。

这些关系很容易搞乱，因此我画了这么个图，来帮助理解：

其中，ppo-model 跟 sft-model/base-model 的关系是这样的：

这些实现，都在 rl_training.py 代码中有详细注释，读者可自行查阅。我目前的实现方式跟 Stack-LLaMA 教程中有一个明显不同，Stack-LLaMA 中训练 SFT，RM 之后，会把 lora 跟 base-model 融合成一个模型，这样每个阶段都要保存一个大模型，很占空间；我的实现则是希望不要进行合并，本地只用保存一个 base-model 即可，然后全部通过 lora 来实现各种训练和计算。但存在一个小问题，目前 trl 框架中对 ref-model 不支持 lora 等 peft 模型，所以我在这一步进行了合并。

我们通过下面的命令来开启训练：

CUDA_VISIBLE_DEVICES=0,1,2 python rl_training.py \
    --base_model_name baichuan-inc/baichuan-7B \
    --merged_sft_model_path [your_merged_sft_model_path] \
    --sft_model_lora_path [your_sft_lora_path] \
    --reward_model_lora_path [your_reward_lora_path] \
    --adafactor False \
    --save_freq 10 \
    --output_max_length 256 \
    --batch_size 8 \
    --gradient_accumulation_steps 16 \
    --batched_gen True \
    --ppo_epochs 4 \
    --seed 0 \
    --learning_rate 1e-5 \
    --early_stopping True \
    --output_dir [your_output_ppo_lora_path] \
    --log_with wandb

下图展示了训练过程中的 reward 的均值的变化：

（重点）训练遇到的问题。。。

上面那些都是理论、流程上的东西，看似顺理成章，但真训练起来发现问题很大。

比如我前期一直被这个 bug 所困扰：RuntimeError: probability tensor contains either inf, nan or element < 0

这个 bug 指向的是模型在经过一定次数的 ppo 迭代之后的 generate 过程，之前一直查不出是啥情况，以为是 baichuan 模型有啥特殊的地方，都找不到。

直到某一次我详细观察模型每一步训练时的输出时，发现 reward-model 给予了一个回答为空的问答对很高的 score，我才意识到，可能是 reward-model 的问题，导致一些，空回答被给予了高分，导致模型迭代后不输出了，从而报上面那个错误。

对此，我尝试了下面几种补救方式：

• 降低学习率。这种方式确实让训练更稳定了一点，但时间一长还是崩了；
• 调整 kl 惩罚的大小，让惩罚来的更大一些。可惜效果不明显；
• 调整 top-k/top-p 等参数，没啥用...

最终，我发现根本原因还是因为 reward model 不够好，对某些极端错误场景给予了高分，导致模型快速学坏，虽然单纯看 reward model 的训练曲线和准确率感觉还彳亍，但是由于本身使用的数据质量、数量、覆盖面都不够，没有考虑很多极端情况，比如大量重复、回复为空、回复很短等，这些情况没有覆盖到的话，reward model 的鲁棒性就很差，很容易在你意想不到的地方突然来一个 Suprise Motherfxxker!

而目前我训练的这个 reward model，容易让输出为空的时候产生高 score，所以在不重新训练 reward model 的情况下，我想了一个绝招：禁止模型回答为空：在 generate 参数中，把 EOS token 加入到 begin_suppress_tokens 中，这样模型就必须得回答了。

果真，加了这个 trick 之后，模型终于不崩了，上面展示的 reward 曲线，也是加了这个 trick 之后的结果。

看曲线的话，一切似乎都很完美，但不要忘记了，我们的 reward model 本来就是有偏的，reward 高不代表模型就真的好。比如，我们看平均回复长度：

模型到后期还是止不住地想变得越来越不想说话。。。你禁止说话是吧，那我就直说几个字敷衍一下。

所以，如果想从根本上去让模型从 ppo 过程中得到更好的结果，reward model 的质量一定要高。这个就留给之后有空再尝试了。

效果展示

虽然上面的训练过程由于 reward model 的问题并不完美，但 ppo 的初期，还是有一些肉眼可见的效果的——模型风格的剧烈变化。

首先这里补充一嘴，在前面我使用 ChatGPT 对一个英文 reward 数据集进行中文的 prompt 中，我特意要求翻译地要“口语化、活泼一点”。当时我只是为了降低“翻译腔”，现在回过头来看模型的训练效果，发现这个居然为后面的模型风格转变埋下了伏笔。

下面是一些例子展示：

我这里对比四个版本模型的输出：

• SFT：就是前面的经过 HC3 数据集 SFT 之后的模型；
• RLHF (10)：就是在 SFT 模型的基础上，经过了 10 个 batch 的 PPO 迭代的模型；
• RLHF (16)：同上，经过了 16 次迭代；
• RLHF (100)：同上，经过 100 次迭代。

而这里面，我最推崇的就是 RLHF (16) 这个版本，因为它有明显的风格转换，但是还没有学坏！

总之，可以发现 RLHF (16) 这个版本的模型，在正常回答你问题的时候，变得非常活泼，明显的网友既视感，而这一点，就是 reward model 和 PPO 过程带来的。对比 SFT 模型，它是由 ChatGPT 数据蒸馏而来，因此一股浓浓的 ChatGPT 味儿，我们这里的 reward model 虽然不好，但也成功地改变了原本干巴巴的风格。

总结

总的来说，这个 RLHF 的过程并不成功，但是却让我很有收获。没有真正跑跑 RLHF 之前，对 reward model，PPO 这些没有概念，也不知道这里面的难点在哪里。下面说说我对这个过程的明显感觉：

• RLHF 中的 PPO 训练确实十分不稳定，对 reward model 异常地敏感，甚至根据我观察的结果：如果出现了一个极端样本（比如回答是空的），且 reward model 给出了极高的分（比如达到0.9），即使就这么一个，也可以让模型学坏，直接走向不归路；
• PPO 的优化非常高效，在我的上面训练中，原本的 SFT 模型是浓厚的 ChatGPT 风味，但仅仅经过了 16 个 batch 的 PPO 微调，模型就变成了微博/小红书风味；
• 基于以上两点，我觉得 reward model 直接决定了 RLHF 的效果，因此花时间好好做一个 reward model 十分重要，值得比 SFT 花更多的时间精力来做。

主要参考：

• Stack-LLaMA：https://huggingface.co/blog/stackllama
• trl 包 中的 Stack-LLaMA 例子：https://github.com/lvwerra/trl/tree/main/examples/stack_llama/scripts
• 一些跟trl作者的讨论：
• 关于 reward model：https://github.com/lvwerra/trl/issues/492
• 关于 peft model：https://github.com/lvwerra/trl/issues/536