5.1 KiB

Raw Blame History

PPO approx_kl 这段代码在做什么

对应源码：

trainer/train_ppo.py:178-189

核心代码：

mb_resp_logp = F.log_softmax(res.logits[:, :-1], dim=-1) \
    .gather(2, labels[inds].unsqueeze(-1)) \
    .squeeze(-1) \
    .gather(1, logp_pos[inds])

log_ratio = mb_resp_logp - old_resp_logp[inds]
approx_kl = (
    0.5 * (log_ratio ** 2) * resp_policy_mask[inds]
).sum() / resp_policy_mask[inds].sum().clamp(min=1)

if approx_kl_val > args.early_stop_kl:
    stop_ppo = True

1. 先看两个 logprob

old_resp_logp：

rollout 时旧 actor 对 response token 的 logprob

数学上：

old_resp_logp[b, t] = log pi_old(a_{b,t} | s_{b,t})

mb_resp_logp：

PPO update 时当前 actor 对同一批 response token 重新算出来的 logprob

数学上：

mb_resp_logp[b, t] = log pi_theta(a_{b,t} | s_{b,t})

这里的 a_{b,t} 是第 b 条样本第 t 个 response token。

这里的 s_{b,t} 是：

prompt + response 前 t-1 个 token

2. log_ratio 是什么

源码：

log_ratio = mb_resp_logp - old_resp_logp[inds]

公式：

log_ratio[b, t]
= log pi_theta(a_{b,t} | s_{b,t})
  - log pi_old(a_{b,t} | s_{b,t})

也就是：

log_ratio[b, t]
= log( pi_theta(a_{b,t} | s_{b,t}) / pi_old(a_{b,t} | s_{b,t}) )

所以后面：

ratio = torch.exp(log_ratio)

就是：

ratio[b, t]
= pi_theta(a_{b,t} | s_{b,t}) / pi_old(a_{b,t} | s_{b,t})

直观理解：

ratio = 1:
  当前 actor 和 old actor 对这个 token 的概率一样。

ratio > 1:
  当前 actor 比 old actor 更倾向这个 token。

ratio < 1:
  当前 actor 比 old actor 更不倾向这个 token。

3. approx_kl 是什么

源码：

approx_kl = (
    0.5 * (log_ratio ** 2) * resp_policy_mask[inds]
).sum() / resp_policy_mask[inds].sum().clamp(min=1)

公式：

approx_kl
= sum_{b,t} m_{b,t} * 0.5 * log_ratio[b,t]^2
  / sum_{b,t} m_{b,t}

展开 log_ratio：

approx_kl
= sum_{b,t} m_{b,t} * 0.5 *
   (log pi_theta(a_{b,t}|s_{b,t})
    - log pi_old(a_{b,t}|s_{b,t}))^2
  / sum_{b,t} m_{b,t}

其中：

m_{b,t} = resp_policy_mask[b,t]

resp_policy_mask 的作用：

只统计有效 response token；
不统计 prompt；
不统计 padding；
不统计 EOS 后面的 token。

4. 为什么叫 approx_kl

PPO 需要监控：

当前 actor pi_theta 离 rollout 时的 old actor pi_old 有多远。

KL 本来是衡量两个分布距离的量。

当新旧策略很接近时，可以用二阶近似：

KL(pi_old || pi_theta)
approx 0.5 * (log pi_theta - log pi_old)^2

所以源码用：

0.5 * (log_ratio ** 2)

来近似新旧策略的 KL。

直观判断：

log_ratio 接近 0:
  new actor 和 old actor 很接近，approx_kl 小。

log_ratio 绝对值变大:
  new actor 和 old actor 差得更远，approx_kl 变大。

5. 它是不是训练 loss

这段 approx_kl 不是 PPO 的主训练 loss。

PPO 的 actor loss 在后面：

policy_loss = (
    clipped_policy_loss
    + args.kl_coef * kl_ref_penalty
)

approx_kl 的主要用途是 early stop：

if approx_kl_val > args.early_stop_kl:
    stop_ppo = True

意思是：

如果当前 actor 相对 old actor 改得太多，
就停止继续用这批 rollout 做 PPO 更新。

因为 PPO 是 on-policy 算法。

这批 response 是 old actor 生成的，如果当前 actor 已经离 old actor 太远，再继续拿这批旧 rollout 训练就不稳定。

6. 和 kl_ref_penalty 的区别

这两个很容易混：

approx_kl:
  比较当前 actor 和 rollout 时的 old actor。
  用途是 early stop / 监控 PPO 更新幅度。
  源码位置: train_ppo.py:180-189

kl_ref_penalty:
  比较当前 actor 和 frozen reference model。
  用途是放进 actor loss，防止模型偏离 SFT reference。
  源码位置: train_ppo.py:194-199

一句话：

approx_kl 控制“别离 old actor 太远”。
kl_ref_penalty 控制“别离 SFT reference 太远”。

7. 最小例子

假设某个 token：

old_resp_logp = -2.0
mb_resp_logp  = -1.8

那么：

log_ratio = -1.8 - (-2.0) = 0.2
ratio = exp(0.2) = 1.221
approx_kl_token = 0.5 * 0.2^2 = 0.02

含义：

当前 actor 比 old actor 更倾向这个 token；
但变化不大，approx_kl_token 只有 0.02。

如果：

old_resp_logp = -2.0
mb_resp_logp  = -0.5

那么：

log_ratio = 1.5
ratio = exp(1.5) = 4.48
approx_kl_token = 0.5 * 1.5^2 = 1.125

含义：

当前 actor 对这个 token 的概率提高太多；
如果 batch 平均也很大，就可能触发 early stop。

8. 记忆方式

old_resp_logp:
  rollout 时候的概率，固定不变。

mb_resp_logp:
  当前 actor 重新 forward 的概率，会随着训练变化。

log_ratio:
  当前 actor 相对 old actor 的变化量。

approx_kl:
  用 log_ratio 的平方估计新旧 actor 的距离。

early_stop_kl:
  如果距离太大，就不要继续拿这批 old rollout 训练。

5.1 KiB Raw Blame History Unescape Escape