总目录›TRL›第 9 章
Chapter 09

Reward Model 训练

📌 commit 8221a9fBradley-Terry 偏好建模

PPO 和 GRPO（神经版 reward）都需要一个能打分的 RM。本章用 TRL 的 RewardTrainer 训一个出来。

9.1RM 是什么

RM 把语言模型的 lm_head 换成 value_head：对完整 (prompt, response) 序列输出一个标量分数。训练 loss 是 Bradley-Terry pairwise：

$$\mathcal{L}_\text{RM} = -\mathbb{E}_{(x,y_w,y_l)\sim D}\Big[\log\sigma\big(R(x,y_w) - R(x,y_l)\big)\Big]$$

$y_w$ 是被偏好回复、$y_l$ 是被拒绝回复。

9.2用 RewardTrainer

from trl import RewardTrainer, RewardConfig
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset

# 1) 模型（带 1-class classification head）
model = AutoModelForSequenceClassification.from_pretrained(
    "./sft_out/final", num_labels=1, torch_dtype="bfloat16",
)
tok = AutoTokenizer.from_pretrained("./sft_out/final")
tok.pad_token = tok.eos_token

# 2) 数据
ds = load_dataset("trl-lib/ultrafeedback_binarized", split="train[:5000]")

# 3) 配置
cfg = RewardConfig(
    output_dir="./rm_out",
    per_device_train_batch_size=4,
    num_train_epochs=1,
    learning_rate=2e-5,
    max_length=2048,
    bf16=True,
)

# 4) 训
trainer = RewardTrainer(
    model=model, args=cfg, processing_class=tok,
    train_dataset=ds,
)
trainer.train()

9.3数据格式

RewardTrainer 期待 dataset 有这两列之一：

格式列说明
conversationalchosen / rejected（each = list[messages]）推荐
plainchosen / rejected（each = str）简易

注意 prompt 必须相同（包含在两个列里），TRL 内部会做 tokenize + pad。

9.4RewardConfig 核心字段

字段含义建议
max_length截断2048
center_rewards_coefficientL2 把 rewards 拉到 0 附近0.01 推荐（防漂移）
disable_dropout关闭 dropoutTrue 推荐（RM 用 max-likelihood 不需要）
per_device_train_batch_size注意：内部会变成 2× pad（chosen+rejected）4-8

9.5RM 训练监控指标

训完 accuracy 太低（< 0.6）通常意味着偏好数据有噪声或base 模型不够强。

9.6RM 训完怎么用

用法 A：传给 PPOTrainer

rm = AutoModelForSequenceClassification.from_pretrained("./rm_out/final")
trainer = PPOTrainer(..., reward_model=rm, value_model=rm, ...)

用法 B：传给 GRPOTrainer（用 RM 当 reward function）

trainer = GRPOTrainer(
    ..., reward_funcs=["./rm_out/final"], ...,
)
# TRL 自动把字符串当 HF model 加载

用法 C：直接当 API 用

rm.eval()
with torch.no_grad():
    inputs = tok(prompt_and_response, return_tensors="pt").to(rm.device)
    score = rm(**inputs).logits[0].item()

9.7常见踩坑

现象处理
accuracy 一直 ~0.5（瞎猜水平）偏好数据有问题；查 chosen / rejected 是否搞反
chosen / rejected 分数都很大（同正/同负）开 center_rewards_coefficient=0.01
RM 训完用到 PPO 上 reward 全是负正常，PPO 只在乎相对差距；可以训前 normalize
训长 prompt OOMRM 的 batch 内部翻倍，per_device_bs 折半

9.8RewardConfig 关键字段（源码锚）

来源 trl/trainer/reward_config.py:23。RewardTrainer 字段很少，但有几个不调容易翻车：

字段含义建议
center_rewards_coefficient 奖励中心化项的权重（line 67–70）；推动 mean(reward)≈00.01（论文 2312.09244 推荐值）
disable_dropout 关掉 dropout 让 logits 确定（line 49–50）True（RM 训练时基本都这么写）
max_length 序列截断跟 SFT 一致或更长
remove_unused_columns HF 默认 True 会删 chosen/rejected 列；RewardTrainer 内部会重置不用动

9.9过程奖励 PRM：PRMTrainer

除了"对整段回复打一个分"的 outcome reward model（ORM），TRL 还提供过程奖励 PRM —— 对推理链每一步打分。位置 trl/experimental/prm/prm_trainer.py。

维度ORM（本章主线）PRM
训练数据 (prompt, chosen, rejected) 对 (prompt, [step1, step2, ...], [label1, label2, ...])
label 形式 chosen 比 rejected 好（隐式标签） 每步显式标签（true / false / mixed）
训练 loss Bradley-Terry pairwise 逐 token cross-entropy（在分隔 token 上算）
用于 RLHF reward signal R1 风格"评估推理步骤" / Best-of-N reranking

PRM 数据集格式（来自 PRMConfig）：

{
  "prompt": "What is 12 * 13?",
  "completions": [
      "Step 1: 12 * 13 = 12 * 10 + 12 * 3",   # 步骤 1
      "Step 2: = 120 + 36",                      # 步骤 2
      "Step 3: = 156"                            # 步骤 3
  ],
  "labels": [True, True, True]                    # 每步是否正确
}

PRM 已经在 Math-Shepherd、Qwen2.5-Math-PRM、Skywork-o1 等开源 PRM 项目里落地。TRL 给的就是训练入口，跑通 PRM 后可以当 reranker。

9.10训完之后：把 RM 喂给 GRPO / PPO

ORM 训完可以走两条路：

用法代码
GRPO/PPO 当 reward_modelreward_funcs="path/to/rm"（GRPO 自动 load AutoModelForSequenceClassification）
外部打分服务 trl chat --model_name_or_path rm-path 起 HTTP，外部调用
Best-of-N 评估 SFT 模型采样 N 条，用 RM 选最高分 → 单步无训练就提升效果
data filter 训练数据过 RM 算分，保留 top X% 高分样本再做 SFT

实战：4 种用法里 data filter 最被低估——同一个 RM 同时改进 SFT 数据质量和 RLHF 信号，省下一份模型也能用。

9.11这章你需要带走的