【AI大模型面试真题】倒腾了一波RLHF，把我整蒙圈了...

本文链接：https://blog.youkuaiyun.com/Android23333/article/details/145132937

前言

最近两月倒腾了一波 RLHF，从 ColossalAI 到 TRLX 以及 DeepSpeed-Chat，最后基于 DeepSpeed-Chat 成功训练上了自己的模型，最后效果也是肉眼可见的提升。对这一部分进行下总结，包括原理，代码以及踩坑与解决方案。

1、基本概念

首先还是解释一下一些概念，从 NLP 的角度举一些例子。

首先是 RL 中的 Policy，State，Action：

Policy π：这是我们需要学习的策略函数，我们使用语言模型来近似这个函数。
State S：模型接受的输入句子，这里的话状态是无限的，因为输入可能为任意句子。
Action A：根据所处的状态，策略函数做出动作。这里的话就是模型接受上文，预测下一个字符。这里的动作是有限集合，大小为 vocab size。

RL 概念图

接下来介绍 Reward，Return，Q，V：

PS：这里要注意区分价值和奖励：价值是未来累计奖励的期望。奖励是我们做出该动作后立即获取的收益。

2、RLHF过程

整个过程主要是分为三步：SFT，Training Reward Model，RLHF。

这里主要介绍一下 Training Reward Model 和 RLHF。

Step2：Training Reward Model

（1）数据

首先是这一部分数据的构成，每一条数据由 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个专门机构（如教科文组织）都没有设在总部。然而，有一些“自治附属机构”（如联合国儿童基金会）的总部设在联合国总部。"
}

（2）模型结构

Reward Model 相较于原始的 SFT Model，在后面加上了一个 value head，value head 是一个 Linear，输入维度为模型的 hidden_dim，输出维度为 1，输出表示模型预测每一字符获取的得分。

DeepSpeed-Chat 中使用最后一个字符的得分作为整个 response 的得分（当然也可以使用整个句子中每个字符的平均分作为整体的得分）。

Reward Model

（3）训练目标

训练 Reward Model 是一个排序任务，针对 query，输入 chosen 和 rejected response，训练目标尽可能的使得 chosen 和 rejected 的差值更大，损失函数为：

以上就是第二步 Training Reward Model 的全部过程，基于 rank loss 训练了一个打分模型。

在第三步强化学习中，reward 模型将扮演环境的角色，针对模型预测的字符给出奖励分数。

Step3：RLHF

（1）整体结构

首先来从整体上看一下这部分（这里就只介绍 RL 部分，PTX 就是加上了预训练任务）：

DeepSpeed-Chat RLHF

RLHF 基于 A2C 方法，这一步包含了四个模型：

Actor Model：由 SFT 之后的模型初始化而来。作为策略（policy）模型，用于接收上文，做出动作，预测下一个字符。学习完毕之后，我们最终使用的就是这个模型。

Reference Model：和 Actor Model 同样初始化自 SFT Model，训练过程中冻结参数，用于和 Actor Model 做对比，保证模型不要偏离原始 SFT Model 太多。

Reward Model：作为环境（env），训练过程中冻结参数，针对每一个状态，给出奖励分数。

Critic Model：由 Reward Model 初始化而来，用于近似价值函数，输入为状态 s，估计当前状态的价值 V。

（2）训练过程

接下来梳理一遍训练过程。训练过程整体分为两步：maker experience 和 learn。

首先是 make_experience，首先在训练数据中抽取一部分 query，然后 Actor Model 生成答案。

然后我们依据这条答案获取我们所需要的经验：

actor_logits：由 Actor Model 产生，包含对答案所有词的概率分布。
reference_logits：由 Reference Model 产生，包含对答案所有词语的概率分布，用于和 actor logits 进行对比，防止 actor model 偏离 SFT Model 太远。
reward_score：由 Reward Model 产生，为当前句子状态下，立即获取的收益分数。
values：由 Critic Model 产生，估计当前句子状态下，到完成生成，可以获取的回报。

整体流程如下：

make experience

然后在 learn 的时候，通过所产生的经验进行学习。我们通过 Actor Model 与 Critic Model 近似策略函数和价值函数，整体流程如下：

learn

关于 learn 这部分，详细介绍下 Critic Model 训练和 Actor Model 训练过程。

这里对下标 new 和 old 做一下说明（最近自己复习一下，看自己的文章，想了半天也没想起，最后被迫又去看代码了，可恶）

问：这里获得的 actor prob 和之前 make experience 中获得的有什么区别，为什么要区分 old 和 new，包括 critic 也是？

答：因为模型会有 dropout 所以会有一些区别。此外如果 ppo_epoch 大于 1，那么 actor 产生经验之后，在 forward，backward 之后，也会有不同（https://github.com/microsoft/DeepSpeedExamples/issues/533）

（3）Critic Loss

Critic Model 估计当前状态可以获取的价值，也就是我们前面所说的 V 值。模型的输入为状态 s，也就是当前模型生成的句子，输出为状态价值 V（s）。

这里举一个例子：

早上起床去公司，地图预计通行时间是 50min，然后走了 10 分钟到马连洼，地图预计用时还需要三十分钟。

那么一开始预估的总体时间是 50min，而基于一部分真实观测的总体时间是 10+30=40min，那么我们最终的误差就是(10+30 - 50) **2啦。

critic loss example

（4）Actor Loss

举个例子，如果我们模型有“上，下，左，右”四个动作，分别有累计奖励“10，20，30，40”，我们做出任意动作，都会获取正向的累计奖励，因此模型也会向这个动作更新参数。

而实际上，我们累计奖励的平均期望为 25，对于动作“上，下”我们都应该受到惩罚。

最后，我们在通过 PPO 算法来进一步优化。关于 PPO 详细讲解，可以看我的文章 PPO 详解。

PPO 的想法是希望限制训练过程中对 policy 进行的更改，避免 policy 进行过大的更新，来提升 policy 训练的稳定性。

基于 PPO 算法，我们的目标函数变为：

以上就是第三步的核心内容，RL 过程整体分为两步，make experience 和 learn。

我们首先采样数据，然后生成结果，Reward Model 给出环境的奖励，以及 Critic Model 对结果进行评判。之后我们依据所获取的经验来对模型进行更新。

3、DeepSpeed-Chat实践与踩坑

Step2 Training Reward Model

这个步骤基本就是全部按照 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

Step3 RLHF

（1）踩坑&解决方案

在这个步骤中，从跑通到收敛还是有不少麻烦，分享一些比较重要的点：

训练过程中发现 make experience 的时候，model.generate() 产生的答案全是重复的无意义字。

这里就很奇怪了，最后发现是开启了 DeepSpeed Hybrid Engine 所导致的，查看了 issue 之后，也发现了有类似的问题。

不过目前 DeepSpeed-Chat 也没有解决，需要关闭 Hybrid Engine 进行训练。

DeepSpeed-Chat 还有一个很严重的问题就是，在 make experience 的时候，强制 Actor Model 生成到最大长度（设置max_length=min_length=max_min_length），这样子导致模型生成偏差很大。

对于一个简单的问题，模型可能本来生成简单的一句话就可以完美回答了，但是却必须强制生成到最大长度，这样训练的模型和我们实际用起来的模型是有区别的。

对于这个问题，可以通过修改 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 了。

（2）Trick

trick 方面主要参考了 The 37 Implementation Details of Proximal Policy Optimization 以及影响 PPO 算法性能的 10 个关键技巧（附 PPO 算法简洁 Pytorch 实现）。对于部分 trick，进行了尝试。

Normalization of Advantages：

将 Advantage 进行归一化：adv=(adv-mean)/std
在 mini-batch 上进行
没有指标的量化，从我个人看结果而言，感觉提升不大

Overall Loss and Entropy Bonus：

为了提高算法的探索能力，在 actor 的 loss 中增加一项策略熵，并乘以一个系数 entropy_coef，使得在优化 actor_loss 的同时，让策略的熵尽可能大。一般我们设置 entropy_coef=0.01
总体损失变为：loss = policy_loss - entropy * entropy_coefficient + value_loss * value_coefficient
没有指标的量化，从我个人看结果而言，感觉没有太多提升

Reward Scale：