How Much Training-Inference Mismatch Can On-Policy Distillation Tolerate?
Abstract
On-policy distillation tolerates substantial training-inference mismatch from quantized inference and asynchronous rollouts, with model- and layer-specific limits, and a lightweight admission check plus FP8 recipes can exploit this tolerance to improve training efficiency.
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 mismatch OPD tolerates, where its effects arise, and how training systems can use this tolerance. Measuring all mismatch sources on one scale, we examine quantization across three student-teacher pairs, with additional tests of weight staleness and injected noise on one pair. OPD tolerates substantial mismatch without meaningful loss of training quality but degrades at higher levels, revealing an empirical mismatch budget. Both the quality cost of mismatch and the mismatch introduced by quantization depend on the model, with some layers particularly sensitive to reduced precision. In controlled comparisons, OPD remains stable under rollout configurations that destabilize RL methods. Under weight noise, mismatch also shifts the prefixes reached during generation, introducing errors that per-token corrections cannot address. To put these findings into practice, we introduce a lightweight admission check that screens rollout configurations against a calibrated budget without requiring a full training run for each candidate. Model-specific FP8 recipes preserve training quality while providing capacity for at least 1.93times as many KV-cache tokens as BF16. Combining FP8 with asynchronous rollouts reaches the same training quality 1.62times sooner than synchronous BF16 training. Code is available at https://github.com/benyucong/train-inference-mismatch-opd
Community
Cheaper rollouts (quantized inference, asynchronous rollouts from older weights) make the rollout engine differ from the model being trained. We measure how much of this training–inference mismatch on-policy distillation (OPD) can absorb, and how training systems can use that tolerance.
- On a 7B student, OPD keeps its training quality at over 200× the mismatch of a standard BF16 engine and fails at about twice that level. Under the same engines, two RL methods break at about 12× and 60×.
- The budget depends on the model, and some layers are far more sensitive to reduced precision than others.
- A lightweight check screens a rollout configuration against the budget in minutes on one GPU, without a full training run.
- Within the budget, FP8 KV caches hold at least 1.93× as many tokens as BF16, and FP8 with asynchronous rollouts reaches the same training quality 1.62× sooner than synchronous BF16 training.
Code: https://github.com/benyucong/train-inference-mismatch-opd
Really useful measurement study, and the common nats-per-token scale is what makes it work: once staleness, FP8 and weight noise sit on one axis, you can see that eight versions of staleness only reaches 0.0028, while keeping just the first layer's KV cache in BF16 drops the vLLM FP8 cache estimate on the 7B student from 0.52 to 0.0037. The finding that this first layer carries 93 to 99% of the FP8 cache mismatch, tied to the key-projection bias in the DeepSeek students, is a very practical tip for anyone running quantized rollouts on limited hardware. I am curious whether an admission line calibrated on math transfers to other domains like code or tool-use rollouts, or whether each new task needs its own calibration runs per student.
Thanks! We only tested math, so this is our best guess. Checking a setup is cheap: you generate a few answers with the quantized model and compare its token probabilities with the full-precision model’s. No teacher and no training are needed. So for code or tool use, we would just rerun this check on prompts from that task. We expect the first-layer problem to show up again, because it comes from how these models are built, not from the task. The expensive part is finding the limit itself, since that takes real training runs. We would not reuse our limit from the paper: it already differs between our 1.5B and 7B models, and a new task may move it. So we would pick a setup that sits well below the limit, then confirm it with one training run on the new task.
Mismatch tolerance is the number I actually want on hand when I'm deciding how hard to quantize the rollout engine — fp8 KV cache, 4-bit weights on the inference side — since that's the only way OPD stays affordable at a batch size worth running. The trouble is a threshold measured on one student/teacher pair doesn't obviously transfer to mine; my guess is tolerance tracks how peaked the teacher's distribution is, which varies token to token, not run to run. Give me it as a knob I set per run — quant scheme, staleness in steps, async depth — plus a curve of where the student stops improving, not a single pass/fail number. I'd also want some evidence the student is getting good for reasons other than faithfully matching the teacher's noisiest tokens. And does that budget shrink as the student drifts away from the teacher, or stay flat? Fixed early and moving late would mean my scheduler needs to know about it.
Thanks! You're right that the limit doesn't carry over from one model pair to another. Even our 1.5B and 7B students have different limits. What does carry over is the method. Each setting (quantization, stale weights, added noise) is turned into one number: how far it pushes the rollouts away from the full-precision model. Figure 3 plots training quality against that number over 18 runs. Quality holds steady up to 0.184 and drops at 0.39. We also check benchmark accuracy, not just how close the student gets to the teacher. Below the limit, accuracy stays within 1.4 points of normal training. The limit also gets stricter later in training: one setting was fine at step 100 but had failed by step 200. So it's worth measuring this number throughout the run.
Get this paper in your agent:
hf papers read 2611.11450 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper