解析 RLHF 微调三阶段

国内现在有不少表现亮眼的基础语言模型，但它们本质上还是个“自回归生成器”——接收到一段文本，就按部就班地往下接续，很多时候根本抓不住人类真正想表达什么。想要让模型真正“对齐”人类意图，即理解指令背后的期望，目前业界公认比较有效的路径就是基于人类反馈的强化学习（RLHF）。

InstructGPT 那篇经典论文——Training language models to follow instructions with human feedback——把 RLHF 拆成了三个环节：

。下面咱们就借着论文和 TRL 代码的视角，看看这三个阶段各自是怎么落地的。

prompt 收集

为什么非要把 prompt 收集单独拎出来讲？因为整个 InstructGPT 的数据集构建，都是建立在这批 prompt 之上的。再说一个容易被忽略的点：这个 prompt 集是张动态表，不是一次性写完就完事的。

OpenAI 先让标注工手动写了一批 prompt，涵盖对话、开放问答、分类、信息抽取、文本生成、头脑风暴等任务。接着用这批 prompt 和对应的标注答案微调出一个 Demo 版的 InstructGPT，然后开放 Playground 给真实用户使用。用户在与模型交互的过程中产生的 prompt，反过来又被收集起来补充到原始 prompt 集里。这才是“人在回路中”的真正含义——用户输入的那些指令，比工人们“编”出来的更能反映真实应用场景。可以说，这是任何 AI 产品走向成熟都绕不开的路。

基于最终的 prompt 集，OpenAI 构建了三份数据集：SFT 数据集、RM 数据集、PPO 数据集。这三个数据集在下面三个阶段里会分别派上用场。

有监督微调

SFT 数据集是标准的 (prompt, answer) 对。从 prompt 集里挑出一批指令，让标注工写答案——一共 13k 条 prompt，来源既有工人手写的，也有用户真实输入的。

SFT 的技术本身没什么好讲的，就是在预训练模型基础上，用这些人工标注的数据再做一轮微调。这里有个概念值得厘清：指令微调（IFT）和有监督微调（SFT）到底有什么区别？

Hugging Face 的一篇博文说得比较明白：IFT 的核心目标是让模型学会“听话”——比如 prompt 里出现 “summarize”，模型就知道当前任务是总结。IFT 可以看作是 SFT 的子集或前置步骤。经过 IFT 后，模型确实能听从指令了，但生成的内容不一定安全可靠。为了提升输出内容的帮助性、降低有害性，人们接着做 SFT，通过向模型展示“无害且有用”的回答样例，约束模型的生成行为。

通俗点说，IFT 的数据比较好收集——针对每个 NLP 任务写一段模板化的 prompt，把现有的数据集直接套进去就行。而 SFT 的回答需要人工逐条编写，收集成本高出不少。此外，从论文的分布来看，Flan-T5 这类研究中对 IFT 讨论得更多，而 OpenAI 的系列论文里强调的是 SFT。

训练奖励模型

奖励模型的输入是 (prompt + answer)，输出是一个标量分数，用来评估回答质量。通常的做法是把 SFT 模型最后的分类头去掉，换成一个回归头，作为奖励模型的初始化。但考虑到训练稳定性和计算成本，一般会选择参数更少的小模型——比如 SFT 是 175B 参数，奖励模型就只用到 6B。

训练奖励模型本质上是个回归任务。理想情况下，数据集的格式是 (input: prompt+answer, label: 人工打的分数)。问题是，怎么量化地给“好”和“坏”下定义？一段回答看起来不错，但到底是打 7 分、8 分还是 9 分？很难拿得准。

OpenAI 换了个思路：不要求绝对分数，而是让标注员对同一个 prompt 的不同回答做排序。具体来说，从 prompt 集里取样，对每个 prompt 让 SFT 模型生成 K=9 个回答，然后让标注工把它们排个序。有了排序关系，就有了相对好坏。每个 prompt 对应 9 个回答，可以组合出 C(9,2)=36 对 (好回答, 差回答)。

损失函数用的是排序任务里常见的 pairwise ranking loss。训练时，InstructGPT

：只需要做 9 次前向传播（算出 9 个分数），就能完成 36 次参数更新，避免了重复计算。RM 数据集包含 33k 个 prompt，既有标注工人手写的，也有用户真实输入的。注意每一条 prompt 对应 36 对回答，所以训练用到的数据量其实比 prompt 数大了一个数量级。

利用 PPO 微调 LLM

PPO 数据集里全是用户输入的 prompt，目的是与人机交互的真实场景尽可能贴合。

这一块实现细节挺多，但 OpenAI 论文写得很简略。下面结合论文和 TRL 代码仔细拆解，会用到一些强化学习术语。

先搞清楚“状态”和“动作”在语言模型语境下怎么对应：

，动作空间就是整个词表。和很多强化学习环境不同，这里的状态转移没有随机性——给定上一个状态和当前动作，下一个状态完全确定。

PPO 微调阶段涉及四个模型：

• SFT LM
  ：第一阶段的产物，在 PPO 阶段参数冻结，作为 reference model（参考模型）。

• RL LM
  ：我们要微调的语言模型，初始参数与 SFT LM 一致，强化学习术语里叫它 actor model。

• reward model
  ：第二阶段训练好的奖励模型，PPO 阶段参数冻结。

• critic model
  ：用于估算状态值函数，本质上是一种值函数近似，参数随 RL LM 一起更新。

PPO 里通常让 actor model 和 critic model 共享参数一起训练，然后把 policy gradient 的损失函数和 critic model 的损失函数（一般是 MSE）加权求和作为总损失。那actor model 和 critic model 具体怎么共享参数？直接看 TRL 的源码：

base_model_output = self.pretrained_model(
      input_ids=input_ids,
      attention_mask=attention_mask,
      **kwargs,
)
...
# 取最后一层的隐藏状态
last_hidden_state = base_model_output.hidden_states[-1]
# v_head 是一个输出维度为 1 的全连接层，将每个 token 映射成一个标量
# value shape: (B, L)
value = self.v_head(last_hidden_state).squeeze(-1)

关键点一目了然：

，让每个 token 对应一个标量，作为该 token 状态下未来可获得的累积（折扣）回报。所以 critic model 的输出形状和输入一致： (batch size, seq_len)。

需要说明的是，InstructGPT 论文里提到 value function（critic）和 reward model 架构相同（都是 6B），说明它和 actor model（175B）并没有共享参数。这里我们先沿 TRL 的思路往下走。

PPO 中的奖励函数

很多人以为 PPO 的奖励就是 reward model 直接输出的分数，实际没那么简单。论文给出了目标函数——也就是

。先忽略最后那一项，第一项 rθ(x,y) 是 reward model 的原始分数；第二项 log(πϕRL(y|x) / πSFT(y|x)) 相当于额外加的惩罚——如果当前的 RL LM 和 reference model（SFT LM）偏差太大，奖励就会减少。

这里有两个细节值得注意。第一，x 是 prompt，y 是 RL LM 生成的回答，πϕRL(y|x) 的形状与 y 相同，是一个序列长度 L 的向量（假设 batch size 为 1），表示每个 token 位置的生成概率。第二，公式里的最后一项和 PPO 论文里 policy gradient 目标函数中的 KL 散度完全是两回事——那个 KL 是为了限制重要性采样中采样策略和当前策略不要差距过大。

还有个麻烦：对于每个生成的 y，reward model 只输出一个标量 rθ(x,y)；但 log(πϕRL(y|x) / πSFT(y|x)) 是一个与 y 等长的向量，每个位置上都有一段 KL 散度。奖励函数怎么对齐？答案很巧妙：只在 log(πϕRL(y|x) / πSFT(y|x)) 的最后一个元素上加上 rθ(x,y)，得到的总向量就是完整的奖励信号。论文里把这种做法叫作“在每 token 位置添加来自 SFT 模型的 KL 惩罚，以防止奖励模型过度优化”。

TRL 的源码实现清晰地反映了这一点：

# 求 KL 散度
kl = self._kl_penalty(logprob, ref_logprob)
non_score_reward = -self.kl_ctl.value * kl
reward = non_score_reward.clone()
# y 中非 padding 的最后一位
last_non_masked_index = mask.nonzero()[-1]
# 最终的奖励 = KL 散度 + reward model 的输出（在 KL 散度最后一位相加）
# score 就是 reward model 的输出
reward[last_non_masked_index] += score

理解了奖励函数的构成，再回头看 TRL 文档里的示意图就清晰多了。后面会讲到如何用 policy gradient 优化模型，现在先锁定奖励函数的计算方式。

优势函数的计算

PPO 中优势函数的计算公式里，δt 是时间差分误差（TD error）。通用的计算方法是：先算出最后一个位置的 AT−1，然后用递推公式 At = δt + γλ At+1

。

TRL 源码的实现：

# reversed：倒序计算
for t in reversed(range(gen_len)):
   # nextvalues 就是 TD error 中的 V(s_{t+1})，即下一个状态的状态值函数
   nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
   # 计算 TD error，这里的 rewards 数组就是上一部分求的奖励函数
   delta = rewards[:, t] + self.config.gamma * nextvalues - values[:, t]
   # 利用递推公式求当前优势函数
   lastgaelam = delta + self.config.gamma * self.config.lam * lastgaelam
   advantages_reversed.append(lastgaelam)
# 由于是倒序计算，需要把列表翻转
advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

优势函数的计算用到了前面求出的奖励函数和 critic model 的输出。换言之，奖励函数是通过影响优势函数，间接作用于损失函数（梯度）计算的。InstructGPT 中有一句值得留意：“No discount is applied when estimating the generalized advantage”，即上面公式中的 λ=1。

actor model 和 critic model 的更新

有了优势函数，PPO 的目标函数就可以拿出来用了。rt(θ) 是因重要性采样而引入的当前策略与采样策略的概率比值。前面提到 actor model 和 critic model 共享参数和损失函数，一起训练。两部分损失函数需要各自计算：policy gradient 的损失函数和 critic model 的损失函数。

# ratio 对应当前策略与采样策略的概率比值
ratio = torch.exp(logprobs - old_logprobs)
# 由于是求损失函数，所以在目标函数前加个负号
pg_losses = -advantages * ratio
# clip 部分
pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - self.config.cliprange, 1.0 + self.config.cliprange)
# 求两部分的最大值
pg_loss = masked_mean(torch.max(pg_losses, pg_losses2), mask)

上面是 policy gradient 损失函数的源码。

vpredclipped = clip_by_value(
     vpreds,
     values - self.config.cliprange_value,
     values + self.config.cliprange_value,
)
vf_losses1 = (vpreds - returns) ** 2
vf_losses2 = (vpredclipped - returns) ** 2
vf_loss = 0.5 * masked_mean(torch.max(vf_losses1, vf_losses2), mask)

上面是 critic model 损失函数的源码。整体上是 MSE 的形式，但分别用预测值 vpreds 和裁剪过的预测值 vpredclipped 与 returns 求 MSE，然后取大者。有一点没搞明白的是：为什么最小化 vpreds - returns？这里的 returns = advantages + old_v_values，等价于最小化 vpreds - old_v_values - advantages。如果哪位更清楚这里的用意，欢迎交流。

最后把两部分 loss 求和，一起反向传播：

loss = pg_loss + self.config.vf_coef * vf_loss
self.accelerator.backward(loss)

整体流程

至此，示意图里的最后一块拼图——”Policy gradients optimize model“也合上了。Hugging Face 博客里有一张流程图，现在回看应该格外清晰。需要提醒的是：那张图看起来像是 RL LM 和 SFT LM 对同一个 prompt 分别生成不同的回答，实际流程是 RL LM 生成回答 y，然后送进 SFT LM 计算概率，再求得 KL 散度。

Deepspeed 代码实现的差异

再来看 Deepspeed 对第三阶段 PPO 微调的实现，和 TRL 存在一些差异。从 Deepspeed 的三阶段流程图可以看到，第三阶段开始时，actor model 拿 SFT LM 做初始化；但

——这意味着 critic model 和 actor model 并没有共享参数，而是分开训练的。源码中 self.actor_model.backward(actor_loss) 和 self.critic_model.backward(critic_loss) 分别调用，也能印证这一点。

Deepspeed 的文档里提到一个非常值得注意的问题：他们的训练极其不稳定——“we are not able to update the actor model multiple times after generating experimental data.” 按照标准 PPO 流程，每次按策略 πθold 收集数据组成大小为 N 的 data buffer，然后用 buffer 内的数据训练 πθ 多个 epochs；但 Deepspeed 的实际操作不得不把 epochs 和 buffer 大小都设为 1，即每次收集完数据只能更新一次参数，否则训练就会发散。

当初我猜测训练不稳定的原因是 actor 和 critic 分开训练导致的，但后来仔细看了 InstructGPT 论文，发现它的 value function（即这里的 critic model）同样是直接拿 reward model 初始化的：”As previously mentioned, for all PPO models we use a 6B RM and a 6B value function, and the latter is initialized from the former.“所以说，InstructGPT 里的 critic model 和 actor model 也不是共享参数的。从这个角度看，Deepspeed 的实现反而更贴合论文。

Bonus：奖励函数的另一项

InstructGPT 还在 PPO 的奖励函数里额外加了一项关于预训练损失的函数：Ex∼Dpretrain[log(πϕRL(x))]。它的作用是让 RL LM 在标准 NLP 任务上保持能力，避免在 RLHF 微调过程中产生“灾难性遗忘”。

TRL 里没看到对这一项的实现，但 Deepspeed 做了：

def train_unsupervised(self, inputs, unsup_coef):
    outputs = self.actor_model(**inputs, use_cache=False)
    loss = outputs.loss
    self.actor_model.backward(unsup_coef * loss)
    self.actor_model.step()

参数更新的顺序上：先根据 (x,y)∼DπϕRL 更新 actor 和 critic；如果开启了无监督训练，再单独更新 actor model。个人觉得这个实现思路很清晰——我们没法把无监督学习的损失直接放进奖励函数，所以只能把它当作和 PPO 损失函数“平行”的独立损失来处理。也就是说，公式里的第一个期望是作为奖励函数影响优势计算，进而影响 PPO 损失函数；第二个期望则直接作为损失函数参与参数更新。

后来翻到 InstructGPT 论文附录，果然验证了这个做法：”For each minibatch, we compute the PPO gradients and pretraining gradients in consecutive steps and accumulate them both into the gradient buffers.“

RLHF 的下一个探索方向

现阶段 RLHF 最大的痛点是训练不稳定——这其实也是强化学习的通病。加上 RLHF 对训练数据的质量要求很高，数据收集成本居高不下。Anthropic 发布过的 16 万条公开数据集 hh-rlhf，至今还是为数不多可以直接用的 RLHF 训练资源。

另外，PPO 是 2017 年提出的，目前没有任何证据表明它在 RLHF 场景下比其他强化学习算法有明显优势。可能只是 OpenAI 对 PPO 比较熟悉，所以拉上算法作者 Schulman 一起搞 RLHF。在第三阶段 PPO 微调时，reward model 也不是非冻不行的——它完全可以和 actor model 交替更新，只是每次更新后都需要人工为新的回答重新排序，成本极高。

总之一句话：RLHF 的很多环节都还是开放的，留给后来者的空间远比想象中大。