本文详细解析了大模型 RLHF 的完整流程，涵盖奖励模型训练与强化学习微调两个核心阶段。文章基于微软 DeepSpeed 代码实现，阐述了 Reward Model 与 Critic Model 的初始化关系及 Pairwise Loss 训练方法。重点讲解了 PPO 算法中 Actor、Ref、Reward、Critic 四个模型的协作机制，包括 KL 散度约束、Advantage 计算及策略梯度损失函数的具体实现。通过代码解读，展示了如何利用强化学习对齐大模型输出，使其更符合人类偏好。
一直对大模型的强化学习微调（RLHF）过程感到好奇，网上相关文章虽多，但往往浅尝辄止。本文结合微软 DeepSpeed 的 RLHF 实现代码，深入讲解奖励模型训练和强化学习微调部分。
在语言模型中使用强化学习微调时，Prompt 对应状态（State），输出的一系列单词对应动作（Action）。与游戏场景不同，这里通常只输入一次 Prompt，得到一串回答后获得一个 Reward。不过，对于回答的第二个词，可以将 Prompt+ 第一个词作为新的 State，状态转移蕴含在 Transformer 内部。
大模型中的 PPO 涉及 Critic Model 的概念，用于预测从 t 时刻到最后一个时刻的累加奖励值，而非通过实际累加得到，这样可以降低奖励方差。RLHF 主要包含两个阶段：有监督微调（SFT）和强化学习微调。
在强化学习阶段，使用的 Reward Model 和 Critic Model 通常使用同一个模型初始化。Reward Model 的输入是 prompt+answer 的形式，让模型学会对 prompt+answer 进行打分。
Reward Model 最后一层隐藏层的输出维度为 (B, L, D)，通过一个 D×1 的全连接层将维度变为 (B, L)。L 维度上第 i 个位置的数据表示从第 i 个位置到最后一个位置输出所能获得的奖励分值的累加和。这种形式的输出满足了 Critic Model 的输出要求。
# huggingface 模型返回值是个 list，第 0 位是模型最后输出的 hidden state
hidden_states = transformer_outputs[0]
# v_head 为 Dx1 的全连接网络对最后一维压缩
rewards = self.v_head(hidden_states).squeeze(-1)
对于一个 Reward Model 来说，目标是给一个句子进行打分。目前对于长度为 L 的句子，Reward Model 输出了 L 个值。我们用 L 维度上的最后一个位置的值作为本句话的奖励得分。
Reward Model 训练优化采用 Pairwise Loss，即同时输入模型关于同一个问题的两个回答，让模型学会这两个句子哪个分高哪个分低。这是因为在给 Reward Model 进行数据标注的过程中，给同一个问题的不同回答量化具体分值比较难，但是对他们进行排序相对简单。
# 同一个 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]
Pairwise Loss 代码如下，如果给 pair 里边好的句子打分高，坏的句子打分低，loss 就会小：
loss += -torch.log(torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()
在训练强化学习的过程中，会用到 Reward Model（Critic Model）的推理过程，通过调用 forward_value 实现。返回的值中有两种值：values 表示每个位置 i，从第 i 个位置到最后一个位置的奖励累加值；chosen_end_scores 指的是对每个 prompt+answer 的打分。
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 是本文第一部分训练出来的模型的两个副本。
首先说明 Actor Model 的训练模式和推理模式的区别。训练模式是用 Teacher Force 的方式，将整句话输入到模型中，并通过 mask 机制在保证不泄漏未来的单词情况下预测下一个单词。推理模式是真正的自回归，预测出下一个单词之后，当作下一步输入再预测下下个单词。
首先用 Actor Model 在推理模式下根据 Prompt 生成一个 Answer，代码如下：
# 保证不触发反向传播
with torch.no_grad():
seq = self.actor_model.module.generate(prompts,
max_length=max_min_length,
min_length=max_min_length)
然后利用 Reward Model 和 Critic 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 Model 跑偏太远。我们在训练模式下，将 prompt+answer 分别输入到 Actor Model 和 Ref Model，用 KL 散度来衡量 Ref Model 和 Actor Model 输出的差别。同时将 KL 散度纳入损失函数，进而来约束 Ref_Model 和 Actor Model 的输出分布别差距太大。
# 得到两个模型的输出
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_logprobs)
PPO 训练时候的奖励值综合考虑 KL 散度和 Reward Model 的输出，只考虑 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 的输出。
def get_advantages_and_returns(self, values, rewards, start):
# values（B，L）critic model 输出
# rewards（B，）reward model 输出
# start answer 开始的位置
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 和新 Actor 在正确单词上出现的概率，这块是 PPO Importance Sampling 相关的知识。
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。通过这种方式，大模型能够在保持原有知识的同时，更好地遵循人类指令，提升回答质量与安全性。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online