REINFORCE → Actor-Critic → TRPO → PPO → GRPO → Train–Inference Mismatch → Score Centering
Yichuan Wang · yichuan-w.github.io/blog/GDN-train-inference-mismatch-asyncRL
Running example adapted from Tyler Romero’s Policy Gradient for LLMs / Off-Policy RL for LLMs (2026)
← → to navigate · N for notes · F for fullscreen
Prompt: “17 × 24 = ?”. After writing “340 +”, the next token is 68 (correct, prob. 0.82) or 58 (slip, prob. 0.07).
$\pi_\theta$: the policy (LLM) · $R\in\{0,1\}$: verifier reward
$A$: advantage, “how much better than expected”
“… = 340 + 68 = 408” → $R=1$
“… = 340 + 58 = 398” → $R=0$
Written out, the objective is a sum over every possible completion:
Goal: rewrite $\nabla_\theta J$ as an average over samples of something we can differentiate.
Chain rule: $\nabla\log\pi=\nabla\pi/\pi$. Multiply both sides by $\pi$.
REINFORCE: sample N, grade them, average “score × reward”. Unbiased but noisy.
A special case of the policy gradient theorem — Sutton, McAllester, Singh & Mansour (NeurIPS 1999); estimator: REINFORCE (Williams, 1992).
4 samples for “17 × 24 = ?”:
Correct answers get pushed up. Probabilities sum to 1, so the wrong ones lose share.
An easy prompt the model always solves — 4 correct, each worded differently:
Fix: train a critic $V_\phi(s_t)$ — expected reward from this prefix — and subtract it as a baseline:
$V \approx 0.85$ → pick 68: $A\approx+0.15$, pick 58: $A\approx-0.85$
The slip is pushed down where it happened; an always-solved prompt gives $A\approx0$: no noise. (illustrative numbers)
Fix a prefix $s_t$. There $V(s_t)$ is one number; the next token $a$ is sampled from $\pi(\cdot\mid s_t)$ (we drop $s_t$ below).
Subtracting $V$ changes each sample, not the average: same expected gradient, less noise. Needs $a\sim\pi$ itself — remember this for part 6.
For LLMs the critic is usually as big as the policy, and hard to train well.
$\mu$: the policy that generated the rollout. At the first update $\pi=\mu$, so $r$ should be exactly 1 — remember this for part 5.
PPO’s question: can plain Adam + one simple trick get most of this?
Sample $G=4$ answers for one prompt:
Follow-ups (Dr. GRPO, DAPO) fix details: std / length normalization, all-correct groups, clip-higher.
$r$ is an exponential of a difference: an error of 0.1 in log-prob is ~10% in $r$ — and PPO’s clip is only ±0.2. In MoE RL, router mismatch alone can collapse training (R3).
They don’t match. Why can two systems with the same weights and the same tokens produce different log-probs? The answer is three layers down:
In math: $(a+b)+c=a+(b+c)$. In floating point, $\mathrm{{fl}}(\cdot)$ rounds to a representable value after every operation:
Every rounding step throws away bits. Which bits depends on what you added first — the error is a function of the order.
bf16 — what we train in:
Same three numbers — the two orders discard different terms, so the results differ.
Here extra terms are simply dropped. GPUs round to nearest-even — same effect, e.g. $a=127.5,\ b=1$: $(a+b)+b=129$, $a+(b+b)=130$.
Every linear layer is a matmul; every output number is one big sum:
Qwen3.5-9B: 4096 additions per output number.
Both are correct; neither is more correct. A heuristic picks one at launch time, from the shapes. So the reduction order is a property of the shape, not of the math.
4096×4096 fp32 matmul, one B300. Row 0’s input bytes are identical in every row of this table.
| batch M | max |diff| vs M=1 | cuBLAS strategy |
|---|---|---|
| 1 | 0 | split-K |
| 2 | 384 | split-K |
| 8 | 384 | split-K |
| 128 | 416 | split-K |
| 256 | 1249.25 | no split-K |
| 2048 | 1249.25 | no split-K |
No randomness: rerun it and every number reproduces exactly. cuBLAS is deterministic given a shape — the shape just picks the algorithm.
You can watch cuBLAS decide (torch.profiler, kernel names):
splitKreduce only adds up the partial sums (the tree) — gone at M ≥ 256.
At the model level — Qwen3.5-9B, bf16, same prompt repeated in the batch:
An 18% error in $r$ — from nothing but batch composition. PPO’s clip range is ±0.2.
Trainer batch: [ x, train sample 1, train sample 2, … ]
Rollout batch: [ x, request 17, request 42, request 93, … ]
They never see the same batch → determinism doesn’t help.
It always comes down to one thing: a reduction along the batch axis. Thinking Machines’ batch_invariant_ops:
| Op | Why batch-variant | Fix |
|---|---|---|
| GEMM | tile size and split-K picked per shape | fixed partition |
| RMSNorm | few rows → reduction split across the batch axis | don’t reduce along batch |
| attention | split-KV count depends on the batch | num_splits=1 |
DP / FSDP add no mismatch — they don’t reduce along the batch axis.
Batch-invariant kernels make $f_A(x)$ consistent with itself — but we need $f_A=f_B$. If the two sides implement the forward differently, aligning the kernels underneath is aligning the wrong layer.
radixark/miles#975: Qwen3.5’s A_log was rounded to bf16 on the Megatron side while SGLang kept fp32. Fix: keep it fp32 through Megatron’s Float16Module wrap. Found by hand — per op, per model, per release.
Async RL loops (slime, open-instruct) compute $r$ from raw $\mu$ and $\pi$ — they only see the sum. When a run breaks, you can’t tell async from infra.
| Layer | Fix |
|---|---|
| ① Floating point | Can’t change the arithmetic → fix the order, same dtype on both sides |
| ② Kernel | Batch-invariant kernels (the table two slides back) |
| ③ Engine | One model definition: vLLM keeps its KV cache and scheduling, every layer is the trainer’s op |
| New in Qwen3.5 | Gated DeltaNet ships chunked (train) and recurrent (decode) kernels — two algorithms. → recurrent on every forward, chunked for backward |
Result: $\log\pi=\log\mu$ — not small, exactly zero, bit for bit. Parity only needs the forward to match ($\pi$, $\mu$ are forward quantities), so backward stays fast.
red vLLM native (two engines) · brown unified model, stock kernels · grey unified + BI kernels
Step 0: exactly 0. Grey keeps the lowest gap and reaches the highest train reward (and held-out peak).
Unified model is free — always do it. BI costs 2–3× (≈5× on agents) → turn it on for 20 steps to debug: gap = 0 and still broken → not infra.
Constant reward $+1$: no signal, so there should be no learning. On-policy that holds at every prefix:
Under mismatch the token is sampled from $\mu$ but the score goes through $\pi$, so $\bar s=\mathbb{E}_\mu[s_{{y_t}}]\neq0$. With $\mathbb{E}[AB]=\mathbb{E}[A]\,\mathbb{E}[B]+\mathrm{{Cov}}(A,B)$:
Only the signal sees which token led to which reward. The drift depends on the rewards only through their mean.
On-policy ($\mu=\pi$): $\bar s=0$, no drift. Group centering shrinks it but can’t remove it: the expected advantage at a prefix isn’t 0.
$y_t$: sampled token · $s_v=\nabla_\theta\log\pi_v$: score of token $v$ (through the trainer) · $\pi_v$ / $\mu_v$: trainer / sampler prob. · $\bar s$: expected score under the sampler · $\mathrm{{Cov}}_\mu$: covariance with $y_t$ and the rest of the rollout drawn from $\mu$
| Importance sampling | Score centering | |
|---|---|---|
| Correction | × ratio $\pi/\mu$ | − $\bar s$ (additive) |
| Depends on token? | yes, huge on rare tokens | no, fixed per prefix |
| In practice | clipped → drift leaks back | exact, no hyperparameters |
They compose (SC on top of TIS / MIS). Qwen3-30B-A3B, INT8 + INT4-KV sampler: SC 30%, TIS 12%, others < 5%.
$\bar s$ needs the sampler’s full distribution: 152K vocab × fp32 = 608 KB per token → 1024 rollouts × 32K tokens = 20 TB. So log only the sampler’s top-$k$ ($k=128$) as the head $H$, and fill the tail with the trainer, rescaled:
Same $\bar s$ as before; only the tail is approximated, $\mu_v\approx\rho\,\pi_v$ — and it collapses to only $k$ terms:
$k=128$ and even $k=32$ match full SC; <1% wall-clock overhead. vLLM / SGLang already return top-$k$ logprobs.
As a scalar loss (sg = stop-gradient):
Adapted from the paper’s JAX snippet (Sec. 4.1). Differentiating gives $A\,(s_{y_t}-\bar s)$. $\pi_v$ / $\mu_v$: trainer / sampler prob. of token $v$ · $y_t$: sampled token · $A$: advantage · $H$: sampler’s top-$k$ tokens · sg: stop-gradient.
Subtracting the expected score is an obvious idea — why wasn’t it used before?
“RL training of LLMs under TIM is unusual in that the expected score is both nonzero and computable exactly.” — Marek & Ryabinin (2026)
Rollouts ∥ training · buffer bounds off-policyness · TorchStore weight sync
In-place update: pause → load $w_{{k+1}}$ → resume, KV kept → one sequence, two weight versions
Theory. Signal is $\mathrm{{Cov}}_\mu$, not $\mathrm{{Cov}}_\pi$ — under staleness, a bias nobody has bounded.
Practice. Only test: 0.6B, sampler synced every 64 steps — no real async pipeline; SC alone < SC + TIS.
Key questions. ① How many async (off-policy) steps can SC tolerate? ② Can SC replace IS in async RL — removing the bias that clipped IS introduces?
| Method | Problem it fixes | Core idea |
|---|---|---|
| REINFORCE | can’t differentiate sampling | $R\,\nabla\log\pi$ |
| Actor-Critic | noise, credit assignment | critic baseline $R-V(s_t)$ |
| TRPO | steps that go too far | ratio $r$ + KL trust region |
| PPO | TRPO is too expensive | clip $r$ to $[0.8,\,1.2]$ |
| GRPO | critic is expensive | group mean as baseline |
Fix all three layers (floating point → kernel → engine) → bitwise parity. Exact, but 2–3× slower → use it to debug.
IS (TIS / MIS) or score centering: subtract $\mathbb{E}_\mu[\nabla\log\pi]$. Cheap (log top-k probs). + async: still open.
A_log in fp32 (2026) · W. Ma et al. Rollout Routing Replay (R3) (2025) · F. Yao et al. TIS (2025)