从策略梯度到训推一致性

从策略梯度到训推一致性#

返回实验室。精度、确定性、异步、PD 和 HiCache 见系统篇。

下面按照仓库实现推导各个目标，区分数学恒等式、有偏 surrogate 和工程近似。论文链接标注原始工作；实现链接说明向导具体导出的变体。

统一符号#

\(x\) 表示 prompt，\(y=(a_1,\ldots,a_T)\) 表示回答，\(h_t=(x,a_{<t})\) 表示 prefix。奖励 \(R(x,y)\) 在求导时不依赖参数。必须区分三种策略：

符号

含义

来源

\(q(a\mid h)\)

真正产生 token 的行为分布

SGLang，包含量化、采样变换以及可能过期的权重

\(p_0(a\mid h)\)

更新前冻结的训练策略

Megatron 重新计算的 old log-probs

\(p_\theta(a\mid h)\)

当前可训练策略

Megatron 的可微 forward

\(\operatorname{sg}(z)\)

数值为 \(z\)、梯度为零

对权重、奖励和修正系数停止梯度

固定 prefix，定义 score 为 \(s_a=\nabla_\theta\log p_\theta(a\mid h)\)。在概率可微、归一化且可交换求和与求导的前提下：

\[ \sum_a p_\theta(a\mid h)s_a =\sum_a\nabla_\theta p_\theta(a\mid h) =\nabla_\theta 1=0. \]

从 \(J(\theta)=\sum_y p_\theta(y\mid x)R(x,y)\) 出发，应用 \(\nabla p=p\nabla\log p\) 和自回归分解：

\[ \nabla J =\mathbb E_{y\sim p_\theta}[R\nabla\log p_\theta(y\mid x)] =\mathbb E_{y\sim p_\theta}\left[R\sum_t s_{a_t}\right]. \]

在 \(p_\theta\) 分布下，与当前 action 无关的 baseline \(b(h_t)\) 和 score 相乘后的期望为零，因此可以降低方差而不改变这个期望。换成 \(q\) 采样，一般就不能这样消去。两个引擎加载了同名 checkpoint，并不证明训练是 on-policy。

GRPO#

原论文：DeepSeekMath；裁剪目标：PPO。

对同一 prompt 采样 \(G\) 个回答，用组内统计构造 advantage：

\[ \bar R=\frac1G\sum_i R_i,\quad \sigma_R=\sqrt{\frac1G\sum_i(R_i-\bar R)^2},\quad A_i=\frac{R_i-\bar R}{\sigma_R+\epsilon}. \]

同一回答中的 response token 共用这个 advantage。标准差是否使用样本修正以及具体 epsilon 取决于实现；关闭标准差归一化后，\(A_i=R_i-\bar R\)。组均值包含当前样本，不能直接套用“独立 baseline 无偏”的证明。独立采样时，仅减均值会让 on-policy 梯度期望缩放为 \(1-1/G\)；再除标准差还引入随机尺度。

令 \(r_{i,t}(\theta)=p_\theta(a_{i,t}\mid h_{i,t})/p_0(a_{i,t}\mid h_{i,t})\)。忽略可选 KL 项，要最大化的 clipped surrogate 为：

\[ J_{\rm clip}=\frac1G\sum_i\frac1{T_i}\sum_t \min\left(r_{i,t}A_i, \operatorname{clip}(r_{i,t},1-\epsilon_-,1+\epsilon_+)A_i\right). \]

\(A_i>0\) 时，上界限制继续提高 token 概率的收益；\(A_i<0\) 时，下界限制继续降低概率的收益。在 \(\theta=\theta_0\) 且未触发 clipping 时，\(\nabla r=r\nabla\log p_\theta\) 还原为策略梯度方向。这是局部 surrogate，不是任意过期数据上的恒等式。

向导使用 --advantage-estimator grpo、每个 prompt 8 次采样；不开 SC 时采用 PPO clipping。--calculate-per-token-loss 将聚合方式改成“总 loss / 总 response token 数”，与“先求每个回答均值，再平均回答”并不相同。

TIS#

原始工作：On the Rollout-Training Mismatch in Modern RL Systems 和作者技术文章。下面的代数也对应仓库内置 callback。

固定 prefix，只要 \(p_0(a)f(a)\ne0\) 的位置都有 \(q(a)>0\)，就有换测度恒等式：

\[ \mathbb E_{p_0}[f(a)] =\sum_a q(a)\frac{p_0(a)}{q(a)}f(a) =\mathbb E_q[\rho_a f(a)],\quad \rho_a=\exp(\log p_0(a)-\log q(a)). \]

PPO 的更新比值 \(r=p_\theta/p_0\) 与行为修正比值 \(\rho=p_0/q\) 负责不同问题。未裁剪时两者乘积为 \(p_\theta/q\)。slime 将 PPO loss 乘上停止梯度的 \(w(\rho)=\operatorname{clip}(\rho,L,U)\)，向导取 \(L=0,U=2\)。

减去精确表达式，就能直接得到偏差：

\[ \mathbb E_q[w(\rho)f]-\mathbb E_{p_0}[f] =\mathbb E_q[(w(\rho)-\rho)f]. \]

\(L=0\) 时，偏差为 \(-\mathbb E_q[(\rho-U)f\mathbf1_{\rho>U}]\)。限制权重能抑制离群值，却不能保持无偏；\(L>0\) 时还会抬高小权重。

这个恒等式固定了 prefix，没有修正 prefix 本身的分布。完整序列的换测度需要 \(p_0(y)/q(y)=\prod_t\rho_t\)，方差可能很大。逐 token TIS 是实用 surrogate，不能宣称等价于完整乘积。

实现：vanilla_tis_function、importance_weights。REINFORCE（包括 SC）使用停止梯度的当前训练策略作为修正分子，而不是 \(p_0\)。

ICE-POP#

原论文：Every Step Evolves: Scaling Reinforcement Learning for Trillion-Scale Thinking Model。

向导选择 slime 的 icepop_function，其行为修正规则是：

\[ w_{\rm ice}(\rho)=\rho\mathbf1_{L\le\rho\le U},\qquad L=0.5,\quad U=2. \]

区间内保留重要性权重，区间外贡献为零。因此相对于同一个固定 prefix 的目标：

\[ \mathbb E_q[w_{\rm ice}(\rho)f]-\mathbb E_{p_0}[f] =-\mathbb E_q[\rho f\mathbf1_{\rho<L\lor\rho>U}]. \]

例如 \(\rho=4\) 时，向导中的 TIS 给权重 2，IcePop 给权重 0。两者不是同一种估计器。Callback 将权重置零，不会按保留下来的 token 数重新归一化 loss；PPO clipping 仍可以独立作用于 \(r\)。论文完整训练流程还有其他部分，启用这个 callback 并不等于复现整篇论文。

导出参数为 --use-tis --custom-tis-function-path slime.backends.megatron_utils.loss.icepop_function --tis-clip-low 0.5 --tis-clip 2。观察 tis、tis_clipfrac、reward 以及 mask 后是否还有足够的学习信号。

R3#

原论文：Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers。

把 MoE 层写为 \(F(x)=\sum_{i\in I(x)}g_i(x)E_i(x)\)，其中 \(I=\operatorname{TopK}(u(x))\) 是离散 expert 集合。Top-k 边界附近，微小扰动也可能换掉一个专家。加减同一个中间项，可将输出差异分解为：

\[ F_{I_q}(x_q)-F_{I_p}(x_p) =\underbrace{F_{I_q}(x_q)-F_{I_q}(x_p)}_{\text{相同路由下的数值与输入差异}} +\underbrace{F_{I_q}(x_p)-F_{I_p}(x_p)}_{\text{专家选择差异}}. \]

训练时固定为 rollout 记录的 \(I_q\)，就能在此层、此比较中消去第二项，但第一项仍然存在。以 softmax gate 为例：

\[ g_i=\frac{\mathbf1_{i\in I_q}\exp u_i}{\sum_{j\in I_q}\exp u_j},\qquad \frac{\partial g_i}{\partial u_j}=g_i(\mathbf1_{i=j}-g_j),\quad i,j\in I_q. \]

重放的是专家身份，不是 detach 的专家输出；梯度仍可流过选中的 gate 和 expert。GLM 使用 sigmoid 路由及自己的归一化与缩放，以上 softmax 只是示例，实际必须保留模型的 gate 规则。

向导添加 --use-rollout-routing-replay。slime 管理的 SGLang 自动返回路由；external 引擎需自行启用 --enable-return-routed-experts。存储随 token 数 × MoE 层数 × 选中专家数增长。R3 可独立组合 TIS/SC，对 dense 模型没有意义。实现见 sglang_engine.py 和 megatron_utils 的 replay 逻辑。

SC#

原论文：Score Centering Stabilizes Off-policy Reinforcement Learning。以下按照 slime 的 score_centering_correction 和 loss 实现展开推导。

从非零均值到中心化 score#

固定 prefix，令 \(\mu=\sum_a q_as_a\)。即使两侧 checkpoint 相同，一般仍有 \(\mu\ne0\)。对停止梯度的 advantage \(A\)，应用协方差分解：

\[ \mathbb E_q[As]=\mathbb E_q[A]\mu+\operatorname{Cov}_q(A,s). \]

令 \(\tilde s_a=s_a-\mu\)，则 \(\mathbb E_q[\tilde s]=0\)，因此：

\[ \mathbb E_q[A\tilde s]=\operatorname{Cov}_q(A,s). \]

Prefix 条件下的均值项被消去了。但协方差仍在 \(q\) 下计算，这并不证明得到了 \(p_\theta\) 下的无偏梯度。对完整回答做组内 reward centering，也不保证每个 prefix 上都有 \(\mathbb E[A\mid h]=0\)。

与重要性修正组合#

令 \(w_a=f(p_a/q_a)\)，\(f\) 可以是原始比值、TIS、mask 或常数 1。定义 \(\mu_w=\sum_aq_aw_as_a\)，对加权 score进行中心化：

\[ \tilde s_a^w=w_as_a-\mu_w,\qquad \mathbb E_q[\tilde s^w]=0. \]

实现这个局部梯度方向的一种最小化 loss 是：

\[ \ell=-\operatorname{sg}(A)\left[ \operatorname{sg}(w_a)\log p_a- \sum_v\operatorname{sg}(q_vw_v)\log p_v\right]. \]

只对 \(\log p\) 求导，得到 \(-\nabla\ell=A(w_as_a-\mu_w)\)。如果让梯度流过 \(w\)，就会出现额外项。完整 support 且 \(w=p/q\) 时，\(\mu_w=\sum p s=0\)；经过裁剪或 mask 后通常不再为零。

Top-k 如何避免保存整个词表#

slime 保存 sampler 的 head \(H\)，向导默认 128 个 token。这些概率保留完整词表上的质量，不按 head 重新归一化。定义：

\[\begin{split} P_H=\sum_{v\in H}p_v,\quad Q_H=\sum_{v\in H}q_v,\quad \alpha=\frac{1-Q_H}{1-P_H},\quad \hat q_v=\begin{cases}q_v&v\in H\\\alpha p_v&v\notin H.\end{cases} \end{split}\]

该近似的概率和为一。利用 \(\sum_vp_vs_v=0\)，把 tail 项替换掉：

\[ \mu_{\hat q}=\sum_Hq_vs_v+\alpha\sum_{v\notin H}p_vs_v =\sum_H(q_v-\alpha p_v)s_v. \]

Tail 中 \(p_v/\hat q_v=1/\alpha\)，因此加权形式为：

\[ \mu_{\hat q,w}=\sum_{v\in H} [q_vf(p_v/q_v)-\alpha f(1/\alpha)p_v]s_v. \]

这正是代码中的 head residual。完整分布的中心化是精确的，使用重建 tail 时则是近似；很小的 tail mass 还需要数值保护。可以用一个 logit 小例子验证：\(p=(0.6,0.4)\)、\(q=(0.5,0.5)\)、\(s_a=e_a-p\)，未中心化的均值为 \(q-p=(-0.1,0.1)\)；减去它之后 score 的期望为零。

必须匹配真实采样分布#

向导固定 temperature=1、top-p=1、top-k=-1，并使用 --pg-loss-type reinforce、--disable-grpo-std-normalization、--calculate-per-token-loss。不能把 SC 直接塞进 GSPO / CISPO / PPO clipping；此路径不支持 streaming、约束解码、采样惩罚和逐请求修改 temperature。

slime 还支持记录 top-p support \(S\)：将训练概率也归一化到 \(S\)，使用所有记录的 \(q_v\)。此时 \(\mu_w=\sum_{v\in S}q_vw_v\nabla\log p_v^S\) 无需 tail 近似。这只相对于记录的 support 精确，不等同于未截断策略；数据量可能远大于 top-128。详见采样与兼容条件。

如何验证选择#

固定 prompt 数据、seed、回复长度预算和评估集。从短基线开始，每次改一个机制，同时记录 reward、train_rollout_logprob_abs_diff、适用时的 tis_clipfrac，以及异步的 staleness/mean、staleness/max。差异更小不自动意味着学得更好。保存固定 batch 做 debug replay，将优化器行为与采样变化分开定位。