PPOTrainer:经典 RLHF 三模型
PPO 是 RLHF 的"祖宗算法"。TRL 2024 中后期对 PPOTrainer 做了一次大重构(v0.10+),把数据流改成对话格式 + 显式 reward model,跟 SFTTrainer / DPOTrainer 保持一致。本章讲新版用法。
8.1PPO objective 回顾
$$\mathcal{L}_\text{PPO} = -\mathbb{E}_t\Big[\min\!\big(\rho_t \hat A_t,\,\text{clip}(\rho_t,1-\epsilon,1+\epsilon)\hat A_t\big)\Big] + \beta\mathrm{KL}(\pi\|\pi_\text{ref})$$
$\hat A_t$ 是 GAE advantage、$\rho_t$ 是策略比、$\epsilon$ 一般 0.2。
8.2用 PPOTrainer
from trl import PPOTrainer, PPOConfig
cfg = PPOConfig(
output_dir="./ppo_out",
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
num_ppo_epochs=4, # 每个 batch update 几次
num_mini_batches=1,
learning_rate=1e-6,
response_length=512,
kl_coef=0.05,
cliprange=0.2,
cliprange_value=0.2,
vf_coef=0.1,
bf16=True,
)
trainer = PPOTrainer(
args=cfg,
processing_class=tok,
model=policy_model,
ref_model=ref_model,
reward_model=rm_model, # 标量 reward head
value_model=value_model, # critic
train_dataset=prompt_ds,
)
trainer.train()
8.3四模型显存
policy / reference / reward / value 都要在 GPU 上:
| 模型 (bf16 7B) | 显存 | 说明 |
|---|---|---|
| policy | ~14 GB + 梯度 + AdamW 累计 ~40 GB | 训练 |
| reference | ~14 GB(无梯度) | 冻结 |
| reward | ~14 GB(无梯度) | 冻结 |
| value (critic) | ~14 GB + 梯度 + AdamW ~40 GB | 训练 |
| 合计 | ~110+ GB | 必须多卡 / LoRA |
8.4核心字段
| 字段 | 含义 | 建议 |
|---|---|---|
response_length | generate 的 max_new_tokens | 512-2048 |
num_ppo_epochs | 每批 rollout 用几次 | 4 |
num_mini_batches | 每 epoch 切几个 mini-batch | 1-4 |
cliprange | 策略比 clip ε | 0.2 |
cliprange_value | value clip | 0.2 |
kl_coef | KL 惩罚 | 0.05-0.2 |
vf_coef | value loss 权重 | 0.1 |
gamma / lam | GAE 参数 | 1.0 / 0.95 |
missing_eos_penalty | 没 EOS 时扣 reward | 1.0 推荐 |
8.5典型 PPO 训练循环(内部)
# 简化
for batch in dl:
# 1) Rollout
responses = policy.generate(prompts)
# 2) Reward
rewards = reward_model(responses)
# 3) KL
log_ratio = compute_log_ratio(policy, ref, responses)
rewards = rewards - kl_coef * log_ratio
# 4) Value
values = value_model(responses)
# 5) GAE
advantages = compute_gae(rewards, values, gamma, lam)
# 6) PPO update (multiple epochs)
for epoch in range(num_ppo_epochs):
for mini_batch in split(batch):
# policy loss
ratio = exp(log_pi_new - log_pi_old)
pg_loss = -min(ratio * adv, clip(ratio, 1-e, 1+e) * adv).mean()
# value loss
vf_loss = (value_new - returns) ** 2
loss = pg_loss + vf_coef * vf_loss
loss.backward(); optimizer.step()
8.6新版 vs 老版 PPOTrainer
| v0.7 之前老版 | v0.10+ 新版 | |
|---|---|---|
| API 风格 | 命令式(一次 step 一次) | 声明式(同 SFTTrainer) |
| 数据格式 | list[str] tokenize | HF Dataset + chat 模板 |
| reward 接口 | 外部传 list[float] | HF 模型对象 |
| vLLM 加速 | 需自己写 | 原生支持 |
| 多 GPU | 需手工 accelerate | auto |
新版位于 trl/trainer/ppo_trainer.py,老版被改名 ppo_trainer_legacy.py 保留兼容。
8.7什么时候选 PPO 而不是 GRPO / DPO
- 有 reward signal 不能写成规则,且必须用神经网络 RM;
- 需要严格的 critic 估值(如对长 horizon 任务);
- 历史包袱 / 复现 InstructGPT 论文;
- 显存够。
否则推荐 GRPO(去 critic 省显存)或 DPO(去 rollout 省时间)。
8.9v1:PPOTrainer 已经移到 experimental
PPOTrainer 从 trl/trainer/ 挪到了
trl/experimental/ppo/ppo_trainer.py:305。主线 from trl import PPOTrainer 现在导的是 experimental 的版本,并附带 "API 不稳定" 警告。
这个变更意味着:
- 新功能优先进 GRPO / RLOO,PPO 不再扩展;
- experimental 的 API 可能在 patch 版本之间变;
- 想关警告:
export TRL_EXPERIMENTAL_SILENCE=1; - 2025 中之后所有正式 RLHF 训练框架(OpenRLHF / verl)也把 PPO 视为"老路径",GRPO 是新标准。
本章后续内容仍然适用(公式没变),但选型时默认走 ch07 的 GRPO。
8.10PPOConfig 关键字段
来自 trl/experimental/ppo/ppo_config.py。重要的几个:
| 字段 | 含义 |
|---|---|
kl_coef | KL 散度惩罚权重,动态调整时与 target_kl 配合 |
whiten_advantages | 每 batch 把 advantage 做 z-score 归一化,稳训练 |
gamma / lam | GAE 的 γ 和 λ;语言模型场景常用 γ=1, λ=0.95 |
cliprange | importance ratio clip 范围,默认 0.2 |
cliprange_value | value head 输出 clip 范围 |
num_ppo_epochs | 每 batch 数据重复更新轮数,常 4 |
num_mini_batches | 每 ppo epoch 切几个 mini batch |
vf_coef | value loss 权重 |
8.11value head 是怎么加上去的
PPO 需要 critic。TRL 通过 AutoModelForCausalLMWithValueHead(trl/models/modeling_value_head.py)在 base 模型顶上挂一个 nn.Linear(hidden, 1) 充当 value 头:
| 行为 | 实现 |
|---|---|
| forward 同时返回 (logits, loss, values) | 三元组 tuple,便于 PPO 用 |
| save 时 value head 单独保存 | v_head.bin 或 safetensors |
| load 时按名字加载 | 没找到则警告(warm-start 友好) |
| 支持 PEFT | LoRA 时 value head 也走 LoRA 风格训练 |
所以 PPO checkpoint 的目录里能看到 v_head.safetensors 这种额外文件。merge LoRA 时value head 不被合并到 base —— 推理引擎只需要 LM head,value head 是训练副产物。
8.8这章你需要带走的
- PPO 维护 4 个模型,显存约 SFT 的 4×;
- 核心超参:
kl_coef、cliprange、num_ppo_epochs; - TRL v0.10+ 把 PPOTrainer 改成跟 SFTTrainer 一致的声明式;
- 2025 实战中 PPO 多数被 GRPO 取代,PPO 仅在"必须用 critic"时还有价值。