Diagnosing the training-inference mismatch

Context. In reinforcement-learning (RL) fine-tuning of a language model, an inference enginesoftware built for fast generation, such as vLLM or SGLang; this note calls it the sampler generates the responses and a training enginesoftware that computes losses and gradients, such as PyTorch FSDP or Megatron; this note calls it the trainer computes the gradient, and the two give slightly different probabilities to the same tokens under the same weights because their kernels round differently. This training-inference mismatchthe difference between the token probabilities that the sampler and the trainer assign to the same response under the same weights has been linked to collapsing runs, and earlier work reduced it by moving both engines from BF16bfloat16: a 16-bit floating-point format with 7 mantissa bits, the default for large-model training to FP16IEEE half precision: a 16-bit floating-point format with 10 mantissa bits, so its rounding is 8 times finer than BF16's (Qi et al., 2025) or corrected for it with importance weights. What it does by itself has been hard to say, because PPOproximal policy optimization: a policy-gradient method that takes several gradient steps on each sampled batch and clips each token's probability ratio so the policy moves only a little per batch-style training is off-policy by design, with several gradient steps on each batch, and because clipping and the corrections change how the mismatch acts. Zhong et al. (2026) build VeXactthe authors' sampler for the VeRL framework whose token log-probabilities equal those of the FSDP trainer bit for bit, a sampler whose token log-probabilities equal the trainer's bit for bit, and use it as a baseline with no mismatch at all.

Against this baseline, the small disagreements of a standard sampler are enough on their own to collapse a REINFORCEthe policy gradient $\nabla_\theta\log\pi_\theta(y\mid x)\,A$ applied once per batch, with no ratio and no clipping run, which has no other source of off-policy error. In PPO and GRPOgroup relative policy optimization: PPO with each response's advantage computed from the rewards of the other responses sampled for the same prompt the mismatch enters through the old policy's probability in the denominator of the ratio. If that probability is recomputed by the trainer, clipping measures each update from a distribution that never generated the data, and the advantage-weighted loss terms become skewed long before a KL monitora running estimate of the KL divergence between the current and the old policy, computed from the ratios of sampled tokens and watched for signs of instability moves. If it is taken from the sampler, the reference is right, but the gradient still flows through the trainer's log-probabilities, and reward still falls. Of the four corrections tested that reweight or discard samples after generation, truncating the per-token weight between the two engines together with rejecting responses whose summed mismatch is large comes closest to the exact baseline.

Setting. A prompt $x$ gets a response $y=(y_1,\dots,y_{|y|})$, and $s_t$the prefix $(x,y_1,\dots,y_{t-1})$ before token $t$ is the prefix before token $t$. Write $\mu_{\theta'}$the sampler's policy at the weights $\theta'$ that generated the batch; the paper writes $\pi^{\rm rollout}_{\rm old}$ for the sampler's policy at the weights that generated the batch, $\pi_{\theta'}$the trainer's recomputation of the same policy at the same weights $\theta'$; the paper writes $\pi^{\rm train}_{\rm old}$ for the trainer's recomputation of it, $\pi_\theta$the trainer's policy at the weights $\theta$ being updated; $\theta=\theta'$ before the first gradient step on a batch for the policy being trained, and $\delta_t$the per-token gap $\log\pi_{\theta'}(y_t\mid s_t)-\log\mu_{\theta'}(y_t\mid s_t)$, in nats; 0 for every token with an exact sampler for the per-token gap between trainer and sampler, $\delta_t=\log\pi_{\theta'}(y_t\mid s_t)-\log\mu_{\theta'}(y_t\mid s_t)$. Three ratios follow: the ratio against the sampler, $r_t$$\pi_\theta(y_t\mid s_t)/\mu_{\theta'}(y_t\mid s_t)$, the PPO ratio with the sampler's probability in the denominator (the paper's bypass ratio), the ratio against the recomputation, $r'_t$$\pi_\theta(y_t\mid s_t)/\pi_{\theta'}(y_t\mid s_t)$, the PPO ratio with the trainer's recomputed probability in the denominator, and the correction weight $w_t$$\pi_{\theta'}(y_t\mid s_t)/\mu_{\theta'}(y_t\mid s_t)=e^{\delta_t}$, the ratio between the two engines at the same weights, related by $r_t=r'_t\,w_t$, with $r_t=\pi_\theta/\mu_{\theta'}$, $r'_t=\pi_\theta/\pi_{\theta'}$ and $w_t=\pi_{\theta'}/\mu_{\theta'}$ at $y_t\mid s_t$. With $\hat A$the advantage, shared by every token of a response; in GRPO it is the response's reward minus the mean reward of the responses sampled for the same prompt the advantage and $\varepsilon$the clip range of PPO and GRPO; the ratio is clipped to $[1-\varepsilon,1+\varepsilon]$ the clip range, each token's PPO loss is$$\begin{aligned}&\ell(r,\hat A)=-\min\big(r\,\hat A,\\&\qquad\operatorname{clip}(r,1-\varepsilon,1+\varepsilon)\,\hat A\big),\end{aligned}$$summed over the response and averaged over the batch. Recomputationcomputing the old policy's probabilities with a trainer forward pass at $\theta'$ before updating, and using the loss $\ell(r'_t,\hat A)$ uses $\ell(r'_t,\hat A)$, and bypassusing the log-probabilities the sampler recorded while generating as the old policy's, giving the loss $\ell(r_t,\hat A)$ uses $\ell(r_t,\hat A)$. With an exact sampler $w_t=1$ and the two are the same.

Main findings (all empirical). (a) On Qwen3-8B in BF16, one sampled token, “that” in “The problem states that”, has log-probability $-0.694$ in the sampler and $-0.827$ in the trainer, and the trainer's most likely token there is a different one, a colon followed by two line breaks; the other tokens of the sentence agree to within 0.01. In a GRPO run of Qwen3-1.7B the batch mean of $|\delta_t|$ stays small while the batch maximum reaches about 1. (b) With REINFORCE on Qwen3-30B-A3B, the run with a standard sampler improves and then, after step 280, falls from 0.574 to 0.255 training reward and from 0.293 to 0.067 validation reward on AIME 2024, while the VeXact run keeps improving, to 0.753 and 0.534. (c) With GRPO on Qwen3-1.7B, VeXact holds training reward near 0.93; recomputation falls from about 0.87 to 0.40, partly recovers, and collapses to near zero after step 1,665; bypass falls to about 0.4 and stays there. (d) Under recomputation, KL monitors computed from the ratio stay at the exact baseline's level for 700 steps of falling reward, while the loss contributions of positive- and negative-advantage tokens are already skewed. (e) Truncating $w_t$ and rejecting responses with a large summed mismatch nearly reproduces VeXact's curves.

Consequences. The mismatch is a cause of instability in its own right, worth checking alongside the hyperparameters and the reward when a run collapses. Where possible make the sampler exact, which takes the same model code in both engines and kernels whose results do not depend on batch composition, and the paper releases such a sampler. Otherwise use the sampler's log-probabilities as the old policy, truncate the correction weight $w_t$, and drop responses whose summed mismatch is large; the same conclusion about the reference distribution comes from a trust-region analysis that measures each update in total variation from the sampler's probabilities (Qi et al., 2026). A flat KL monitor computed from the ratio is no evidence of a healthy run under recomputation, since the recomputation ratio $r'_t$ is 1 before the first update whatever the mismatch is; monitor $\delta_t$ itself. An exact sampler is also a reference for tuning the thresholds of the corrections, which otherwise have to be set by trial and error.

Limits. The models are a 1.7B dense model and a 30B mixture-of-experts model, the GRPO experiments use only the 1.7B model, the corrections are compared on one filtered math set, each configuration is run once, and the curves comparing corrections are smoothed over 25 steps. The account of how recomputation skews the loss is a hypothesis supported by one diagnostic plot. That the corrections track the exact baseline is shown for this setting and these thresholds; the authors do not claim it generalizes, and they note the corrections may have side effects their experiments do not show. VeXact's speed relative to vLLM, which decides whether an exact sampler is practical at scale, is not reported.