作者:ILikeShuhuaMilk
最近两月倒腾了一波RLHF,从ColossalAI到TRLX以及DeepSpeed-Chat,最后基于DeepSpeed-Chat成功训练上了自己的模型,最后效果也是肉眼可见的提升。对这一部分进行下总结,包括原理,代码以及踩坑与解决方案。
基本概念
首先还是解释一下一些概念,从NLP的角度举一些例子。
首先是RL中的Policy,State,Action。
RL概念图
接下来介绍Reward,Return,Q,V。
PS:这里要注意区分价值和奖励:价值是未来累计奖励的期望。奖励是我们做出该动作后立即获取的收益。
RLHF过程
整个过程主要是分为三步:SFT,Training Reward Model,RLHF。
这里主要介绍一下Training Reward Model和RLHF。
首先是这一部分数据的构成,每一条数据由query,chosen response, rejected response构成,chosen response相较于rejected response质量更高。我使用的数据来自Cohere/miracl-zh-queries-22-12,里面对于每条query包含了positive_passages和negative_passages两个部分,positive部分可以视为chosen,negative部分视为rejected,就可以采样构建我们训练Reward Model的数据了:
{
"query": "联合国总部在哪里?",
"chosen": "联合国总部大楼(亦称联合国大厦)是联合国总部的所在地,位于美国纽约市曼哈顿东侧,属于国际领土,因此只要是会员国国民持有护照就可以进入,包括与美国无邦交的联合国会员国。联合国总部大楼位于纽约市,其西侧边界为第一大道、南侧边界为东42街、北侧边界为东48街、东侧边界为东河,从联合国总部大楼可以俯瞰东河。此大楼于1949年和1950年间兴建,土地购自于当时的纽约房地产家,面积阔达17英亩(约6.87973公顷)。在此之前,洛克斐勒家族有意提供其在纽约州威斯特彻斯特郡洛克菲勒庄园的土地,但因距离曼哈顿遥远而作罢;之后,纳尔逊·洛克菲勒便协助新的土地的购买,其父小约翰·戴维森·洛克菲勒则捐助了850万美元协助兴建大楼。",
"rejected": "联合国的15个专门机构(如教科文组织)都没有设在总部。然而,有一些“自治附属机构”(如联合国儿童基金会)的总部设在联合国总部。"
}
Reward Model相较于原始的SFT Model,在后面加上了一个value head,value head是一个Linear,输入维度为模型的hidden_dim,输出维度为1,输出表示模型预测每一字符获取的得分。DeepSpeed-Chat中使用最后一个字符的得分作为整个response的得分(当然也可以使用整个句子中每个字符的平均分作为整体的得分)。
Reward Model
训练目标
训练Reward Model是一个排序任务,针对query,输入chosen和rejected response,训练目标尽可能的使得chosen和rejected的差值更大,损失函数为:
以上就是第二步Training Reward Model的全部过程,基于rank loss训练了一个打分模型。在第三步强化学习中,reward模型将扮演环境的角色,针对模型预测的字符给出奖励分数。
首先来从整体上看一下这部分(这里就只介绍RL部分,PTX就是加上了预训练任务):
DeepSpeed-Chat RLHF
RLHF基于A2C方法,这一步包含了四个模型:
训练过程
接下来梳理一遍训练过程。训练过程整体分为两步:maker experience和learn。
首先是make_experience,首先在训练数据中抽取一部分query,然后Actor Model生成答案。然后我们依据这条答案获取我们所需要的经验:
整体流程如下:
make experience
然后在learn的时候,通过所产生的经验进行学习。我们通过Actor Model与Critic Model近似策略函数和价值函数,整体流程如下:
learn
关于learn这部分,详细介绍下Critic Model训练和Actor Model训练过程。
Critic Model估计当前状态可以获取的价值,也就是我们前面所说的V值。模型的输入为状态s,也就是当前模型生成的句子,输出为状态价值V(s)。
举个例子,如果我们模型有「上,下,左,右」四个动作,分别有累计奖励「10,20,30,40」,我们做出任意动作,都会获取正向的累计奖励,因此模型也会向这个动作更新参数。而实际上,我们累计奖励的平均期望为25,对于动作「上,下」我们都应该受到惩罚。
以上就是第三步的核心内容,RL过程整体分为两步,make experience和learn。我们首先采样数据,然后生成结果,Reward Model给出环境的奖励,以及Critic Model对结果进行评判。之后我们依据所获取的经验来对模型进行更新。
DeepSpeed-Chat实践与踩坑
这个步骤基本就是全部按照DeepSpeed-Chat代码来了,使用cohere-zh的数据,大约构造了2w条chosen-rejected pair用于训练。最后训练训练了一个epoch,在验证集上准确率到了0.79左右。
***** Evaluating reward, Epoch 1/1 *****
step: 499/2287, chosen_last_scores (higher is better) : 5.074454307556152,reject_last_scores (lower is better) : 0.5599770545959473, acc (higher is better) : 0.812000036239624
step: 999/2287, chosen_last_scores (higher is better) : 5.084388732910156,reject_last_scores (lower is better) : 0.7938708662986755, acc (higher is better) : 0.7940000295639038
step: 1499/2287, chosen_last_scores (higher is better) : 5.106724262237549,reject_last_scores (lower is better) : 0.7971451878547668, acc (higher is better) : 0.7986666560173035
step: 1999/2287, chosen_last_scores (higher is better) : 5.0183587074279785,reject_last_scores (lower is better) : 0.672178328037262, acc (higher is better) : 0.7955000400543213
chosen_last_scores (higher is better) : 5.028912544250488,reject_last_scores (lower is better) : 0.7077188491821289, acc (higher is better) : 0.7936161160469055
踩坑&解决方案
在这个步骤中,从跑通到收敛还是有不少麻烦,分享一些比较重要的点:
不过目前DeepSpeed-Chat也没有解决,需要关闭Hybrid Engine进行训练。
对于这个问题,可以通过修改generate中的参数eos_token_id来解决。设置一个虚假的结束符,然后模型就可能生成正常的结束符。然后我们在构建attention_mask来遮蔽掉结束符后面的回答。例如: seq = [0,0,0,0, prompt, answer, eos_token, other_word] mask = [0,0,0,0,1(prompt),1(answer),1(eos_token),0(other_word)]
通过以上两步基本就可以跑通流程了,但是训练过程中还遇到了一个比较大的问题,就是Critic Loss并不收敛,越来越大。
具体的原因是:训练后期,随着模型输出答案越来越好,我们的reward值也会越来越高,导致我们最终累计回报return的区间也会越来越大。而前面说过,我们Critic Loss是一个MSE损失,因此训练后期,随着return的估计范围越来越大,Critic Model就越难估计。在我训练的过程中,一开始return的范围是在3-4左右,训练后期涨到了18-20,因此我们需要想点办法来约束一下我们的return。
首先的超参的调节,γ参数为折扣回报率,DeepSpeed-Chat中初始设置为1,可以将其调小一些来缓解。其次的话,使用reward scale这一trick帮助也非常大。
通过以上这些步骤,基本上我们正常训练了,模型最后也能看到一些效果,但是需要取得更好的效果,我们就需要引入一些trick了。
trick方面主要参考了The 37 Implementation Details of Proximal Policy Optimization以及影响PPO算法性能的10个关键技巧(附PPO算法简洁Pytorch实现)。对于部分trick,进行了尝试。
关于其他的一些trick,如学习率衰减、梯度裁剪、Value Clipping等,本身框架就包含了,就不进行特别说明了。当然,这些trick在不同的数据或者模型上都会有不同的效果,需要自己进行探索。
通过以上这些方法,最后也顺利的完成了训练,希望对大家能有所帮助,也希望大家能分享自己有用的经验与理解,关于RL自己了解的也很少,还需要更多的学习。
参考资料:
https://zhuanlan.zhihu.com/p/635569455