minimind/course/notes/23_ppo_update_loss_block.md

PPO update loss block 源码解析

本文只解释 MiniMind PPO 训练中的这一段：

trainer/train_ppo.py:171-214

它完成一次 PPO mini-batch update 的核心计算：

当前 actor logprob
-> ratio
-> clipfrac
-> reference KL penalty
-> clipped policy loss
-> clipped value loss
-> 合成最终 loss
-> backward

0. 变量约定

这一节先按源码变量读。数学符号只用来把 [batch, response_len] 里的某个元素写清楚。

0.1 下标怎么读

b = batch 里的第几条样本
t = response 里的第几个 token

所以：

x[b, t]

就是：

第 b 条样本，第 t 个 response token 位置上的值。

0.2 action 和 state

第 b 条样本第 t 个 response token 记作：

a_{b,t}

它对应源码里的：

completion_ids[b, t]

生成这个 token 时，模型看到的上下文记作：

s_{b,t}

它表示：

prompt + response 前 t-1 个 token

也就是说，PPO 里的：

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})

对应源码：

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})

对应源码：

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})

对应源码：

ref_resp_logp[b, t]

0.4 其它源码变量

有效 response token mask：

resp_policy_mask[b, t]

公式里写成：

m_{b,t}

GAE 算出的 advantage：

advantages[b, t]

公式里写成：

A_{b,t}

当前 critic value：

mb_resp_values[b, t]

公式里写成：

V_{\theta,b,t}

rollout 阶段保存的 old critic value：

old_resp_values[b, t]

公式里写成：

V_{\mathrm{old},b,t}

critic 的训练目标 return：

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

源码：

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})

含义：

rollout 阶段已经保存了 old actor 的 old_resp_logp。
现在 PPO update 阶段，要用当前 actor 对同一批 response token 重新算 logprob。

注意：

这里没有重新生成 response。
它是在同一批 old rollout 上重新 forward，计算当前 actor 对这些 token 的概率。

2. log_ratio 和 ratio

源码：

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})
}

含义：

ratio 表示当前 actor 相对 rollout 时 old actor，对同一个 token 的概率改了多少。

判断方式：

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

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

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

3. approx_kl：监控 actor 是否离 old actor 太远

源码：

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)

用途：

approx_kl 不直接作为 policy_loss 的主项。
它主要用于 early stop。

源码：

if approx_kl_val > args.early_stop_kl:
    stop_ppo = True

含义：

如果当前 actor 已经离 old actor 太远，就停止继续用这批 rollout 更新。

4. clipfrac：统计有多少 token 超出 PPO clip 范围

源码：

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}

含义：

clipfrac 不是 loss。
它是一个监控指标，表示有多少有效 response token 的 ratio 已经超过 PPO 裁剪范围。

例子：

clip_epsilon = 0.2
ratio 合理范围 = [0.8, 1.2]

ratio = 1.05 -> 不计入 clipfrac
ratio = 1.35 -> 计入 clipfrac
ratio = 0.70 -> 计入 clipfrac

如果 clipfrac 很大，说明：

这批 update 里很多 token 的策略变化已经被 PPO clip 限制住了。

5. reference KL penalty：约束 actor 不要偏离 reference

源码：

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)

含义：

reference model 是 frozen SFT 模型。
这个项惩罚当前 actor 偏离 reference model。

注意它和 approx_kl 不同：

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

源码：

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}

直观理解：

advantage > 0:
  这个 token 比 critic 预期更好，希望提高它的概率。
  但 ratio 不能无限变大，最多按 1 + epsilon 裁剪。

advantage < 0:
  这个 token 比 critic 预期更差，希望降低它的概率。
  但 ratio 不能无限变小，最多按 1 - epsilon 裁剪。

所以 policy_loss 同时做两件事：

1. 按 advantage 更新 actor。
2. 用 reference KL 约束 actor 不要偏离 SFT reference 太远。

7. value_loss：critic 的 clipped value loss

源码：

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]

含义：

critic 要学习预测 returns。
但源码也限制 critic value 不要一次离 old value 太远。

这里的 0.5 是 MSE 常见写法，方便梯度形式更简洁。

8. kl 和 kl_ref 只是日志变量

源码：

kl = approx_kl_val
kl_ref = kl_ref_penalty.detach()

含义：

kl:
  记录当前 actor vs old actor 的 approx_kl。

kl_ref:
  记录当前 actor vs reference model 的 KL penalty。

.detach() 表示：

这里只用于记录，不让日志变量继续保留计算图。

9. stop_ppo 时为什么 loss 乘 0

源码：

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？

源码注释已经说明：

早停时必须保证 forward-backward 闭环，故只截断 loss 不中断 DDP 通信

含义：

分布式训练里，各卡最好执行相同的 forward/backward 结构。
如果某张卡提前 break，而其它卡还在 backward，可能导致 DDP 等通信卡死。

乘 0.0 的效果：

仍然执行 backward；
但梯度为 0；
不会真正更新参数；
同时保持 DDP 通信路径完整。

10. 这段代码的完整执行顺序

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。

approx_kl: 当前 actor vs old actor
kl_ref_penalty: 当前 actor vs reference model

2. clipfrac 不是 loss。

它只是监控有多少 token 的 ratio 超过 [1-eps, 1+eps]。

3. policy_loss 更新 actor。

它使用 new/old logprob、advantage、reference KL。

4. value_loss 更新 critic。

它使用 current value、old value、returns。

5. stop_ppo 时仍然 backward。

loss 乘 0 是为了保留 DDP forward-backward 闭环。