该博文内容摘自:https://blog.csdn.net/chaishen10000/article/details/131232948
https://github.com/microsoft/DeepSpeedExamples/tree/master/applications/DeepSpeed-Chat
一、RLHF微调三阶段
参考:https://huggingface.co/blog/rlhf
1)使用监督数据微调语言模型,和fine-tuning一致。
2)训练奖励模型
奖励模型是输入一个文本序列,模型给出符合人类偏好的奖励数值,这个奖励数值对于后面的强化学习训练非常重要。构建奖励模型的训练数据一般是同一个数据用不同的语言模型生成结果,然后人工打分。如果是训练自己领域的RLHF模型,也可以尝试用chatgpt打分,效果也不错。
3)训练RL模型
在训练强化学习模型时,需要搞清楚状态空间、动作空间、策略函数、价值函数这些东西,动作空间就是所有的token,状态空间就是输入的序列的分布,价值函数由第二步的奖励模型和策略约束结合,策略函数就是微调的大模型。
从上图可以看出,给定一个输入x,会生成两个文本y11和y22,一个来自于初始的模型,另一个来自于微调的模型,微调的模型生成的文本还会进入到奖励模型中打分输出rθ,而初始模型和微调的模型生成的结果会用KL散度约束它们的分布,确保模型不会太偏离原来的模型,并且能输出高质量的回复。
值得注意的是三个阶段的训练数据尽量是分布一致的,否则后面的训练会很不稳定。所以在第一步微调时不要一味地使用大量的训练数据(这一步的数据比较容易获得),尽量和后面两步的数据分布保持一致。
二 、训练奖励模型的解析
其他步骤理解起来很简单,重点是要理解如何训练奖励模型。
代码在training/step2_reward_model_finetuning文件夹下,奖励模型可以选择一个较小的模型,如opt-350M,在chosen和rejected这种样本对上训练。奖励模型的代码实现在training/utils/model/reward_model.py中。reward model的输出类似于回归任务,将大模型的输出,然后经过N ✖️ 1 的线性层,得到一个batch size ✖️ seq len ✖️ 1的输出。在训练过程中,使用到的loss是二元交叉熵,确保每个prompt 的 chosen分数都是要大于rejected。
loss += -torch.log(torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()
上面的代码中c_truncated_reward 和 r_truncated_reward 即给定一个prompt,对应的chosen和rejected获得的分数,而且是chosen 和 rejected所有token的分数差值。注意在这里因为chosen和rejected的长度不一致,而且还有padding的部分,所以c_truncated_reward和r_truncated_reward要做阶段,主要是截取chosen_id和rejected_id不等的部分出来,去除共同padding的部分。