Diagnosing Training Inference Mismatch in LLM Reinforcement Learning

Authors: Tianle Zhong, Neiwen Ling, Yifan Pi, Zijun Wei, Tianshu Yu, Geoffrey Fox, Peng Wu, Xiao Yu (ByteDance, University of Virginia)

---

Summary (Overview)

• Identifies Training–Inference Mismatch (TIM) as a critical infrastructure-level confounder in LLM reinforcement learning: implementation differences between training engines (FSDP) and inference engines (vLLM) cause divergent token probabilities for identical inputs and weights, which can independently trigger training collapse.
• Introduces VeXact, a lightweight zero-mismatch rollout engine that achieves bit-wise alignment with the FSDP training engine by unifying kernel implementations and employing deterministic, batch-invariant kernels—providing a TIM-free diagnostic baseline.
• Demonstrates that TIM alone destabilizes RL training: In REINFORCE experiments, vLLM-based rollout shows reward collapse (MoE: 0.574→0.255 training reward) while VeXact remains stable (0.753), isolating TIM as the sole causal factor.
• Analyzes failure modes of recomputation vs. bypass strategies in GRPO, showing that TIM changes the effective optimization objective through sign-imbalanced, advantage-weighted loss contributions rather than uniform noise—explaining why KL estimators fail to detect early-stage instability.
• Ablates algorithmic TIM compensations (truncated importance sampling, rejection sampling), finding that a combination of token-level truncation with $r_{corr}$-based sequence-level rejection can closely track the zero-mismatch baseline.

---

Introduction and Theoretical Foundation

Background

LLM reinforcement learning has become central to post-training foundation models, yet RL training remains difficult to stabilize—policies may rapidly degrade over short training windows. Diagnosing collapses is complicated because multiple failure modes are entangled: hyperparameter issues, reward misspecification, reward hacking, and infrastructure-level confounders.

The TIM Problem

Training-Inference Mismatch (TIM) arises from implementation differences between training engines (FSDP, Megatron) and inference engines (vLLM, SGLang). Even with identical model checkpoints and inputs, divergent model/kernel implementations produce different token probability distributions. This introduces an unintended off-policy bias between sampling and model update—distinct from PPO mini-step off-policy drift, as it cannot be addressed by naive PPO clipping.

Formal Definition

Given context $x$ and sampled response $y = (a_1, \ldots, a_T)$ with state $s_t = (x, y_{<t})$, TIM is defined at the token level as:

$$\delta_{t} = \log \pi_{old}^{\mathrm{train}}(a_{t} | s_{t}) - \log \pi_{old}^{\mathrm{rollout}}(a_{t} | s_{t}) \tag{1}$$

where $\pi_{old}^{\mathrm{train}}$ is the trainer-side reference distribution and $\pi_{old}^{\mathrm{rollout}}$ is the behavioral distribution realized by the rollout engine. This definition is objective-agnostic (applies to REINFORCE, PPO, or GRPO).

Key Empirical Observation

Table 1 shows token-level drift between rollout and training stacks on Qwen3-8B (bf16) weights:

† marks an argmax flip: training side's top-1 token differs from rollout's, potentially ending a clause differently. While mean $|\delta_t|$ is small per batch, max $|\delta_t|$ can reach ~1.0 for extreme tokens.

---

Methodology

VeXact: Zero-Mismatch Rollout Engine

VeXact eliminates TIM from two sources:

1. Model/kernel implementation differences: Uses the same HuggingFace-based model implementation as FSDP, registering VeXact kernels in the FSDP engine initialization.
2. Kernel reduction order and tiling variations: Employs deterministic, batch-invariant kernels that fix tiling and reduction order. Implements batch-invariant versions of:
   • RMSNorm
   • Batched matrix multiplication
   • Fused MoE kernels
   • Attention (with KV splitting disabled for determinism)

To maintain practical throughput despite fixed tiling, VeXact integrates chunked prefill, CUDAGraph, pipeline parallelism, and optimistic KV allocation with preemption fallback. Its codebase is lightweight (similar LOC to nano-vLLM).

Experimental Configurations

PPO/GRPO Surrogate Objectives

The token-level clipped surrogate is:

$$\mathcal{L}_{\mathrm{ppo}}(r_{ppo}, A) = -\min\left(r_{ppo} A, \mathrm{clip}(r_{ppo}, 1-\epsilon, 1+\epsilon) A\right) \tag{2}$$

Under recomputation vs. bypass, the PPO ratios differ in their denominators:

$$r_{ppo}^{\mathrm{train}} = \frac{\pi_{\theta}(a_{t} \mid s_{t})}{\pi_{old}^{\mathrm{train}}(a_{t} \mid s_{t})}, \quad r_{ppo}^{\mathrm{rollout}} = \frac{\pi_{\theta}(a_{t} \mid s_{t})}{\pi_{old}^{\mathrm{rollout}}(a_{t} \mid s_{t})} \tag{3}$$

Zero-Centered Loss Contribution

To analyze objective-space distortion:

$$C(r_{ppo}) = -(r_{ppo} - 1) A_{t}, \quad r_{ppo} \in \{r_{ppo}^{\mathrm{train}}, r_{ppo}^{\mathrm{rollout}}\} \tag{5}$$

This has the same gradient as the standard $-r_{ppo}A_t$ objective but is zero when $r_{ppo} = 1$.

---

Empirical Validation / Results

REINFORCE: TIM Alone Causes Collapse

In Figure 2, with TIM as the only difference:

• MoE (Qwen3-30B-A3B): vLLM run degrades after step 280 (training reward 0.574→0.255, validation 0.293→0.067); VeXact reaches 0.753 training / 0.534 validation reward.
• Dense (Qwen3-1.7B): Similar instability pattern with vLLM; VeXact remains stable.

GRPO: Failure Modes of Recomputation and Bypass

Figure 3 shows:

• VeXact: Stable, training reward ~0.93.
• vLLM recomputation: Degrades from ~0.87 to ~0.40 (first 650 steps), partially recovers, then collapses to near-zero after ~step 1665.
• vLLM bypass: Single-stage degradation to ~0.4 without collapse; not accompanied by comparably large loss spikes.

KL Estimators Fail to Detect Early Instability

KL estimators $K_1(r_{ppo}) = -\log r_{ppo}$ and $K_3(r_{ppo}) = (r_{ppo} - 1) - \log r_{ppo}$:

• Bypass mode: Both $K_1$ and $K_3$ increase noticeably.
• Recomputation mode: Both remain nearly flat during the first 700 steps even as reward degrades—aggregate probability-space divergence is not sufficient to characterize early TIM-induced failure.

Sign-Imbalanced Loss Contributions

Figure 5 reveals that recomputation induces a sign-dependent skew in advantage-weighted update signals:

• Harmful contributions are amplified, while offsetting contributions are not symmetrically amplified.
• Symmetric numerical noise in $\delta_t$ is transformed into skewed, non-zero-mean gradient distortion due to asymmetric interaction with PPO clipping bounds (which react differently for positive vs. negative $A_t$).

Why Bypass Also Fails

In bypass mode, the PPO ratio correctly uses the behavioral distribution in the denominator, but the numerator $\pi_{\theta}^{train}$ is evaluated on the trainer's numerical path. The optimizer exploits numerical artifacts in the trainer's forward pass; weight updates fail to translate into behavioral improvements when deployed back to the rollout engine.

Algorithmic TIM Compensation Results

Three correction configurations were evaluated (Figure 6):

Key findings:

1. $r_{corr}$-based rejection outperforms $r_{ppo}$-based: $r_{corr} = \pi_{old}^{train}/\pi_{old}^{rollout}$ measures system-induced mismatch, while $r_{ppo}$ overlaps with PPO's policy-ratio mechanism (controlling update magnitude, not starting distribution location).
2. TIS is effective: It corrects the PPO ratio from $r_{ppo}^{train}$ to $r_{ppo}^{rollout}$ by multiplying with $r_{corr}$.
3. Combined TIS + sequence rejection tracks VeXact closely, indicating TIM manifests at multiple granularities.
4. The choice between $K_1$ and $K_3$ for sequence rejection has minor impact.

---

Theoretical and Practical Implications

TIM Changes the Effective Optimization Problem

• Recomputation mode: The sampling log-probabilities used in loss computation are not from the actual samplers, creating a skew in advantage-weighted loss contributions—a fundamentally different optimization landscape than intended.
• Bypass mode: The optimization target exists in a different sampling space from the rollout, making policy updates ineffective even though the denominator is correct.

Implications for Asynchronous LLM RL

Large-scale asynchronous RL (e.g., agentic tasks) is off-policy by design. TIM remains critical because it creates different probability landscapes for optimization and sampling spaces. Zero-mismatch RL preserves the full learning signal, unlike algorithmic corrections that mask or discard high-mismatch tokens/sequences.

VeXact as a Calibration Tool

VeXact enables:

• Scientific benchmarking of algorithmic patches in a noise-free environment.
• Accurate tuning of sensitive filtering thresholds ($\tau_{seq}$, $\tau_{tok}$) before deployment to large-scale pipelines.
• Causal diagnosis of whether a given stabilization technique addresses TIM, off-policy drift, or introduces new optimization bias.

---

Conclusion

This work demonstrates that TIM is not benign numerical noise but a systems-level perturbation that should be treated as a first-order factor in analyzing LLM RL stability. Key takeaways:

1. TIM alone can destabilize RL training across REINFORCE and GRPO setups.
2. Common implementation choices (recomputation, bypass) fail to eliminate TIM's impact because they change the effective optimization objective.
3. Algorithmic corrections (TIS + $r_{corr}$-based rejection) can closely approach the zero-mismatch reference but require careful design and calibration.
4. A joint system-algorithm perspective is essential for RL stability, highlighting the need for zero-mismatch RL execution.

Limitations and Future Directions

• Evaluation is limited in scale and coverage (finite models, tasks, system configurations).
• Generalization of mitigations across broader RL settings remains unclear.
• Potential optimization side effects of algorithmic corrections are not fully visible in current experiments.
• Future work should explore zero-mismatch execution as the default for both synchronous and asynchronous RL pipelines.