How Much Training-Inference Mismatch Can On-Policy Distillation Tolerate?

AI summary

Overview

Research area: Machine learning systems and post-training for large language models — specifically on-policy distillation (OPD) and the engineering of cheap rollout generation (quantized inference, asynchronous rollouts).

Technical level: Advanced. The paper assumes familiarity with policy-gradient updates, KL divergence, quantization formats (FP8, NVFP4), KV caching, and inference engines such as vLLM.

Scope: A measurement study of how much training–inference mismatch on-policy distillation can absorb before training quality degrades, how to estimate that limit cheaply before training, and what efficiency gains lie within it, across three student–teacher pairs (1.5B, 3B, 7B students).

What This Paper Is About

On-policy distillation trains a student model on responses it generates itself, with a fixed teacher scoring every token. Generating those responses is expensive, so practitioners use quantized inference and asynchronous (stale-weight) rollouts to speed it up — but both make the rollout engine differ from the model actually being trained. This paper asks how large that difference can get before OPD training quality suffers, whether that limit can be predicted in minutes rather than by a full training run, and why OPD tolerates mismatch that breaks reinforcement learning methods.

Key Contributions

1. A common scale and methodology for measuring mismatch and locating a "mismatch budget." Mismatch is defined as the per-token KL between the rollout engine's next-token distribution and the student's, in nats per token, averaged over the prefixes the engine actually visits. This one measure covers numerical kernel differences, quantization, and stale weights, which are otherwise expressed in incompatible units (bits versus model versions). An empirical bracket is defined from the passing and failing runs.

2. A lightweight admission check that takes minutes. A candidate engine generates a small set of responses, the standard BF16 engine re-scores the same tokens, and the resulting offline estimate is compared against a calibrated admission line with a safety margin. All eight decisions fixed before full training matched the observed outcomes (five of them from the offline check), and checking six configurations took 24 minutes on one GPU versus 5.5 to 6 hours on eight GPUs for one full run.

3. A quantified efficiency envelope within the budget. Model-specific FP8 recipes preserve training quality while providing capacity for at least 1.93× as many KV-cache tokens as BF16; FP8 combined with asynchronous rollouts reaches the synchronous BF16 run's step-200 held-out KL 1.62× sooner.

4. An explanation of why OPD tolerates mismatch that destabilizes RL, and where per-token corrections fall short. Because the OPD advantage is the teacher–student log ratio, the mismatch-induced bias in the gradient vanishes as the student approaches the teacher, so matching the teacher remains a stationary point whatever engine produces the rollouts. Under weight noise, mismatch also shifts which prefixes generation reaches, an error per-token corrections cannot address.

Main Findings

• OPD tolerates large mismatch. On the 7B student, all twelve configurations with mismatch D(c) ≤ 0.184 nats/token passed the primary quality rule, with relative held-out KL change between −5.0% and +6.4%. That level is over 200 times the standard engine's mismatch of 0.0008 nats/token. The two highest passing levels finish within 0.5% of the standard run's held-out KL.

• Quality degrades beyond a bracket. The empirical bracket at threshold τ = 0.20 is [0.184, 0.39] nats/token. Two weight-noise runs at D(c) ≈ 0.39 fail with Δ(c) ≈ +23%. The three every-layer quantized configurations reach D(c) = 0.56–0.68, fail with Δ(c) = +54% to +60%, and lose 3–5 accuracy points; under every-layer FP8, rollout answers are almost never correct from the start. At τ = 0.25 the bracket shifts to [0.39, 0.56]. The authors stress the bracket is not an exact failure threshold.

• Contrast with RL is stark. Under the same engines, REINFORCE and GRPO break at about 12 and 60 times the standard engine's mismatch respectively. In direct comparisons, REINFORCE broke under all three OPD-passing configurations (accuracy falls of 21, 59 and 63 points) and GRPO under two of the three (a 15-point fall with NVFP4).

• Staleness alone is cheap. Rollout weights up to eight versions old yield only D(c) = 0.0028 nats/token, far below the failure range — so neither bit width nor staleness alone predicts danger.

• Tolerance is model-specific. Two every-layer FP8 recipes on the 1.5B student give D(c) = 0.36–0.39 over steps 1–100 (versus 0.65–0.69 for comparable 7B quantization) yet finish with Δ(c) = +420% to +479% and about 11 points lower accuracy. Keeping the 1.5B student's first-layer KV cache in BF16 avoids the degradation (Δ(c) = −4.3%). Every-layer Model Optimizer FP8 on Llama-3.2-3B yields only D(c) ≈ 0.01 and passes (Δ(c) = +5.6%).

• Some layers are far more sensitive than others. On the 7B student an FP8 KV cache dominates mismatch at every tested weight and activation precision, and the first layer's cache accounts for 93%–99% of it across seven FP8-cache recipes. A keys-only emulation attributes about 82% of that layer's effect to first-layer key rounding. Retaining the first-layer KV cache in BF16 reduces vLLM's FP8-cache estimate from 0.52 to 0.0037 nats/token.

• The admission check works but is untested near its boundary. Offline estimates for the trained engines fall at ≤ 0.030 or ≥ 0.408 nats/token, with the admission boundary L/α ≈ 0.27 lying between the two groups. Pilot-derived settings proved optimistic: admitted 7B engines reach 1.56–1.68 times their estimates during steps 1–100, exceeding the margin α = 1.5, and a noise level at D(c) = 0.39 fails at step 200 despite lying below the pilot line L = 0.41. Retrospectively setting L = 0.184 would preserve all eight decisions but is not established as generally safe.

• Asynchrony helps until the trainer is the bottleneck. Throughput rises to 1.62× and 1.59× relative to synchronous BF16 at k = 4 and k = 8 stale versions, then levels off; all runs pass and finish within 3% of the standard held-out KL. From k = 2 on, the trainer's wait for rollouts is at most three seconds and training limits throughput to about 55 updates/hour. FP8 combined with k = 4 and 8 gives 1.63× and 1.55×, close to asynchrony alone.

• FP8 weights speed generation; FP8 KV cache mainly adds capacity. At 64 responses/step, Model Optimizer FP8 with BF16 first/last layers raises updates/hour by 1.27×, while vLLM's FP8 cache alone gives 1.00×. At a memory-bound batch of 1,536 responses per step, BF16 runs out of cache and must evict and recompute, while FP8 does not; FP8 variants then give 1.11× and 1.16× throughput, smaller because generation and teacher scoring take only 42% of each step.

• Mismatch matters more, not less, late in training. At three weight-noise levels, the mismatch-induced error at step 200 is 0.11 to 0.13 of its size at the untrained student, but the update shrinks faster: the error-to-update ratio at the middle level is 0.93 at the untrained student and 1.88 at step 200.

• Per-token corrections do not help OPD. Score centering rescues the tested RL failures (raising REINFORCE held-out accuracy by 61 points at step 100 and GRPO's by 16.5 points at step 200) but does not improve OPD. Under weight noise, the prefix-shift error these corrections cannot address grows 5.3 times from a passing to a failing noise level, versus 2.1 times for per-token error.

Methodology in Plain English

The authors run controlled distillation experiments and treat the rollout engine as the variable under study.

Setup. Three students — DeepSeek-R1-Distill-Qwen-7B, DeepSeek-R1-Distill-Qwen-1.5B and Llama-3.2-3B-Instruct — are each distilled for 200 steps from an RL-fine-tuned version of themselves acting as teacher. Training uses 8 responses per problem for 8 DAPO-Math-17k problems, temperature 1, up to 16,384 tokens per response, and a learning rate of 10⁻⁶. The testbed is two nodes of four NVIDIA GH200 GPUs each, with vLLM as the rollout engine and verl for training.

Varying the mismatch. On the 7B student the authors train 17 rollout configurations: six quantized ones from three production recipes (NVIDIA Model Optimizer FP8, NVFP4, and vLLM's FP8 KV cache, each with and without a BF16 escape hatch for the sensitive layer), five with different degrees of asynchrony (weights up to 1, 2, 4 or 8 versions old, plus colocated generation), four with injected Gaussian weight noise at σ ∈ {0.0665, 0.15, 0.234, 0.325} times each projection matrix's standard deviation, and two combining FP8 with asynchrony. The 1.5B and 3B pairs add three and two FP8 runs. Only the rollout engine is modified; the trainer, teacher and checkpoints stay in BF16.

Judging quality. The main metric is held-out KL: the student's per-token KL to the teacher on 530 held-out math problems from AIME24 and MATH-500, always generated by the standard BF16 engine rather than the configuration's own engine. A run fails if that KL is at least 20% worse than the standard run's (τ = 0.20) and a 95% paired bootstrap interval excludes zero. A secondary rule checks accuracy on AIME24, AIME25 and a non-overlapping subset of AMC23, with 16 samples per problem. For scale, standard training on the 7B student reduces mean held-out KL from 0.0965 to 0.0481, and repeated runs differ by at most 1.5% in final held-out KL.

Cheap screening. Instead of running full training for every candidate engine, the admission check has the candidate generate two responses for 16 math problems, re-scores those exact tokens with the standard BF16 engine, and compares the resulting estimate to a threshold derived from known passing runs, with a 1.5× safety margin.

Comparison with RL. To test whether OPD's tolerance is unusual, the same 7B student is trained with GRPO and REINFORCE using a 0/1 correctness reward and the same rollout configurations, and a run counts as broken if held-out DAPO accuracy on 1,000 problems drops at least five points or if entropy or response length collapses.

Analysis. A first-order analysis of the gradient separates the bias that comes from per-token probability error from the error that comes from the engine visiting different prefixes, which is what per-token importance corrections cannot fix.

Why This Matters

Impact on research. The paper reframes training–inference mismatch from a bug to be eliminated into a budget to be spent, and provides a shared unit (nats per token) that lets quantization and staleness be compared directly for the first time. It also gives a concrete, mechanism-level explanation for the anecdotal observation that OPD pipelines tolerate stale or quantized rollouts far better than RL pipelines, and it shows that corrections borrowed from RL may be solving the wrong part of the problem.

Real-world applications:

• Cheaper post-training pipelines. Teams running OPD as a post-training stage can adopt FP8 rollout engines and multi-version asynchronous rollouts within a calibrated budget rather than defaulting to BF16-synchronous training.
• Higher rollout throughput per GPU. FP8 KV caches hold at least 1.93× as many KV-cache tokens as BF16 on the 7B student, letting long chain-of-thought responses fit without eviction and recomputation.
• Faster model releases. FP8 with eight-version-stale weights reaches the synchronous BF16 run's quality 1.62× sooner, shortening the wall-clock critical path of a post-training stage.
• Safer configuration choices. A screening check that costs minutes on one GPU rather than hours on eight lets engineers reject a bad quantization recipe before committing a training budget.

Industry relevance. The recipes tested (NVIDIA Model Optimizer FP8, NVFP4, vLLM FP8 KV cache) are production formats, and the workloads span both compute-bound training batches and memory-bound large-batch generation. The finding that a single layer — the first layer's KV cache, tied to DeepSeek's key-projection bias and lack of query/key normalization — carries 93%–99% of the mismatch gives practitioners a specific, cheap mitigation rather than a blanket "don't quantize" rule.

Future Directions

• Establishing full budget brackets for smaller students. The 1.5B and 3B tests are partial; the paper explicitly states they do not establish complete brackets, and the 1.5B result (Δ(c) = +420% to +479% at mismatch levels the 7B tolerated) argues that 7B budgets cannot be transferred without calibration.
• A cheaper, more principled calibration procedure. The admission line L and margin α were derived from short pilots that proved optimistic (admitted engines reached 1.56–1.68 times their offline estimates). A method that calibrates these without representative full training runs remains open.
• Better handling of prefix-shift error. Per-token corrections do not improve OPD, and the prefix-shift error grows 5.3 times from a passing to a failing noise level. What intervention could address errors in which prefixes generation reaches is not established here.
• Testing near the admission boundary and on other model families. The offline estimates cluster at ≤ 0.030 or ≥ 0.408 with the boundary at ≈ 0.27, so the check is untested in between; the paper's own Qwen3-4B check is offline only and the NVFP4 runs are emulated in BF16 for quality without speed, leaving generalisation across architectures and true 4-bit hardware unmeasured.

Target Audience

Machine learning systems researchers and engineers who build or operate on-policy post-training pipelines — particularly those choosing rollout-engine precision, KV-cache formats, or asynchronous scheduling for LLM distillation. It is also valuable to RL infrastructure teams, since the RL-versus-OPD contrast bears directly on how much effort to invest in exact engine–trainer matching and per-token importance corrections. Readers without background in policy-gradient methods, quantization formats, or inference-engine internals will find the paper demanding.

Authors’ abstract

On-policy distillation (OPD) trains a student language model on its own responses, guided by a fixed teacher. Generating these responses is expensive, so systems use quantized inference and asynchronous rollouts to reduce the cost. Both, however, introduce training-inference mismatch between the rollout engine and the model being trained. Such mismatch can destabilize reinforcement learning (RL), but its effects on OPD remain unclear. This paper investigates how much…