305 lines
5.1 KiB
Markdown
305 lines
5.1 KiB
Markdown
# PPO approx_kl 这段代码在做什么
|
||
|
||
对应源码：
|
||
|
||
```text
|
||
trainer/train_ppo.py:178-189
|
||
```
|
||
|
||
核心代码：
|
||
|
||
```python
|
||
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`：
|
||
|
||
```text
|
||
rollout 时旧 actor 对 response token 的 logprob
|
||
```
|
||
|
||
数学上：
|
||
|
||
```text
|
||
old_resp_logp[b, t] = log pi_old(a_{b,t} | s_{b,t})
|
||
```
|
||
|
||
`mb_resp_logp`：
|
||
|
||
```text
|
||
PPO update 时当前 actor 对同一批 response token 重新算出来的 logprob
|
||
```
|
||
|
||
数学上：
|
||
|
||
```text
|
||
mb_resp_logp[b, t] = log pi_theta(a_{b,t} | s_{b,t})
|
||
```
|
||
|
||
这里的 `a_{b,t}` 是第 `b` 条样本第 `t` 个 response token。
|
||
|
||
这里的 `s_{b,t}` 是：
|
||
|
||
```text
|
||
prompt + response 前 t-1 个 token
|
||
```
|
||
|
||
## 2. log_ratio 是什么
|
||
|
||
源码：
|
||
|
||
```python
|
||
log_ratio = mb_resp_logp - old_resp_logp[inds]
|
||
```
|
||
|
||
公式：
|
||
|
||
```text
|
||
log_ratio[b, t]
|
||
= log pi_theta(a_{b,t} | s_{b,t})
|
||
- log pi_old(a_{b,t} | s_{b,t})
|
||
```
|
||
|
||
也就是：
|
||
|
||
```text
|
||
log_ratio[b, t]
|
||
= log( pi_theta(a_{b,t} | s_{b,t}) / pi_old(a_{b,t} | s_{b,t}) )
|
||
```
|
||
|
||
所以后面：
|
||
|
||
```python
|
||
ratio = torch.exp(log_ratio)
|
||
```
|
||
|
||
就是：
|
||
|
||
```text
|
||
ratio[b, t]
|
||
= pi_theta(a_{b,t} | s_{b,t}) / pi_old(a_{b,t} | s_{b,t})
|
||
```
|
||
|
||
直观理解：
|
||
|
||
```text
|
||
ratio = 1:
|
||
当前 actor 和 old actor 对这个 token 的概率一样。
|
||
|
||
ratio > 1:
|
||
当前 actor 比 old actor 更倾向这个 token。
|
||
|
||
ratio < 1:
|
||
当前 actor 比 old actor 更不倾向这个 token。
|
||
```
|
||
|
||
## 3. approx_kl 是什么
|
||
|
||
源码：
|
||
|
||
```python
|
||
approx_kl = (
|
||
0.5 * (log_ratio ** 2) * resp_policy_mask[inds]
|
||
).sum() / resp_policy_mask[inds].sum().clamp(min=1)
|
||
```
|
||
|
||
公式：
|
||
|
||
```text
|
||
approx_kl
|
||
= sum_{b,t} m_{b,t} * 0.5 * log_ratio[b,t]^2
|
||
/ sum_{b,t} m_{b,t}
|
||
```
|
||
|
||
展开 `log_ratio`：
|
||
|
||
```text
|
||
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}
|
||
```
|
||
|
||
其中：
|
||
|
||
```text
|
||
m_{b,t} = resp_policy_mask[b,t]
|
||
```
|
||
|
||
`resp_policy_mask` 的作用：
|
||
|
||
```text
|
||
只统计有效 response token；
|
||
不统计 prompt；
|
||
不统计 padding；
|
||
不统计 EOS 后面的 token。
|
||
```
|
||
|
||
## 4. 为什么叫 approx_kl
|
||
|
||
PPO 需要监控：
|
||
|
||
```text
|
||
当前 actor pi_theta 离 rollout 时的 old actor pi_old 有多远。
|
||
```
|
||
|
||
KL 本来是衡量两个分布距离的量。
|
||
|
||
当新旧策略很接近时，可以用二阶近似：
|
||
|
||
```text
|
||
KL(pi_old || pi_theta)
|
||
approx 0.5 * (log pi_theta - log pi_old)^2
|
||
```
|
||
|
||
所以源码用：
|
||
|
||
```python
|
||
0.5 * (log_ratio ** 2)
|
||
```
|
||
|
||
来近似新旧策略的 KL。
|
||
|
||
直观判断：
|
||
|
||
```text
|
||
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 在后面：
|
||
|
||
```python
|
||
policy_loss = (
|
||
clipped_policy_loss
|
||
+ args.kl_coef * kl_ref_penalty
|
||
)
|
||
```
|
||
|
||
`approx_kl` 的主要用途是 early stop：
|
||
|
||
```python
|
||
if approx_kl_val > args.early_stop_kl:
|
||
stop_ppo = True
|
||
```
|
||
|
||
意思是：
|
||
|
||
```text
|
||
如果当前 actor 相对 old actor 改得太多，
|
||
就停止继续用这批 rollout 做 PPO 更新。
|
||
```
|
||
|
||
因为 PPO 是 on-policy 算法。
|
||
|
||
这批 response 是 old actor 生成的，如果当前 actor 已经离 old actor 太远，再继续拿这批旧 rollout 训练就不稳定。
|
||
|
||
## 6. 和 kl_ref_penalty 的区别
|
||
|
||
这两个很容易混：
|
||
|
||
```text
|
||
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
|
||
```
|
||
|
||
一句话：
|
||
|
||
```text
|
||
approx_kl 控制“别离 old actor 太远”。
|
||
kl_ref_penalty 控制“别离 SFT reference 太远”。
|
||
```
|
||
|
||
## 7. 最小例子
|
||
|
||
假设某个 token：
|
||
|
||
```text
|
||
old_resp_logp = -2.0
|
||
mb_resp_logp = -1.8
|
||
```
|
||
|
||
那么：
|
||
|
||
```text
|
||
log_ratio = -1.8 - (-2.0) = 0.2
|
||
ratio = exp(0.2) = 1.221
|
||
approx_kl_token = 0.5 * 0.2^2 = 0.02
|
||
```
|
||
|
||
含义：
|
||
|
||
```text
|
||
当前 actor 比 old actor 更倾向这个 token；
|
||
但变化不大，approx_kl_token 只有 0.02。
|
||
```
|
||
|
||
如果：
|
||
|
||
```text
|
||
old_resp_logp = -2.0
|
||
mb_resp_logp = -0.5
|
||
```
|
||
|
||
那么：
|
||
|
||
```text
|
||
log_ratio = 1.5
|
||
ratio = exp(1.5) = 4.48
|
||
approx_kl_token = 0.5 * 1.5^2 = 1.125
|
||
```
|
||
|
||
含义：
|
||
|
||
```text
|
||
当前 actor 对这个 token 的概率提高太多；
|
||
如果 batch 平均也很大，就可能触发 early stop。
|
||
```
|
||
|
||
## 8. 记忆方式
|
||
|
||
```text
|
||
old_resp_logp:
|
||
rollout 时候的概率，固定不变。
|
||
|
||
mb_resp_logp:
|
||
当前 actor 重新 forward 的概率，会随着训练变化。
|
||
|
||
log_ratio:
|
||
当前 actor 相对 old actor 的变化量。
|
||
|
||
approx_kl:
|
||
用 log_ratio 的平方估计新旧 actor 的距离。
|
||
|
||
early_stop_kl:
|
||
如果距离太大，就不要继续拿这批 old rollout 训练。
|
||
```
|