Chapter 08

PPOTrainer:经典 RLHF 三模型

📌 commit 8221a9fInstructGPT 同款,TRL 2024 版重构

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_lengthgenerate 的 max_new_tokens512-2048
num_ppo_epochs每批 rollout 用几次4
num_mini_batches每 epoch 切几个 mini-batch1-4
cliprange策略比 clip ε0.2
cliprange_valuevalue clip0.2
kl_coefKL 惩罚0.05-0.2
vf_coefvalue loss 权重0.1
gamma / lamGAE 参数1.0 / 0.95
missing_eos_penalty没 EOS 时扣 reward1.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] tokenizeHF Dataset + chat 模板
reward 接口外部传 list[float]HF 模型对象
vLLM 加速需自己写原生支持
多 GPU需手工 accelerateauto

新版位于 trl/trainer/ppo_trainer.py,老版被改名 ppo_trainer_legacy.py 保留兼容。

8.7什么时候选 PPO 而不是 GRPO / DPO

否则推荐 GRPO(去 critic 省显存)或 DPO(去 rollout 省时间)。

8.9v1:PPOTrainer 已经移到 experimental

重要 TRL v1(2026-03 发布)把 PPOTrainertrl/trainer/ 挪到了 trl/experimental/ppo/ppo_trainer.py:305主线 from trl import PPOTrainer 现在导的是 experimental 的版本,并附带 "API 不稳定" 警告。

这个变更意味着:

本章后续内容仍然适用(公式没变),但选型时默认走 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 通过 AutoModelForCausalLMWithValueHeadtrl/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这章你需要带走的