862 lines
12 KiB
Markdown
862 lines
12 KiB
Markdown
# PPO update loss block 源码解析
|
||
|
||
本文只解释 MiniMind PPO 训练中的这一段：
|
||
|
||
```text
|
||
trainer/train_ppo.py:171-214
|
||
```
|
||
|
||
它完成一次 PPO mini-batch update 的核心计算：
|
||
|
||
```text
|
||
当前 actor logprob
|
||
-> ratio
|
||
-> clipfrac
|
||
-> reference KL penalty
|
||
-> clipped policy loss
|
||
-> clipped value loss
|
||
-> 合成最终 loss
|
||
-> backward
|
||
```
|
||
|
||
## 0. 变量约定
|
||
|
||
这一节先按源码变量读。数学符号只用来把 `[batch, response_len]` 里的某个元素写清楚。
|
||
|
||
### 0.1 下标怎么读
|
||
|
||
```text
|
||
b = batch 里的第几条样本
|
||
t = response 里的第几个 token
|
||
```
|
||
|
||
所以：
|
||
|
||
```text
|
||
x[b, t]
|
||
```
|
||
|
||
就是：
|
||
|
||
```text
|
||
第 b 条样本，第 t 个 response token 位置上的值。
|
||
```
|
||
|
||
### 0.2 action 和 state
|
||
|
||
第 `b` 条样本第 `t` 个 response token 记作：
|
||
|
||
$$
|
||
a_{b,t}
|
||
$$
|
||
|
||
它对应源码里的：
|
||
|
||
```python
|
||
completion_ids[b, t]
|
||
```
|
||
|
||
生成这个 token 时，模型看到的上下文记作：
|
||
|
||
$$
|
||
s_{b,t}
|
||
$$
|
||
|
||
它表示：
|
||
|
||
```text
|
||
prompt + response 前 t-1 个 token
|
||
```
|
||
|
||
也就是说，PPO 里的：
|
||
|
||
```text
|
||
state = s_{b,t}
|
||
action = a_{b,t}
|
||
```
|
||
|
||
### 0.3 logprob 符号
|
||
|
||
本文用小写希腊字母 ell：
|
||
|
||
$$
|
||
\ell
|
||
$$
|
||
|
||
表示 logprob。
|
||
|
||
当前 actor 的 token logprob：
|
||
|
||
$$
|
||
\ell^{\theta}_{b,t}
|
||
=
|
||
\log \pi_{\theta}(a_{b,t} \mid s_{b,t})
|
||
$$
|
||
|
||
对应源码：
|
||
|
||
```python
|
||
mb_resp_logp[b, t]
|
||
```
|
||
|
||
rollout 时 old actor 的 token logprob：
|
||
|
||
$$
|
||
\ell^{\mathrm{old}}_{b,t}
|
||
=
|
||
\log \pi_{\mathrm{old}}(a_{b,t} \mid s_{b,t})
|
||
$$
|
||
|
||
对应源码：
|
||
|
||
```python
|
||
old_resp_logp[b, t]
|
||
```
|
||
|
||
frozen reference model 的 token logprob：
|
||
|
||
$$
|
||
\ell^{\mathrm{ref}}_{b,t}
|
||
=
|
||
\log \pi_{\mathrm{ref}}(a_{b,t} \mid s_{b,t})
|
||
$$
|
||
|
||
对应源码：
|
||
|
||
```python
|
||
ref_resp_logp[b, t]
|
||
```
|
||
|
||
### 0.4 其它源码变量
|
||
|
||
有效 response token mask：
|
||
|
||
```python
|
||
resp_policy_mask[b, t]
|
||
```
|
||
|
||
公式里写成：
|
||
|
||
$$
|
||
m_{b,t}
|
||
$$
|
||
|
||
GAE 算出的 advantage：
|
||
|
||
```python
|
||
advantages[b, t]
|
||
```
|
||
|
||
公式里写成：
|
||
|
||
$$
|
||
A_{b,t}
|
||
$$
|
||
|
||
当前 critic value：
|
||
|
||
```python
|
||
mb_resp_values[b, t]
|
||
```
|
||
|
||
公式里写成：
|
||
|
||
$$
|
||
V_{\theta,b,t}
|
||
$$
|
||
|
||
rollout 阶段保存的 old critic value：
|
||
|
||
```python
|
||
old_resp_values[b, t]
|
||
```
|
||
|
||
公式里写成：
|
||
|
||
$$
|
||
V_{\mathrm{old},b,t}
|
||
$$
|
||
|
||
critic 的训练目标 return：
|
||
|
||
```python
|
||
returns[b, t]
|
||
```
|
||
|
||
公式里写成：
|
||
|
||
$$
|
||
G_{b,t}
|
||
$$
|
||
|
||
### 0.5 masked mean
|
||
|
||
后面很多 loss 都只在有效 response token 上求平均。统一写成：
|
||
|
||
$$
|
||
\operatorname{mean}_m(x)
|
||
=
|
||
\frac{
|
||
\sum_{b,t} m_{b,t} x_{b,t}
|
||
}{
|
||
\sum_{b,t} m_{b,t}
|
||
}
|
||
$$
|
||
|
||
其中 \(m_{b,t}\) 对应源码里的 `resp_policy_mask[inds]` 或 `resp_value_mask[inds]`。
|
||
|
||
## 1. 当前 actor 重新计算 response token logprob
|
||
|
||
源码：
|
||
|
||
```python
|
||
res = actor_unwrapped(input_ids=gen_out[inds], attention_mask=full_mask[inds])
|
||
|
||
mb_resp_logp = F.log_softmax(res.logits[:, :-1], dim=-1) \
|
||
.gather(2, labels[inds].unsqueeze(-1)) \
|
||
.squeeze(-1) \
|
||
.gather(1, logp_pos[inds])
|
||
```
|
||
|
||
公式：
|
||
|
||
$$
|
||
\ell^{\theta}_{b,t}
|
||
=
|
||
\log \pi_{\theta}(a_{b,t} \mid s_{b,t})
|
||
$$
|
||
|
||
含义：
|
||
|
||
```text
|
||
rollout 阶段已经保存了 old actor 的 old_resp_logp。
|
||
现在 PPO update 阶段，要用当前 actor 对同一批 response token 重新算 logprob。
|
||
```
|
||
|
||
注意：
|
||
|
||
```text
|
||
这里没有重新生成 response。
|
||
它是在同一批 old rollout 上重新 forward，计算当前 actor 对这些 token 的概率。
|
||
```
|
||
|
||
## 2. log_ratio 和 ratio
|
||
|
||
源码：
|
||
|
||
```python
|
||
log_ratio = mb_resp_logp - old_resp_logp[inds]
|
||
ratio = torch.exp(log_ratio)
|
||
```
|
||
|
||
公式：
|
||
|
||
$$
|
||
\log r_{b,t}
|
||
=
|
||
\ell^{\theta}_{b,t}
|
||
-
|
||
\ell^{\mathrm{old}}_{b,t}
|
||
$$
|
||
|
||
$$
|
||
r_{b,t}
|
||
=
|
||
\exp(\log r_{b,t})
|
||
=
|
||
\frac{
|
||
\pi_{\theta}(a_{b,t} \mid s_{b,t})
|
||
}{
|
||
\pi_{\mathrm{old}}(a_{b,t} \mid s_{b,t})
|
||
}
|
||
$$
|
||
|
||
含义：
|
||
|
||
```text
|
||
ratio 表示当前 actor 相对 rollout 时 old actor，对同一个 token 的概率改了多少。
|
||
```
|
||
|
||
判断方式：
|
||
|
||
```text
|
||
ratio = 1:
|
||
当前 actor 和 old actor 对这个 token 概率一样。
|
||
|
||
ratio > 1:
|
||
当前 actor 更倾向这个 token。
|
||
|
||
ratio < 1:
|
||
当前 actor 更不倾向这个 token。
|
||
```
|
||
|
||
## 3. approx_kl：监控 actor 是否离 old actor 太远
|
||
|
||
源码：
|
||
|
||
```python
|
||
approx_kl = (
|
||
0.5 * (log_ratio ** 2) * resp_policy_mask[inds]
|
||
).sum() / resp_policy_mask[inds].sum().clamp(min=1)
|
||
```
|
||
|
||
公式：
|
||
|
||
$$
|
||
\operatorname{approx\_kl}
|
||
=
|
||
\operatorname{mean}_m
|
||
\left(
|
||
\frac{1}{2}
|
||
(\log r_{b,t})^2
|
||
\right)
|
||
$$
|
||
|
||
展开：
|
||
|
||
$$
|
||
\operatorname{approx\_kl}
|
||
=
|
||
\operatorname{mean}_m
|
||
\left(
|
||
\frac{1}{2}
|
||
\left[
|
||
\log \pi_{\theta}(a_{b,t} \mid s_{b,t})
|
||
-
|
||
\log \pi_{\mathrm{old}}(a_{b,t} \mid s_{b,t})
|
||
\right]^2
|
||
\right)
|
||
$$
|
||
|
||
用途：
|
||
|
||
```text
|
||
approx_kl 不直接作为 policy_loss 的主项。
|
||
它主要用于 early stop。
|
||
```
|
||
|
||
源码：
|
||
|
||
```python
|
||
if approx_kl_val > args.early_stop_kl:
|
||
stop_ppo = True
|
||
```
|
||
|
||
含义：
|
||
|
||
```text
|
||
如果当前 actor 已经离 old actor 太远，就停止继续用这批 rollout 更新。
|
||
```
|
||
|
||
## 4. clipfrac：统计有多少 token 超出 PPO clip 范围
|
||
|
||
源码：
|
||
|
||
```python
|
||
clipfrac = (
|
||
(((ratio - 1.0).abs() > args.clip_epsilon).float()
|
||
* resp_policy_mask[inds]).sum()
|
||
/ resp_policy_mask[inds].sum().clamp(min=1)
|
||
)
|
||
```
|
||
|
||
公式：
|
||
|
||
$$
|
||
\operatorname{clipfrac}
|
||
=
|
||
\operatorname{mean}_m
|
||
\left(
|
||
\mathbf{1}
|
||
\left[
|
||
\left|r_{b,t} - 1\right| > \epsilon
|
||
\right]
|
||
\right)
|
||
$$
|
||
|
||
其中：
|
||
|
||
$$
|
||
\epsilon = \texttt{args.clip\_epsilon}
|
||
$$
|
||
|
||
含义：
|
||
|
||
```text
|
||
clipfrac 不是 loss。
|
||
它是一个监控指标，表示有多少有效 response token 的 ratio 已经超过 PPO 裁剪范围。
|
||
```
|
||
|
||
例子：
|
||
|
||
```text
|
||
clip_epsilon = 0.2
|
||
ratio 合理范围 = [0.8, 1.2]
|
||
|
||
ratio = 1.05 -> 不计入 clipfrac
|
||
ratio = 1.35 -> 计入 clipfrac
|
||
ratio = 0.70 -> 计入 clipfrac
|
||
```
|
||
|
||
如果 `clipfrac` 很大，说明：
|
||
|
||
```text
|
||
这批 update 里很多 token 的策略变化已经被 PPO clip 限制住了。
|
||
```
|
||
|
||
## 5. reference KL penalty：约束 actor 不要偏离 reference
|
||
|
||
源码：
|
||
|
||
```python
|
||
kl_ref_penalty = (
|
||
(torch.exp(ref_resp_logp[inds] - mb_resp_logp)
|
||
- (ref_resp_logp[inds] - mb_resp_logp)
|
||
- 1.0)
|
||
* resp_policy_mask[inds]
|
||
).sum() / resp_policy_mask[inds].sum().clamp(min=1)
|
||
```
|
||
|
||
先定义：
|
||
|
||
$$
|
||
\Delta_{b,t}
|
||
=
|
||
\ell^{\mathrm{ref}}_{b,t}
|
||
-
|
||
\ell^{\theta}_{b,t}
|
||
=
|
||
\log
|
||
\frac{
|
||
\pi_{\mathrm{ref}}(a_{b,t} \mid s_{b,t})
|
||
}{
|
||
\pi_{\theta}(a_{b,t} \mid s_{b,t})
|
||
}
|
||
$$
|
||
|
||
单个 token 的 penalty：
|
||
|
||
$$
|
||
k_{b,t}
|
||
=
|
||
\exp(\Delta_{b,t})
|
||
-
|
||
\Delta_{b,t}
|
||
-
|
||
1
|
||
$$
|
||
|
||
masked mean：
|
||
|
||
$$
|
||
\operatorname{kl\_ref\_penalty}
|
||
=
|
||
\operatorname{mean}_m
|
||
\left(
|
||
\exp(\Delta_{b,t}) - \Delta_{b,t} - 1
|
||
\right)
|
||
$$
|
||
|
||
含义：
|
||
|
||
```text
|
||
reference model 是 frozen SFT 模型。
|
||
这个项惩罚当前 actor 偏离 reference model。
|
||
```
|
||
|
||
注意它和 `approx_kl` 不同：
|
||
|
||
```text
|
||
approx_kl:
|
||
当前 actor vs old actor
|
||
用于 early stop
|
||
|
||
kl_ref_penalty:
|
||
当前 actor vs reference model
|
||
加进 policy_loss
|
||
```
|
||
|
||
## 6. policy_loss：PPO clipped surrogate + reference KL
|
||
|
||
源码：
|
||
|
||
```python
|
||
policy_loss = (
|
||
(
|
||
torch.max(
|
||
-advantages[inds] * ratio,
|
||
-advantages[inds] * torch.clamp(
|
||
ratio,
|
||
1.0 - args.clip_epsilon,
|
||
1.0 + args.clip_epsilon,
|
||
)
|
||
)
|
||
* resp_policy_mask[inds]
|
||
).sum() / resp_policy_mask[inds].sum().clamp(min=1)
|
||
+ args.kl_coef * kl_ref_penalty
|
||
)
|
||
```
|
||
|
||
先定义裁剪后的 ratio：
|
||
|
||
$$
|
||
\bar{r}_{b,t}
|
||
=
|
||
\operatorname{clip}
|
||
\left(
|
||
r_{b,t},
|
||
1-\epsilon,
|
||
1+\epsilon
|
||
\right)
|
||
$$
|
||
|
||
PPO 原始最大化目标通常写成：
|
||
|
||
$$
|
||
\mathcal{J}_{\mathrm{clip}}
|
||
=
|
||
\mathbb{E}
|
||
\left[
|
||
\min
|
||
\left(
|
||
r_{b,t} A_{b,t},
|
||
\bar{r}_{b,t} A_{b,t}
|
||
\right)
|
||
\right]
|
||
$$
|
||
|
||
源码是在最小化 loss，所以写成：
|
||
|
||
$$
|
||
\mathcal{L}_{\mathrm{policy\_clip}}
|
||
=
|
||
\operatorname{mean}_m
|
||
\left[
|
||
\max
|
||
\left(
|
||
-A_{b,t} r_{b,t},
|
||
-A_{b,t} \bar{r}_{b,t}
|
||
\right)
|
||
\right]
|
||
$$
|
||
|
||
MiniMind 最终 policy loss：
|
||
|
||
$$
|
||
\mathcal{L}_{\mathrm{policy}}
|
||
=
|
||
\mathcal{L}_{\mathrm{policy\_clip}}
|
||
+
|
||
\lambda_{\mathrm{KL}}
|
||
\operatorname{kl\_ref\_penalty}
|
||
$$
|
||
|
||
其中：
|
||
|
||
$$
|
||
\lambda_{\mathrm{KL}} = \texttt{args.kl\_coef}
|
||
$$
|
||
|
||
直观理解：
|
||
|
||
```text
|
||
advantage > 0:
|
||
这个 token 比 critic 预期更好，希望提高它的概率。
|
||
但 ratio 不能无限变大，最多按 1 + epsilon 裁剪。
|
||
|
||
advantage < 0:
|
||
这个 token 比 critic 预期更差，希望降低它的概率。
|
||
但 ratio 不能无限变小，最多按 1 - epsilon 裁剪。
|
||
```
|
||
|
||
所以 `policy_loss` 同时做两件事：
|
||
|
||
```text
|
||
1. 按 advantage 更新 actor。
|
||
2. 用 reference KL 约束 actor 不要偏离 SFT reference 太远。
|
||
```
|
||
|
||
## 7. value_loss：critic 的 clipped value loss
|
||
|
||
源码：
|
||
|
||
```python
|
||
value_loss = 0.5 * (
|
||
torch.max(
|
||
(mb_resp_values - returns[inds]) ** 2,
|
||
(
|
||
torch.clamp(
|
||
mb_resp_values,
|
||
old_resp_values[inds] - args.cliprange_value,
|
||
old_resp_values[inds] + args.cliprange_value,
|
||
)
|
||
- returns[inds]
|
||
) ** 2
|
||
)
|
||
* resp_value_mask[inds]
|
||
).sum() / resp_value_mask[inds].sum().clamp(min=1)
|
||
```
|
||
|
||
当前 critic value：
|
||
|
||
$$
|
||
V_{\theta,b,t}
|
||
=
|
||
\texttt{mb\_resp\_values}_{b,t}
|
||
$$
|
||
|
||
旧 critic value：
|
||
|
||
$$
|
||
V_{\mathrm{old},b,t}
|
||
=
|
||
\texttt{old\_resp\_values}_{b,t}
|
||
$$
|
||
|
||
return target：
|
||
|
||
$$
|
||
G_{b,t}
|
||
=
|
||
\texttt{returns}_{b,t}
|
||
$$
|
||
|
||
裁剪后的 value：
|
||
|
||
$$
|
||
\bar{V}_{b,t}
|
||
=
|
||
\operatorname{clip}
|
||
\left(
|
||
V_{\theta,b,t},
|
||
V_{\mathrm{old},b,t} - c_v,
|
||
V_{\mathrm{old},b,t} + c_v
|
||
\right)
|
||
$$
|
||
|
||
其中：
|
||
|
||
$$
|
||
c_v = \texttt{args.cliprange\_value}
|
||
$$
|
||
|
||
两个 value error：
|
||
|
||
$$
|
||
e^{\mathrm{raw}}_{b,t}
|
||
=
|
||
\left(
|
||
V_{\theta,b,t} - G_{b,t}
|
||
\right)^2
|
||
$$
|
||
|
||
$$
|
||
e^{\mathrm{clip}}_{b,t}
|
||
=
|
||
\left(
|
||
\bar{V}_{b,t} - G_{b,t}
|
||
\right)^2
|
||
$$
|
||
|
||
最终 value loss：
|
||
|
||
$$
|
||
\mathcal{L}_{\mathrm{value}}
|
||
=
|
||
\frac{1}{2}
|
||
\operatorname{mean}_m
|
||
\left[
|
||
\max
|
||
\left(
|
||
e^{\mathrm{raw}}_{b,t},
|
||
e^{\mathrm{clip}}_{b,t}
|
||
\right)
|
||
\right]
|
||
$$
|
||
|
||
含义：
|
||
|
||
```text
|
||
critic 要学习预测 returns。
|
||
但源码也限制 critic value 不要一次离 old value 太远。
|
||
```
|
||
|
||
这里的 `0.5` 是 MSE 常见写法，方便梯度形式更简洁。
|
||
|
||
## 8. kl 和 kl_ref 只是日志变量
|
||
|
||
源码：
|
||
|
||
```python
|
||
kl = approx_kl_val
|
||
kl_ref = kl_ref_penalty.detach()
|
||
```
|
||
|
||
含义：
|
||
|
||
```text
|
||
kl:
|
||
记录当前 actor vs old actor 的 approx_kl。
|
||
|
||
kl_ref:
|
||
记录当前 actor vs reference model 的 KL penalty。
|
||
```
|
||
|
||
`.detach()` 表示：
|
||
|
||
```text
|
||
这里只用于记录，不让日志变量继续保留计算图。
|
||
```
|
||
|
||
## 9. stop_ppo 时为什么 loss 乘 0
|
||
|
||
源码：
|
||
|
||
```python
|
||
if stop_ppo:
|
||
loss = (policy_loss + args.vf_coef * value_loss + aux_loss) * 0.0
|
||
else:
|
||
loss = (
|
||
policy_loss
|
||
+ args.vf_coef * value_loss
|
||
+ aux_loss
|
||
) / args.accumulation_steps
|
||
|
||
loss.backward()
|
||
```
|
||
|
||
正常训练时：
|
||
|
||
$$
|
||
\mathcal{L}
|
||
=
|
||
\frac{
|
||
\mathcal{L}_{\mathrm{policy}}
|
||
+
|
||
\lambda_v \mathcal{L}_{\mathrm{value}}
|
||
+
|
||
\mathcal{L}_{\mathrm{aux}}
|
||
}{
|
||
\texttt{accumulation\_steps}
|
||
}
|
||
$$
|
||
|
||
其中：
|
||
|
||
$$
|
||
\lambda_v = \texttt{args.vf\_coef}
|
||
$$
|
||
|
||
early stop 时：
|
||
|
||
$$
|
||
\mathcal{L}_{\mathrm{stop}}
|
||
=
|
||
0
|
||
\cdot
|
||
\left(
|
||
\mathcal{L}_{\mathrm{policy}}
|
||
+
|
||
\lambda_v \mathcal{L}_{\mathrm{value}}
|
||
+
|
||
\mathcal{L}_{\mathrm{aux}}
|
||
\right)
|
||
$$
|
||
|
||
为什么不直接 `continue` 或 `break`？
|
||
|
||
源码注释已经说明：
|
||
|
||
```text
|
||
早停时必须保证 forward-backward 闭环，故只截断 loss 不中断 DDP 通信
|
||
```
|
||
|
||
含义：
|
||
|
||
```text
|
||
分布式训练里，各卡最好执行相同的 forward/backward 结构。
|
||
如果某张卡提前 break，而其它卡还在 backward，可能导致 DDP 等通信卡死。
|
||
```
|
||
|
||
乘 `0.0` 的效果：
|
||
|
||
```text
|
||
仍然执行 backward；
|
||
但梯度为 0；
|
||
不会真正更新参数；
|
||
同时保持 DDP 通信路径完整。
|
||
```
|
||
|
||
## 10. 这段代码的完整执行顺序
|
||
|
||
```text
|
||
1. critic forward:
|
||
得到当前 critic value: mb_resp_values
|
||
|
||
2. actor forward:
|
||
得到当前 actor logits
|
||
|
||
3. gather:
|
||
从 logits 中取 response token 的 mb_resp_logp
|
||
|
||
4. log_ratio:
|
||
当前 actor logprob - old actor logprob
|
||
|
||
5. approx_kl:
|
||
估计当前 actor 和 old actor 的距离，用于 early stop
|
||
|
||
6. ratio:
|
||
exp(log_ratio)
|
||
|
||
7. clipfrac:
|
||
统计多少 token 超出 PPO clip 范围
|
||
|
||
8. kl_ref_penalty:
|
||
惩罚 actor 偏离 reference model
|
||
|
||
9. policy_loss:
|
||
PPO clipped policy loss + reference KL
|
||
|
||
10. value_loss:
|
||
clipped critic value loss
|
||
|
||
11. final loss:
|
||
policy_loss + vf_coef * value_loss + aux_loss
|
||
|
||
12. backward:
|
||
反向传播 actor 和 critic
|
||
```
|
||
|
||
## 11. 最容易混的 5 个点
|
||
|
||
1. `approx_kl` 不是 `kl_ref_penalty`。
|
||
|
||
```text
|
||
approx_kl: 当前 actor vs old actor
|
||
kl_ref_penalty: 当前 actor vs reference model
|
||
```
|
||
|
||
2. `clipfrac` 不是 loss。
|
||
|
||
```text
|
||
它只是监控有多少 token 的 ratio 超过 [1-eps, 1+eps]。
|
||
```
|
||
|
||
3. `policy_loss` 更新 actor。
|
||
|
||
```text
|
||
它使用 new/old logprob、advantage、reference KL。
|
||
```
|
||
|
||
4. `value_loss` 更新 critic。
|
||
|
||
```text
|
||
它使用 current value、old value、returns。
|
||
```
|
||
|
||
5. `stop_ppo` 时仍然 backward。
|
||
|
||
```text
|
||
loss 乘 0 是为了保留 DDP forward-backward 闭环。
|
||
```
|