How Much Training-Inference Mismatch Can On-Policy Distillation Tolerate?
Cong Yu, Ziting Wang, Songlin Jiang, Sitong Zhang, Chengliang Chai, Youhe Jiang, Eiko Yoneki, Tuo Shi
cs.LG, cs.AI
2026-11-08
A 7B student tolerates 200x the standard engine's rollout mismatch with no quality loss; FP8 plus async rollouts reach the same training quality 1.62x sooner.
On-policy distillation (OPD) trains a student on responses it samples itself, with a fixed teacher scoring every token; several recent model releases run it as a post-training stage. The cost sits in generating rollouts: with a standard BF16 engine, generation and teacher scoring take 76% of each training step on a 7B student. The two standard ways to cut that cost, quantized inference (FP8) and asynchronous rollouts from stale weights, both make the rollout engine diverge from the model being trained. In RL post-training this mismatch can collapse training, so production OPD pipelines inherited RL-style per-token corrections. Whether OPD actually needs them had never been measured.
The plan is to turn mismatch into one measurable quantity, then locate the level that stops hurting training.
| Configuration | Mismatch (nats/token) | Held-out KL change | Verdict |
| 7B, FP8, first layer's KV cache in BF16 | ≤0.184 | −5.0% to +6.4% | pass |
| 7B, weight noise | 0.39 | +23% | fail |
| 7B, FP8 every layer | 0.56–0.68 | +54% to +60% | fail |
| 1.5B, FP8 every layer | 0.36–0.39 | +420% to +479% | fail |
| Llama-3.2-3B, FP8 every layer | ≈0.01 | +5.6% | pass |
The 7B budget brackets at [0.184, 0.39] nats/token, and the lower bound is more than 200x the standard engine's mismatch. Under the same engines, REINFORCE breaks at about 12x and GRPO at about 60x. Under every-layer FP8, rollout answers are almost never correct from the first step.
The mismatch source can be localized: on the 7B, an FP8 KV cache dominates, and the first layer's cache alone carries 93–99% of it. Keeping that layer in BF16 drops the estimate from 0.52 to 0.0037. DeepSeek students carry a key-projection bias and lack QK-norm, so first-layer keys run unusually large, and about 82% of that layer's effect is key rounding. Llama and Qwen3 do not have this problem.
All eight pre-registered admission decisions matched full-run outcomes. The pilot safety margin of 1.5x was nonetheless too small: in-training mismatch reached 1.56–1.68x the offline estimate, and the initial admission line of 0.41 was in hindsight too high.
The gains are concrete: async rollouts alone give 1.62x throughput at k=4; FP8 plus async at k=8 reaches the synchronous BF16 run's step-200 held-out KL 1.62x sooner, with final accuracy 60.5% versus 61.5%; FP8 KV cache holds at least 1.93x as many tokens.
The RL contrast is stark: REINFORCE broke under all three configurations OPD tolerates (accuracy down as much as 63 points), GRPO under two of three. Score centering prevents the tested RL failures (61 points back for REINFORCE) yet slightly hurts OPD, leaving held-out KL 2.3–3.7% above uncorrected OPD. Mechanistically, OPD's mismatch bias vanishes as the student approaches the teacher while RL's does not. Mismatch also shifts which prefixes generation reaches, an error per-token corrections cannot touch, and it grew 5.3x from a passing to a failing noise level. Offline OPD variants that sample rollouts once sit at an equivalent mismatch of about 0.082, inside the budget, which helps explain why they work.
Teams running OPD can act on this directly: drop the RL-style per-token corrections, use quantized and asynchronous rollouts freely inside a measured budget, and keep the first layer's KV cache in BF16 as a one-line safeguard. A 1.62x end-to-end speedup matters when rollouts are the bottleneck. The counterintuitive part: safety cannot be read off bit width or staleness, it has to be measured per model.
The budget bracket holds only for a given student, training setup, and quality threshold; the 1.5B student loses over 420% held-out KL at mismatch levels the 7B shrugs off, so the 7B numbers do not transfer. The admission check is a screen, not a guarantee: each new student still needs calibration runs, and the authors themselves report the pilot margin and admission line were too loose, with the region near the boundary untested. Coverage is math distillation only (DAPO-Math-17k, 200 steps, up to 7B); other domains, larger models, and longer runs are unverified. The RL baselines showed no measurable accuracy gains within 100 steps even on a standard engine, so the RL comparison is about breaking versus not breaking, not learning quality. NVFP4 runs emulated quantization for quality only, so the throughput numbers exclude it.