首页 > Ai资讯 > Ai知识库 > 详解大模型RLHF过程(配代码解读)

详解大模型RLHF过程(配代码解读)

发布时间:2024年06月06日

出自:https://zhuanlan.zhihu.com/p/624589622

一直都特别好奇大模型的强化学习微调是怎么做的,网上虽然相关文章不少,但找到的文章都是浅尝辄止说到用PPO训练,再细致深入的就没有讲了。。。只能自己看一看代码,以前搞过一点用PPO做游戏,感觉和语言模型PPO的用法不太一样。在游戏场景,每个step给环境一个action之后,agent拿到的state都是会变化的,通常也会设计奖励函数使得每个step都会有reward;但是在用强化学习微调语言模型这里,prompt是state,只输入一次,然后输出一串action(回答的单词),得到一个reward,模型并没有在每个action之后得到新的state(感谢评论区大佬的点拨,对于answer的第二个词,可以把prompt+answer的一个词当作新的state,而不只是把prompt当作state,状态转移蕴含在transformer内部)

本篇文章并不会介绍太多PPO的原理,相关文章已经很多了,比如李宏毅介绍PPO的课程。大模型里边的PPO涉及到了critic model的概念,在李宏毅教程里只提了一下并没有细讲,如果想了解可以看一下这个文章,相当于利用一个critic model预测从t时刻到最后一个时刻的累加奖励值(强化学习里边的第t个时刻对标answer句子里边的第t个单词),而不是通过实际累加得到从t时刻到最后一个时刻的累加奖励值,这样可以降低奖励的方差。下文也结合代码介绍critic model输出的具体含义。同时RLHF是什么也会再详细介绍,相关文章已经很多了。

本篇文章涉及的代码均来自微软的deepspeed对RLHF的实现,可配合huggingface官方的博客一起食用。本文只对算法的一些有特点的关键点进行阐述,并不对整体实现进行介绍。先上一张经典的论文图。本文重点结合代码讲解奖励模型训练和强化学习训练部分。

奖励(reward)模型训练

首先要声明的是,在强化学习阶段,用到的reward model和critic model都使用同一个模型初始化,因此在训练reward模型的过程中,也是在训练critic model。其次对符号进行说明,大模型中间隐藏层的参数维度为(B,L,D),B为batch size大小,L为句子长度,D为embedding维度。在接下来的代码讲解中,我也会标明代码中各个变量的维度,以更好的理解其意义。

在进行RLHF时,需要一个奖励模型来评估语言大模型(actor model)回答的是好是坏,这个奖励模型通常比被评估的语言大模型小一些(deepspeed的示例中,语言大模型66B,奖励模型只有350M)。奖励模型的输入是prompt+answer的形式,让模型学会对prompt+answer进行打分。奖励模型最后一层隐藏层的输出维度为(B,L,D),通过一个D✖️1的全连接层将维度变为(B, L),在L这个维度上,第i个位置的数据表示:从第i个位置到最后一个位置输出所能获得的奖励分值的累加和(和DQN里边的Q值一个意义),这种形式的输出满足了critic model的输出要求。对应代码如下:

#huggingface模型返回值是个list,第0位是模型最后输出的hideen state
hidden_states = transformer_outputs[0]
# v_head为Dx1的全连接网络对最后一维压缩
rewards = self.v_head(hidden_states).squeeze(-1)

对于一个奖励模型来说,目标是给一个句子进行打分,按理说每个句子对应一个分值就行了,但是目前对于长度为L的句子,奖励模型输出了L个值。我们用L维度上的最后一个位置的值当作为本句话的奖励得分。奖励模型训练优化采用pair wiss loss,即同时输入模型关于同一个问题的两个回答,让模型学会这两个句子哪个分高哪个分低。之所以如此训练是因为,在给奖励模型进行数据标注的过程中,给同一个问题的不同回答量化的打具体分值比较难,但是对他们进行排序相对简单,代码如下:

# 同一个batch里边的句子需要等长,短句后边会被padding
# [divergence_ind:end_ind]索引了padding前一个位置的输出分值
# chosen_reward是同一个句子pair里分数高的句子,r_truncated_reward是句子pair里分数低的句子
c_truncated_reward = chosen_reward[divergence_ind:end_ind]
r_truncated_reward = rejected_reward[divergence_ind:end_ind]

pair wise loss代码如下,如果给pair里边好的句子打分高(c_truncated_reward),坏的句子(r_truncated_reward)打分低,loss就会小:

loss += -torch.log(torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()

在训练强化学习的过程中,会用到reward model(critic model,再次提醒,critic model和reward model是同一个模型的两个副本)的推理过程,通过调用forward_value实现,具体代码如下,返回的值中有两种值,values表示每个位置i,从第i个位置到最后一个位置的奖励累加值,供强化学习过程中critic model使用;“chosen_end_scores”指的是对每个prompt+answer的打分,供reward model使用。

def forward_value(...):
    ...
    if return_value_only:
        #(B,L)
        return values
    else:
        ...
        return {
            "values": values,
            # (B,)
            "chosen_end_scores": torch.stack(chosen_end_scores),
        }

强化学习微调

强化学习微调阶段,会用到4个模型,actor model, ref_model,reward model和critic model(好费显存啊!!!)。其中actor model和ref_model是RLHF第一个阶段有监督微调模型的两个副本,reward model和critic model是本文第一部分训练出来的模型的两个副本。整体流程见这篇文档,整体流程图如下所示(没画出critic model):

首先说明actor model的训练模式推理模式的区别( 后边会用到)。训练模式是用teacher force的方式(不明白的同学知乎搜一下),将整句话输入到模型中,并通过mask机制在保证不泄漏未来的单词情况下预测下一个单词。推理模式是真正的自回归,预测出下一个单词之后,当作下一步输入再预测下下个单词,原理如下图所示:

首先用actor model在推理模式下根据prompt生成一个answer(prompt对应强化学习里边的state,answer对应一些列的action),代码如下:

# 保证不触发反向传播
with torch.no_grad():
    seq = self.actor_model.module.generate(prompts,
    max_length=max_min_length,
    min_length=max_min_length)

然后利用reward model和ciric model对输出的prompt+answer进行打分(PPO训练时使用的奖励值并不单单是reward model的输出还要考虑kl散度,后文介绍):

# 奖励模型返回的是个字典,key为chosen_end_scores位置存储数据维度为(B,),表示对于prompt+answer的打分
reward_score = self.reward_model.forward_value(
                seq, attention_mask,
                prompt_length=self.prompt_length)['chosen_end_scores'].detach(
                )
#critic model返回的数据维度为(B,L),L维度上第i个位置代表从i位置到最后的累积奖励
#舍去最后一个位置是因为句子“终止符”无意义 
values = self.critic_model.forward_value(
                seq, attention_mask, return_value_only=True).detach()[:, :-1]

actor model是我们想通过强化学习微调的大模型,但是强化学习过程很容易把模型训练“坏”,因此需要另外一个不会参数更新的 ref_model来当作标的,别让actor mode跑偏太远。我们在训练模式下,将prompt+answer分别输入到actor mode和ref model,用KL散度来衡量 ref model和actor mode输出的差别。同时将KL散度(衡量数据分布差距大小)纳入损失函数(KL散度本质是纳入到奖励值里边的,奖励值被纳入到了损失函数),进而来约束 ref_model和actor mode的输出分布别差距太大。具体代码如下:

# 得到两个模型的输出
output = self.actor_model(seq, attention_mask=attention_mask)
output_ref = self.ref_model(seq, attention_mask=attention_mask)
logits = output.logits
logits_ref = output_ref.logits
...
return {
...
# 分别得到两个模型在真实单词上的预测概率
'logprobs': gather_log_probs(logits[:, :-1, :], seq[:, 1:]),
'ref_logprobs': gather_log_probs(logits_ref[:, :-1, :], seq[:,1:]),
...
}
...
# 计算kl散度,log_probs里边存的数字经过log变化了,因此减法就对应除法
kl_divergence_estimate = -self.kl_ctl * (log_probs - ref_log_probs)

PPO训练时候的奖励值综合考虑KL散度和reward模型的输出,只考虑answer部分的KL散度,将reward model的输出加到KL散度L维度的最后一个位置上,得到最终的奖励值,代码如下:

rewards = kl_divergence_estimate
# 只考虑answer部分的奖励,不考虑prompt
start = prompts.shape[1] - 1
# 不考虑padding部分
ends = start + action_mask[:, start:].sum(1)
reward_clip = torch.clamp(reward_score, -self.clip_reward_value,
                         self.clip_reward_value)
batch_size = log_probs.shape[0]
# 在L维度上,每个位置都有KL散度,但是只在最后一个位置加上奖励值
for j in range(batch_size):
    rewards[j, start:ends[j]][-1] += reward_clip[j]

接下来的内容就是PPO的训练过程的比较核心的内容了,目标是计算PPO更新公示里边的advantage,具体公式如下,V就是critic model的输出。如果原理不懂建议先到这个链接看看。我直接在代码中给注释了。

图片出处:https://huggingface.co/blog/deep-rl-a2c
def get_advantages_and_returns(self, values, rewards, start):
    # values(B,L) critic model输出
    # rewards(B,)reward model输出
    # start answer开始的位置
    # Adopted from https://github.com/CarperAI/trlx/blob/main/trlx/models/modeling_ppo.py#L134
    lastgaelam = 0
    advantages_reversed = []
    length = rewards.size()[-1]
    # 计算每个时刻(序列位置)的critic model预测误差
    for t in reversed(range(start, length)):
        nextvalues = values[:, t + 1] if t < length - 1 else 0.0
        # critic model预测的是t到到最后一个时刻的奖励和,所以变化量delta可以用如下公式表示
        delta = (rewards[:, t] + self.gamma * nextvalues) - values[:, t]
        # self.gamma=1,self.lam=0.95是衰减因子,表示之前计算的delta对现在影响越来越小
        lastgaelam = delta + self.gamma * self.lam * lastgaelam
        advantages_reversed.append(lastgaelam)
    advantages = torch.stack(advantages_reversed[::-1], dim=1)
    # 后续用来更新critic model用
    returns = advantages + values[:, start:]
    return advantages.detach(), returns

以上过程,我们已经拿到了PPO训练所需要的advantage以及actor model的输出,我先现在可以对actor model进行训练啦。具体代码如下。logprobs和old_logprobs这两个参数分别是“老actor(n个epoch才会更新一次)”和新actor(每个batch都会更新它)”在正确单词上出处的概率,这块时PPO import sampling相关的知识,就不在这重复介绍了,不明白的同学补习一下哈。借用一下李宏毅老师的PPO公式:

def actor_loss_fn(self, logprobs, old_logprobs, advantages, mask):
    ## policy gradient loss
    #logprobs, old_logprobs都是经过log变化的单词概率,这里带着log做减法就相当于在做概率除法
    log_ratio = (logprobs - old_logprobs) * mask
    # 指数操作去掉log
    ratio = torch.exp(log_ratio)
    pg_loss1 = -advantages * ratio
    pg_loss2 = -advantages * torch.clamp(ratio, 1.0 - self.cliprange,
                                            1.0 + self.cliprange)
    pg_loss = torch.sum(torch.max(pg_loss1, pg_loss2) * mask) / mask.sum()
    return pg_loss

同样的,我们也要对critic model进行训练,更新,loss就是mse loss。

def critic_loss_fn(self, values, old_values, returns, mask):
    ## value loss
    # 用“老critic model”的输出约束“新critic model”不要步子太大,裁剪一下
    values_clipped = torch.clamp(
        values,
        old_values - self.cliprange_value,
        old_values + self.cliprange_value,
    )
    vf_loss1 = (values - returns)**2
    vf_loss2 = (values_clipped - returns)**2
    vf_loss = 0.5 * torch.sum(
        torch.max(vf_loss1, vf_loss2) * mask) / mask.sum()
    return vf_loss

至此,我们的RLHF训练流程就结束了。第二部分开头我们说过,共涉及actor model, ref_model,reward model和critic model这四个模型,其实更新参数的模型只有actor model和critic model。

如果你想要了解关于智能工具类的内容,可以查看 智汇宝库,这是一个提供智能工具的网站。
在这你可以找到各种智能工具的相关信息,了解智能工具的用法以及最新动态。