大模型 RLHF 强化学习微调过程详解
引言
在大型语言模型（LLM）的训练流程中，通常包含预训练、有监督微调（SFT）和人类反馈强化学习（RLHF）三个阶段。虽然网上关于 PPO（Proximal Policy Optimization）的文章不少，但很多内容浅尝辄止。本文结合微软 DeepSpeed 的 RLHF 实现代码，深入讲解奖励模型训练和强化学习微调的核心逻辑。
在强化学习微调语言模型时，Prompt 对应状态（State），输出一串单词（Action），得到一个 Reward。与传统游戏场景不同，这里的状态转移蕴含在 Transformer 内部，每个动作后并不立即获得新的外部 State，而是通过模型内部机制更新隐状态。对于回答的第二个词，可以将 Prompt+ 第一个词当作新的 State。
奖励（Reward）模型训练
1. 模型初始化说明
在强化学习阶段，使用的 Reward Model 和 Critic Model 使用同一个模型初始化。因此，在训练 Reward 模型的过程中，实际上也在训练 Critic Model。
符号说明：大模型中间隐藏层的参数维度为 (B, L, D)，其中 B 为 Batch Size，L 为句子长度，D 为 Embedding 维度。
2. 奖励模型的作用
RLHF 需要一个奖励模型来评估语言大模型（Actor Model）回答的好坏。该模型通常比被评估的语言大模型小一些（例如 Actor 66B，Reward 350M）。输入为 Prompt + Answer，输出用于打分。
奖励模型最后一层隐藏层的输出维度为 (B, L, D)，通过一个 D×1 的全连接层将维度变为 (B, L)。在 L 维度上，第 i 个位置的数据表示从第 i 个位置到最后一个位置输出所能获得的奖励分值的累加和（类似 DQN 中的 Q 值）。这种形式的输出满足了 Critic Model 的输出要求。
# Huggingface 模型返回值是个 list，第 0 位是模型最后输出的 hidden state
hidden_states = transformer_outputs[0]
# v_head 为 Dx1 的全连接网络对最后一维压缩
rewards = self.v_head(hidden_states).squeeze(-1)
对于一个奖励模型来说，目标是给一个句子进行打分。按理说每个句子对应一个分值即可，但目前长度为 L 的句子，奖励模型输出了 L 个值。我们取 L 维度上的最后一个位置的值作为本句话的奖励得分。
3. Pairwise Loss 训练
奖励模型训练优化采用 Pairwise 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]
Pairwise Loss 代码如下，如果给 pair 里好的句子打分高（c_truncated_reward），坏的句子（r_truncated_reward）打分低，loss 就会小：
loss += -torch.log(torch.sigmoid(c_truncated_reward - r_truncated_reward)).mean()
4. 推理过程中的 Critic Model
在训练强化学习的过程中，会用到 Reward Model（Critic Model，两者是同一模型的两个副本）的推理过程。通过调用 forward_value 实现，返回的值中有两种值：