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
实测在线蒸馏的失配预算:7B 学生容忍 200 倍标准引擎失配不伤质量,FP8 加异步 rollout 提前 1.62 倍达到同质量。
On-policy distillation(OPD,在线蒸馏)让学生模型在自己采样的回答上训练,由固定 teacher 给每个 token 打分,最近的模型发布常把它当 post-training 阶段。成本大头在生成 rollout:标准 BF16 引擎下,7B 学生每个训练步有 76% 的时间花在生成回答和 teacher 打分上。两个标准省钱手段都会引入训练推理失配:量化(FP8)改变引擎的输出分布,异步 rollout 让引擎用落后几个版本的旧权重。这种失配在 RL 后训练里能把训练搞崩,不少 OPD 流水线因此直接继承了 RL 的逐 token 修正。OPD 到底需不需要这层保险,此前没有系统测量。
思路是先把「失配多大」变成一个可测的量,再找出不伤训练质量的上限。
| 配置 | 失配(nats/token) | held-out KL 变化 | 判定 |
| 7B,FP8,首层 KV cache 留 BF16 | ≤0.184 | −5.0%+6.4% | 通过 |
| 7B,权重噪声 | 0.39 | +23% | 失败 |
| 7B,全层 FP8 | 0.560.68 | +54%+60% | 失败 |
| 1.5B,全层 FP8 | 0.360.39 | +420%+479% | 失败 |
| Llama-3.2-3B,全层 FP8 | ≈0.01 | +5.6% | 通过 |
7B 的预算区间是 [0.184, 0.39] nats/token,下限是标准引擎失配的 200 多倍。同一批引擎下 REINFORCE 约 12 倍就崩,GRPO 约 60 倍。全层 FP8 从第一步起就几乎产不出正确回答。
失配来源可以定位:7B 上 FP8 KV cache 主导失配,首层 cache 一层占 93%99%;把首层留成 BF16,估计从 0.52 降到 0.0037。DeepSeek 学生带 key-projection bias 又没有 QK-norm,首层 key 异常大,该层约 82% 的效应来自 key 舍入。Llama 和 Qwen3 没这个问题。
准入检查 8 个预登记决定全部命中;但试点安全系数 1.5 偏小,训练中实际失配达离线估计的 1.561.68 倍,准入门限 0.41 事后看也偏高。
收益实打实:异步 k=4 吞吐 1.62×;FP8 加异步 k=8 提前 1.62× 达到同步 BF16 的 step-200 held-out KL,最终 accuracy 60.5% 对 61.5%;FP8 KV cache 容量至少 1.93×。
RL 对照最扎眼:三个 OPD 稳稳通过的配置,REINFORCE 全崩(accuracy 最多掉 63 点),GRPO 崩两个。score centering 防得住测试的 RL 崩溃(REINFORCE 拉回 61 点),对 OPD 反而略有害,held-out KL 比不修正高 2.3%3.7%。机制上,OPD 的失配偏差随学生逼近 teacher 自然消失,RL 的不会;失配还会把生成带到不同的前缀上,这部分误差逐 token 修正够不着,从通过档到失败档涨了 5.3 倍。只采样一次的离线 OPD 等效失配约 0.082,落在预算内,给 Semi-OPD 这类方法的成功提供了一个解释。
跑 OPD 的团队可以直接改配置:RL 式逐 token 修正可以摘掉,量化加异步在预算内放心用,首层 KV cache 保 BF16 是一行配置的保险。rollout 是瓶颈时,1.62 倍端到端加速是真金白银。安全边界必须按模型逐个测,bit 宽度和延迟版本数都不构成判断依据,这一点和常见直觉相反。
预算区间只对特定学生、训练设置和质量阈值成立,1.5B 上同档失配 held-out KL 恶化超过 420%,7B 的数字不能外推。准入检查是筛查不是保证,每个新学生仍需校准训练,论文自己承认试点参数偏松、准入线附近没有测试。实验只覆盖数学蒸馏(DAPO-Math-17k、200 步、最大 7B),其他任务域、更大模型、更长训练未验证。RL 基线自身在 100 步内没有可测的 accuracy 提升,所以 RL 对照比较的是崩不崩,不是学多好。NVFP4 是模拟量化,只测质量,速度数字不含它。