Train-infer mismatch for Open-weight MoE RL in Open-source code

0 train-infer mismatch for Open-weight MoE RL in Open-source code

Skip to the article

Contents

Where does mismatch come from, and where does it go?

In RL we sample rollouts with an inference engine and update our model by backpropagating gradients through log probabilities computed by prefilling those sampled rollouts through a training engine. These engines can compute different logprobs at many points in the forward pass, and there are multiple ways to correct for the induced off-policyness of those rollouts.

Trainer scoring, inference prefill and decode

Inference processes the prompt in parallel during prefill, writes the model's cached state and generates subsequent tokens with one-token decode calls. The trainer scores the completed sequence in parallel. The kernels in these paths can round or add the same values differently. For recurrent and compressed attention, prefill also creates state that every subsequent decode step reads and updates.

++ ++

2−8 2−8 2−8 2−8

1 + 2−7 (a + b) + c a + (b + c)

Rounding depends on the order of additions

The smallest representable number in BF16 that is larger than 1 is 1 + 2−7. If we compute (1 + 2−8), this rounds to 1. But if we first add (2−8 + 2−8), we get 2−7. Any time we do a reduction, we need to enforce the same order of additions across engines.

KERNEL A KERNEL B

OP 1 BF16 OP 2 OUTPUT A OP 1 FP32 OP 2 OUTPUT B same operations · different intermediate precision

Precision

When a kernel evaluates a sequence of operations, the precision of every intermediate value is part of the calculation. If one engine rounds an intermediate to BF16 while the other keeps it in FP32, later operations receive different values and can produce different output bits.

XORL TRAINING SGLANG PREFILL SGLANG DECODE

COMPLETE SEQUENCE · PARALLEL PROMPT STATE HANDOFF STATE +t STATE +t

Cached state must match across prefill and decode

Decode reads the cache and updates it on each decoded token. In order to achieve agreement for stateful layers like Gated DeltaNet, we need the prefill cache and every decode update to match what the trainer computes. This means we need to eliminate prefill-decode mismatch.

How mismatch changes the GRPO update

Let qθ denote the policy evaluated by the trainer and Ât the detached group-relative advantage in GRPOShao et al., “DeepSeekMath” (2024). Introduces GRPO and its importance ratio.. For a token at sampled after prefix ht, the on-policy gradient is:

∇J(θ) = Σₜ Eₕₜ∼dqθ, ₐₜ∼qθ(·|hₜ)[Âₜ ∇ log qθ(aₜ | hₜ)]

Each token gets scored multiple times. Inference records a probability for the token while sampling. Inference will record a slightly different probability if you ask it to prefill the finished sequence, and the trainer will compute a different one when it scores the rollout. Write qt for the trainer's selected-token probability, sPt for an inference prefill replay at the current weights, sDt for cached decode at those weights and μt for the probability recorded during sampling. We can define the importance sampling ratio in terms of these different forward paths.

ρₜ = qₜ / μₜ = [qₜ / sᴾₜ] trainer scoring vs. serving prefill × [sᴾₜ / sᴰₜ] serving prefill vs. serving decode × [sᴰₜ / μₜ] stale state

The objectives we optimize in RL weight the gradient by the importance sampling ratio to correct for these sources of mismatch.

The importance ratio in detail

At a fixed prefix ht, weighting samples by qt / μt rewrites an expectation over μ(· | ht) as one over qθ(· | ht). The prefix itself still comes from μ:

At fixed hₜ, Eμ[ρₜ f] = Eqθ[f]. Across rollouts, Eμ[ρₜ f] keeps hₜ ∼ dμ, not hₜ ∼ dqθ.

Stale weights and cached state

Let s̃Dt denote decode at the current weights with its cache rebuilt at those weights. The two terms are the weight change, and the stale cache:

sᴰₜ / μₜ = [s̃ᴰₜ / μₜ] stale weights × [sᴰₜ / s̃ᴰₜ] stale cache

We always flush the cache after updating the weights of the inference engine, so the second term is one.

Unclipped importance sampling

Unclipped importance sampling multiplies the token gradient by ρt. We detach the selected-token logprobs used to form the ratio and differentiate log qt.Tinker’s importance-sampling loss uses this formulation.

∇Lᴵˢₜ = −ρₜ Âₜ ∇ log qₜ; Eμ[‖gₜ‖²] = Eμ[ρₜ² Âₜ² ‖∇ log qₜ‖²]

The second moment of the gradient scales with the square of the importance ratio. We're training in finite-sized batches, so there may be some high-variance steps where the gradient is dominated by a few tokens with large ratios.

Clipped importance ratios

GRPO commonly clips this ratio using the PPOSchulman et al., “Proximal Policy Optimization Algorithms” (2017). Introduces the clipped surrogate objective. objective:

Jᴾᴾᴼₜ = min(ρₜ Âₜ, clip(ρₜ, 1 − εlow, 1 + εhigh) Âₜ); ∇Lᴾᴾᴼₜ = −1[not clipped] ρₜ Âₜ ∇ log qₜ

For a positive-advantage token, PPO drops the policy gradient once ρt exceeds 1 + εhigh. For a negative-advantage token, it drops the gradient once ρt falls below 1 −...

Train-infer mismatch for Open-weight MoE RL in Open-source code

Related Articles

Incident with Github.com

(no title)

Scientists reverse brain aging, with a nasal spray

NHS apologises and admits Palantir have access to identifiable patient data

AI has torched the market for junior programmers