RLHF微调大模型---PPO原理和代码实战

一、前言

RLHF: Reinforment Learning Human Feedback。中文含义是： 基于人类反馈的强化学习，用奖励模型Reward Model来训练SFT模型；生成模型使用奖励或惩罚来更新其策略，以便生成更高质量、更符合人类偏好的文本。

为什么需要RLHF，SFT 不够吗？

• 数据层面：SFT的目的是预测值与标签token级别完全一致，模型效果依赖于标注数据的质量，而且标注成本相对较高。SFT只有正反馈，没有负反馈机制，模型只知道下一个token是什么是正确的，而不知道什么是错误的。RLHF则通过直接与人类互动进行学习，不需要依赖大量的标注数据，尤其在处理一些特殊任务时（如多轮对话、创意生成等），这种方式可以更加灵活且高效。
• 模型安全与理论方面：RLHF可以帮助确保模型的行为符合道德标准和社会规范。例如，在处理有潜在危害的内容生成时，RLHF通过反馈机制帮助模型避免生成有害或不恰当的内容。

可以跳过SFT阶段直接进行RLHF么？

• 探索空间巨大： 预训练模型虽然拥有强大的语言能力，但它并不知道如何有效地遵循指令。如果没有 SFT 的初步引导，RLHF
  需要在一个巨大的空间中进行探索，这会导致训练非常缓慢且不稳定。模型可能需要很长时间才能找到一个合理的策略，甚至可能无法收敛。
• 奖励信号稀疏： RLHF 的奖励信号通常比较稀疏，只有在生成了完整的回复后才能获得奖励。如果没有 SFT的引导，模型很难在早期生成有意义的回复，从而难以获得有效的奖励信号。这会导致训练非常困难。
• 训练成本高昂： 由于探索空间巨大和奖励信号稀疏，直接进行 RLHF 需要大量的计算资源和时间。这使得训练成本非常高昂，甚至可能无法完成。

所以我们需要RLHF，而且是在SFT的基础上进行的RLHF

大模型的训练需要经历以下阶段：

• 预训练阶段：PT（Pre training）。使用公开数据经过预训练得到预训练模型，预训练模型具备语言的初步理解；训练周期比较长；
• 微调阶段1：SFT（指令微调/有监督微调）。如果想要预训练模型在某个垂直领域（金融、法律、电商等）有更好的知识储备，就需要使用人工标注的QA问答对进行有监督的微调训练，从而得到精调模型；训练周期较短；
• 微调阶段2：对齐/强化训练。精调模型的输出并不是全部都令人满意的，我们还需要让模型知道回复的接受度。可以在运行日志中收集对齐数据，包含【问题，接受的回复，不接受的回复】，再进行对齐训练，得到最后可使用的模型；

二、RLHF原理

RLHF常用的方式为PPO算法，这里我们不过多讲PPO的原理，而是多讲解如何使用PPO来进行RLHF，PPO算法的原理可以参考如下：
PPO算法原理

使用PPO进行RLHF需要的模型如下：

• Actor Model：用于生成句子的模型，也就是需要训练的模型。
• Critic Model：指导你进步的教练模型，注意，这个教练模型也会随着你的进步来调整自己的指导策略，主要是评判大模型每一个动作，即输出每一个token的好坏。
• Reward Model：用于给出最终分数的模型。虽然教练能够给你一定的指导，但最终游戏获胜与否还是要靠裁判说了算，可以说教练在教你的同时也在尝试学习裁判的偏好。裁判一般是固定的，因此 Reward Model 在整个训练过程中参数是被冻结的，Critic Model和Reward Model是同一个模型的两个副本,只不过一个是要用每一个token的评分，一个是要用整个句子的评分。
• Reference Model：这是 PPO 在 LLM 中独有的概念，目的是为了让 actor 不要训练偏离太远，主要是缓解 reward hacking + 稳定训练使用的，Reference Model是Actor Model的副本，不参与训练。

2.1、利用Reward Model

符号说明：大模型中间隐藏层的参数维度为(B,L,D)，B为batch size大小，L为句子长度，D为embedding维度。

在进行RLHF时，需要一个奖励模型来评估语言大模型（actor model）回答的是好是坏，这个奖励模型通常比被评估的语言大模型小一些。

奖励模型的输入是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]
loss += -torch.log(torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()

loss的目的就是为了c_truncated_reward更大，r_truncated_reward更小

在训练强化学习的过程中，会用到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),
        }

2.2、利用Actor Model

首先用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是我们想通过强化学习微调的大模型，但是强化学习过程很容易把模型训练“坏”，因此需要另外一个不会参数更新的Reference Model来当作标的，别让Actor Mode跑偏太远。我们在训练模式下，将prompt+answer分别输入到Actor Mode和Reference Model，用KL散度来衡量 Reference Model和Actor Mode输出的差别。同时将KL散度（衡量数据分布差距大小）纳入损失函数（KL散度本质是纳入到奖励值里边的，奖励值被纳入到了损失函数），进而来约束 Reference 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]

2.3、优势函数

接下来是计算PPO更新公示里边的advantage，具体公式如下:
V就是Critic Model的输出。
$A_{\theta }(s,a)=Q_{\theta }(s,a)-V_{\theta }(s_t)$

$Q_{\theta }(s,a)$表示在状态s下做出的行为a期望的回报。
$Q_{\theta }(s,a)=r_t+\gamma \ast V_{\theta }(s_{t+1})$
则上面的优势函数可以改写为：
A θ ( s , a ) = r t + γ ∗ V θ ( s t + 1 ) − V θ ( s t