On-Policy Distillation（OPD）学习文档

本文档基于原文《8.2 On-Policy Distillation》以及后续问答整理，只归纳对话中已经出现的内容，不额外编造未讨论的细节。

OPD 一句话总览
核心术语
从离线蒸馏到 OPD
Reverse KL 与逐 token 训练信号
Teacher 必须评价同一条 Student 轨迹
tokenizer 兼容与 chat template
logprob 位置对齐与"一位偏移"
Student 网络一般怎么实现
同步版 OPD 代码流程
采样条数：batch_size 与 group_size
多条 completion 的梯度如何回传
OPD 与强化学习、离线蒸馏的区别
运行、异步版与指标观察
常见误区与检查清单
一句话总结

1. OPD 一句话总览

OPD 可以理解成：让学生模型自己写作业，老师不替它重写，而是逐字批改学生写出来的这条作业；学生根据老师的逐字批改调整自己。

它借用强化学习的"当前策略采样轨迹 → 构造 advantage → 更新策略"的流程，但训练信号不是环境给的 0/1 奖励，而是 Teacher 对每个 token 的概率评价。

OPD 的三步迭代：

Student rollout：Student 使用当前权重生成回答；
Teacher feedback：Teacher 在完全相同的 prompt 与 completion 上计算 logprob；
Student update：根据两者概率差构造逐 token 信号并更新 Student。

Teacher 只做前向计算，不参与反向传播；Student 的采样动作也不参与反向传播；梯度从构造好的训练目标回到当前 Student 参数。

2. 核心术语

术语	通俗解释
Teacher / Student	Teacher 是强模型，Student 是要训练的小模型
on-policy	用"当前 Student 自己生成的轨迹"来训练，而不是固定数据集
rollout	让 Student 实际生成回答的过程
completion	Student 根据 prompt 生成出来的那一段回答，即"补全内容"
logprob	模型给某个 token 的对数概率；越接近 0 概率越高，越负概率越低
KL 散度	衡量两个概率分布差多少
reverse KL	本章使用的分布差异，具有 mode-seeking 特性
advantage	强化学习中表示"这个动作该鼓励还是抑制"的信号；OPD 中是逐 token 的
tokenizer 对齐	Teacher 和 Student 必须把同一段文本切成同一组 token id
chat template	把多轮对话消息拼成模型能读的一长串文本的格式化规则
停止梯度边界	Teacher 只前向打分，不反传；Student 采样动作也不反传；梯度只回到 Student 参数

completion 的例子：

text 复制代码

prompt:      "请解这道题：1+1=？"
completion:  "1+1=2，所以答案是 2。"

在代码里就是：

python 复制代码

prompt_ids + completion_ids

注意：OPD 里的 completion 不是 Teacher 写的，而是 Student 当前策略采样出来的回答。

3. 从离线蒸馏到 OPD

3.1 离线蒸馏

设 Teacher 为 πT\pi_TπT，Student 为 πθ\pi_\thetaπθ。离线蒸馏先固定数据集 DDD，然后最小化：

Loffline=Es∼D $D(πT(\cdot∣s),πθ(\cdot∣s))$ L_{\text{offline}}=E_{s\sim D} $D(\\pi_T(\\cdot\|s),\\pi_\\theta(\\cdot\|s))$ Loffline=Es∼D $D(πT(\cdot∣s),πθ(\cdot∣s))$

通俗说：题库是固定的，学生反复练这套题。

问题是：Student 在真实生成时可能很早写错一个 token，然后进入一个固定数据里从没出现过的状态。固定数据中大多是 Teacher 的高质量轨迹，但 Student 真实生成时早早写错，后面那段"分布外轨迹"没有被蒸馏数据覆盖，Teacher 的指导就覆盖不到。

3.2 OPD 的做法

OPD 把状态采样改为当前 Student：

y∼πθold(⋅∣x) y\sim \pi_{\theta_{\text{old}}}(\cdot|x) y∼πθold(⋅∣x)

然后让 Teacher 计算同一条 Student completion 的条件概率：

log⁡πT(yt∣x,y<t) \log \pi_T(y_t|x,y_{<t}) logπT(yt∣x,y<t)

所以 OPD 是"学生当前走到哪，老师就在哪批改"。老师不需要重新生成一套答案，也不需要参与反向传播，只提供稠密的逐 token 概率反馈。

3.3 OPD 的用途

原文提到：OPD 可用于推理能力迁移、能力恢复、持续学习，也能与强化学习数据流结合：Student rollout 对应 on-policy 轨迹，Teacher logprob 对应逐 token 训练信号，Student 更新后再生成下一批轨迹。

4. Reverse KL 与逐 token 训练信号

本章使用 reverse KL：

DKL(πθ∥πT)=Ey∼πθ $logπθ(y∣x)-logπT(y∣x)$ D_{KL}(\pi_\theta\|\pi_T) =E_{y\sim\pi_\theta} $\\log\\pi_\\theta(y\|x)-\\log\\pi_T(y\|x)$ DKL(πθ∥πT)=Ey∼πθ $logπθ(y∣x)-logπT(y∣x)$

对 Student 实际采样到的第 ttt 个 token，做 Monte Carlo 估计：

dt=log⁡πθold(yt∣x,y<t)−log⁡πT(yt∣x,y<t) d_t=\log\pi_{\theta_{\text{old}}}(y_t|x,y_{<t}) -\log\pi_T(y_t|x,y_{<t}) dt=logπθold(yt∣x,y<t)−logπT(yt∣x,y<t)

4.1 怎么理解 dtd_tdt

dt>0d_t>0dt>0：Student 对这个 token 的偏好高于 Teacher，即 Student 给的概率比老师还高，应抑制；
dt<0d_t<0dt<0：Teacher 给的概率更高，老师更看好这个 token，应鼓励。

为了最小化 reverse KL，训练 advantage 写为：

At=−βdt A_t=-\beta d_t At=−βdt

其中 β\betaβ 控制蒸馏强度。

于是：

如果 dt>0d_t>0dt>0，At<0A_t<0At<0，这个 token 被压低；
如果 dt<0d_t<0dt<0，At>0A_t>0At>0，这个 token 被提高。

这样每个 completion token 都有独立训练信号，信号密度远高于"整条回答最后只给一个 0/1 reward"。

4.2 reverse KL 的特点

reverse KL 有较强 mode-seeking 特性。Student 会优先集中到 Teacher 已经赋予较高概率的区域，适合将 Teacher 的高置信行为压缩到小模型里。但 Teacher 与 Student 的能力差距、推理风格、tokenizer 兼容性都会影响蒸馏效果。

更大的 Teacher 只是候选条件，Teacher 还需要在 Student 当前轨迹上提供有价值且可学习的新信号。

4.3 Monte Carlo 估计是什么

Monte Carlo 估计就是：当一个期望很难精确算时，用"采样 + 平均"来近似它。

理论上：

Ey∼πθ $f(y)$ ≈1N∑i=1Nf(yi) E_{y\sim \pi_\theta} $f(y)$ \approx \frac{1}{N}\sum_{i=1}^N f(y_i) Ey∼πθ $f(y)$ ≈N1i=1∑Nf(yi)

NNN 越大越准，但 N=1N=1N=1 也是一个无偏估计，只是方差大。训练不是只做一次估计，而是每个 step 采样一批、更新一次，下一个 step 用新策略再采样，很多 step 累积起来样本量并不小。

更重要的是：一条 completion 不是只提供一个信号，而是每个 token 都提供一个信号。 一条回答有几百个 token，就有几百个 dtd_tdt。

所以"单步单条"不等于"整个训练只用一条"。

5. Teacher 必须评价同一条 Student 轨迹

OPD 的关键数据约束是：Teacher 与 Student 必须在相同前缀下评价相同目标 token。

如果 Teacher 自己生成一条回答，再让 Student 学这条回答，那就退回普通合成数据蒸馏了。OPD 要的是：Student 写什么，Teacher 就批改什么。

Teacher 的输入是：

text 复制代码

prompt + Student 生成的 completion

顺序是：

Student 先根据 prompt 生成 completion；
把 prompt + completion 拼起来；
Teacher 看这条完整序列，逐 token 给出"如果是我，我会给这个 token 多大概率"；
Student 根据 Student logprob 和 Teacher logprob 的差来更新。

Teacher 不自己生成回答，Teacher 只批改 Student 写出来的回答。 Teacher 的参数固定，不参与反向传播。

5.1 代码：Teacher 对同一条轨迹打分

python 复制代码

def completion_teacher_logprobs(teacher_client, prompt_ids, completion_ids):
    all_ids = prompt_ids + completion_ids
    all_logprobs = teacher_client.compute_logprobs(
        trio.ModelInput.from_ints(all_ids)
    ).result()

    completion_logprobs = all_logprobs[len(prompt_ids):]

    if (len(completion_logprobs) != len(completion_ids)
        or any(value is None for value in completion_logprobs)):
        raise ValueError("Invalid teacher logprobs for completion tokens")

    return [float(value) for value in completion_logprobs]

长度检查与 None 检查用于保护 token 对齐。通过检查后，第 ttt 个 Student completion token、Student 旧 logprob 和 Teacher logprob 才能组成同一个训练位置。

6. tokenizer 兼容与 chat template

6.1 tokenizer 兼容是什么意思

同一段文本，用 Student tokenizer 切出来的 token id 序列，和用 Teacher tokenizer 切出来的 token id 序列，必须完全一样。

不只是长度一样，而是：

每个位置的 token id 相同；
每个 id 在两个模型里对应的字符串也相同；
special tokens、BOS/EOS、chat template 的处理也一致。

验证方法：

python 复制代码

text = "A short tokenizer alignment test."
student_ids = student_tokenizer.encode(text, add_special_tokens=False)
teacher_ids = teacher_tokenizer.encode(text, add_special_tokens=False)

assert student_ids == teacher_ids

也可以反向 decode 检查：

python 复制代码

assert student_tokenizer.decode(student_ids) == teacher_tokenizer.decode(teacher_ids)

6.2 为什么必须一样

OPD 流程中，Teacher 拿到的是 Student 产生的 token id，不是文本。Teacher 不会重新 tokenize，而是直接用自己的词表解释这些 id。

如果两个 tokenizer 不兼容，可能出现：

text 复制代码

Student 认为 id=1234 是 "的"
Teacher 认为 id=1234 是 "hello"

那么 Teacher 看到的序列语义已经完全变了。它返回的 logprob 虽然数字上还能算，但评价的已经不是 Student 实际写的那句话，概率差 dtd_tdt 也就没有正确语义。

6.3 chat template 是什么

chat template 就是"把多轮对话消息拼成模型能读的一长串文本"的格式化规则。

例如：

python 复制代码

messages = [
    {"role": "system", "content": "你是一个数学助手。"},
    {"role": "user", "content": "1+1=?"},
]

套用常见 Qwen 风格模板后，可能变成：

text 复制代码

<|im_start|>system
你是一个数学助手。<|im_end|>
<|im_start|>user
1+1=?<|im_end|>
<|im_start|>assistant

然后 tokenizer 再把这段文本切成 token id。

chat template 负责插入：

角色标记：system / user / assistant；
每轮开始和结束的分隔符；
是否加 generation prompt；
有些模型还会加 thinking 开关、工具调用标记等。

6.4 代码中的 chat template 使用

python 复制代码

messages = [{"role": "user", "content": content}]

prompt = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=enable_thinking,
)
return tokenizer.encode(prompt, add_special_tokens=False)

参数解释：

messages：结构化对话，这里是单轮 user 消息；
apply_chat_template(...)：按 tokenizer 自带模板渲染成文本；
tokenize=False：先返回字符串，不直接切 token；
add_generation_prompt=True：在末尾加"轮到 assistant 了"的标记；
enable_thinking=...：控制模板里是否加思考标记；
最后 tokenizer.encode(...) 把渲染后的文本切成 token id。

6.5 在 OPD 里为什么强调"prompt 只渲染一次"

正确流程：

text 复制代码

Student tokenizer 渲染一次 chat template
        ↓
得到 prompt_ids
        ↓
Student 生成 completion_ids
        ↓
prompt_ids + completion_ids 直接给 Teacher
        ↓
Teacher 不重新渲染 template，只做前向算 logprob

错误做法是 Student 用自己的 template 渲染一次，Teacher 用自己的 template 再渲染一次。因为两套 template 可能插入不同角色标记、系统提示写法不同、分隔符不同、generation prompt 不同、thinking 开关不同，结果就是两边看到的上下文已经变了。

文章原话：

对于同系列但 chat template 不同的模型，应由 Student tokenizer 渲染一次完整 prompt，然后让 Teacher 直接评价对应 token 序列。不要分别套用两份 chat template，否则两边看到的上下文已经发生变化。

通俗说：题目只由学生抄一遍，老师直接看学生抄好的题和学生写的答案；不要让老师再按自己的格式把题目重抄一遍。

6.6 add_generation_prompt=True 的作用

它会在渲染后的 prompt 末尾加上"轮到 assistant 说话"的起始标记。

加了：

text 复制代码

<|im_start|>user
1+1=?<|im_end|>
<|im_start|>assistant

不加：

text 复制代码

<|im_start|>user
1+1=?<|im_end|>

不加的话，模型看到的上下文停在"用户刚说完话"，但没有明确的"现在该你回答了"的信号。后果可能是：

模型可能不知道该由谁继续写；
可能继续生成另一个 <|im_start|>user 开头的内容，模拟用户继续说话；
生成格式混乱或质量下降；
不同模型表现不一样，不能假设所有模型都能自动补上这个信号。

在 OPD 的 build_prompt() 里，这个 prompt 后面要直接交给 Student 采样生成，所以必须加 add_generation_prompt=True。

7. logprob 位置对齐与"一位偏移"

7.1 位移从哪里来

模型是"看着前面的 token，预测下一个 token"。假设输入：

text 复制代码

位置:   0    1    2    3
token:  A    B    C    D

模型内部做的是：

text 复制代码

看到 A        → 预测 B
看到 A,B      → 预测 C
看到 A,B,C    → 预测 D

它不会"看到 A 预测 A"。所以对于输入 [A, B, C, D]，真正有意义的"下一个 token 概率"是：

text 复制代码

P(B | A)
P(C | A,B)
P(D | A,B,C)

一共 3 个，不是 4 个。这就是位移的来源。

7.2 两种常见返回约定

约定一：返回"预测下一个 token"的 logprob，长度比输入少 1。

输入：

text 复制代码

[A, B, C, D]

text 复制代码

[ log P(B|A),  log P(C|A,B),  log P(D|A,B,C) ]

返回数组第 0 个位置对应输入的第 1 个 token B。这就是"偏移一位"。

约定二：返回和输入等长，位置 i 对应"预测第 i 个 token"。

可能返回：

text 复制代码

[ log P(A|BOS),  log P(B|A),  log P(C|A,B),  log P(D|A,B,C) ]

这种约定下，返回数组第 i 个位置对应输入第 i 个 token，没有偏移。还有的 SDK 返回等长数组，但第 0 个位置是占位/None，语义又不一样。

关键：长度看起来差不多，语义可能完全不同，不能靠猜。

7.3 错一位会发生什么

假设真实约定是"约定一"，返回数组第 0 个位置对应 B，但代码以为第 0 个位置对应 A，于是按输入位置去切。

text 复制代码

prompt_ids = [P1, P2]
completion_ids = [C1, C2, C3]
all_ids = [P1, P2, C1, C2, C3]

真实返回：

text 复制代码

all_logprobs = [
    log P(P2 | P1),        # 对应 P2
    log P(C1 | P1,P2),     # 对应 C1
    log P(C2 | P1,P2,C1),  # 对应 C2
    log P(C3 | P1,P2,C1,C2)# 对应 C3
]

如果代码写：

python 复制代码

completion_logprobs = all_logprobs[len(prompt_ids):]

len(prompt_ids)=2，切出来是：

text 复制代码

[ log P(C2 | ...), log P(C3 | ...) ]

长度是 2，但 completion_ids 长度是 3，长度检查会报错，这还算好的。更隐蔽的是长度碰巧相等但内容错位：把 C2 的老师 logprob 当成 C1 的去和学生 C1 概率比较，dtd_tdt 语义完全错。

通俗说：老师批改三道题，返回评语表第一行写的是"关于 C2 的评价"，你却以为第一行是 C1 的，每道题都错位一格。

7.4 文章代码为什么写 `all_logprobs[len(prompt_ids):]`

这说明作者已经确认 PyTRIO 0.2.6 的返回约定不是上面"少一个"的约定 ，而是另一种位置定义，使得 len(prompt_ids): 正好切到 completion 第一个 token。文章也明确提醒：

不同 SDK 对 logprob 的位置定义可能存在一位偏移，因此不能仅凭数组长度猜测切片。最稳妥的方法是构造一个很短的已知序列，确认返回值中每个位置究竟表示"当前 token 的 logprob"还是"下一个 token 的 logprob"，再完成 prompt 与 completion 的边界测试。本章代码已经按照 PyTRIO 0.2.6 的返回约定完成对齐。

所以不是文章代码错了，而是换 SDK 就必须重新确认。

7.5 如何保证位置对齐

文章推荐：构造一个很短的已知序列做探针测试。

步骤：

构造短序列，例如 ids = [A, B, C]；
调用 compute_logprobs(ids)，打印返回数组和每个位置的 token；
看返回长度是输入长度、少 1、还是多 1；
用"概率差异很大"的序列辅助判断哪个索引对应哪个 token；
做 prompt 与 completion 的边界测试，确认 completion_lps[0] 到底是不是 completion 第一个 token 的 logprob。

运行时保护：

python 复制代码

if (len(completion_logprobs) != len(completion_ids)
    or any(value is None for value in completion_logprobs)):
    raise ValueError("Invalid teacher logprobs for completion tokens")

长度检查和 None 检查能挡住"长度不对"和"有 None"的情况，但挡不住"长度对但语义错一位"。语义对齐必须靠短序列探针测试和 SDK 文档确认。

8. Student 网络一般怎么实现

不一定直接在 Teacher 网络上简化。 实际做法通常更灵活。

文章示例：

Student：Qwen/Qwen3.5-4B
Teacher：Qwen/Qwen3.6-27B
Student 用 LoRA training client 训练；
Teacher 用 sampling client 只做前向打分。

也就是说，Student 是一个较小的基座模型，通过 LoRA 微调来学习 Teacher 的逐 token 概率信号。它不是把 27B 教师直接剪枝成 4B，而是选一个现成的小模型作为 Student。

常见 Student 实现方式：

直接选同系列小模型

例如 Teacher 是 Qwen 27B，Student 选 Qwen 4B。优点：tokenizer、chat template 通常更兼容。本文示例就是这种方式。
选不同架构的小模型

只要 tokenizer 兼容、logprob 能对齐，也可以。但风险更高：tokenizer 不一致、chat template 不一致、概率语义对不上。
对 Teacher 做剪枝 / 层丢弃 / 宽度缩小

这是"直接简化教师"的一类做法，但本文没有展开。剪枝后还要继续训练，否则性能可能下降很多。
从零训练或继续预训练一个小模型

也可以，但成本更高，通常不是 OPD 示例的重点。
LoRA / 全参数微调

Student 可以是小基座 + LoRA，也可以全参数训练。本文用 LoRA training client，Teacher 固定。

关键约束不是"Student 必须由 Teacher 简化而来"，而是：

Student 和 Teacher 的 tokenizer 要兼容；
同一条 prompt 只渲染一次，Teacher 直接评价同一组 token；
logprob 的位置定义要对齐；
Teacher 要在 Student 当前轨迹上给出有价值的学习信号。

通俗说：Student 不一定是把 Teacher"砍小"得到的；更常见的是直接找一个小模型当学生，让它模仿大模型在"学生自己写出来的答案"上的逐字概率判断。

9. 同步版 OPD 代码流程

配套目录有同步和异步两个脚本：

文件	执行方式	用途
`01-demo-sync.py`	同步	正文逐段讲解的教学主线
`02-demo-async.py`	异步	并发处理 Student rollout 与 Teacher logprob 请求

同步版流程：

准备参数与 DeepMath 数据；
把问题渲染成 Student prompt；
创建 Student 与 Teacher；
让 Teacher 评价 Student 的同一条轨迹；
构造逐 token OPD Datum；
生成 Student 轨迹并计算 reverse KL；
更新 Student、记录指标并保存权重。

9.1 数据

DeepMath-103K 在 ModelScope 上有 10 个 parquet 分片。--num-shards 1 适合第一次联调，--num-shards 10 覆盖完整数据集，--sample-size 控制加载后实际参与 prompt pool 的样本数。

OPD 使用 prompt-only 数据，不读取 DeepMath 的参考推理过程。build_prompt() 把问题和输出要求交给 Student tokenizer，prompt 只渲染一次。

9.2 创建 Student 与 Teacher

python 复制代码

service_client = trio.ServiceClient()

training_client = service_client.create_lora_training_client(
    base_model=args.base_model,
    rank=args.lora_rank,
    seed=args.seed,
)
tokenizer = training_client.get_tokenizer()

teacher_client = service_client.create_sampling_client(
    base_model=args.teacher_base_model or args.base_model,
    model_path=args.teacher_model_path,
)

Teacher 不调用 forward_backward() 或 optim_step()，参数在 OPD 过程中保持固定。老师只批改，不学习。

9.3 构造逐 token OPD Datum

OPD 与 GRPO 使用相同的 importance_sampling schema，差异在 advantage。GRPO 给整条 completion 复制一个组内 advantage；OPD 给每个 completion token 写入独立的 reverse-KL advantage。

python 复制代码

def build_opd_datum(prompt_ids, completion_ids, old_logprobs, advantages):
    prompt_loss_len = len(prompt_ids) - 1

    input_ids = prompt_ids + completion_ids[:-1]
    target_ids = [0] * prompt_loss_len + completion_ids
    padded_logprobs = [0.0] * prompt_loss_len + list(old_logprobs)
    padded_advantages = [0.0] * prompt_loss_len + list(advantages)

    if not (len(input_ids) == len(target_ids)
            == len(padded_logprobs) == len(padded_advantages)):
        raise ValueError("OPD datum fields must have the same length")

    return trio.Datum(
        model_input=trio.ModelInput.from_ints(input_ids),
        loss_fn_inputs={
            "target_tokens": np.asarray(target_ids, dtype=np.int64),
            "logprobs": np.asarray(padded_logprobs, dtype=np.float32),
            "advantages": np.asarray(padded_advantages, dtype=np.float32),
        },
    )

右移规则：自回归模型是"看前面的 token 预测下一个 token"。

input_ids = prompt_ids + completion_ids[:-1]：输入不包含 completion 最后一个 token；
target_ids = [0]*prompt_loss_len + completion_ids：prompt 区间用 0 占位，不训练；completion 区间放目标 token；
padded_logprobs 和 padded_advantages 在 prompt 区间也用 0 占位；
completion 区间一一放入 Student token、Student old logprob 和该 token 的 OPD advantage。

prompt_loss_len = len(prompt_ids)-1 是为了让右移后的长度对齐。

9.4 生成轨迹并计算 reverse KL

每个 step 开始时，代码按 sampler_refresh_steps 刷新 Student sampler。默认值为 1，因此每一步都使用最新 Student 权重：

python 复制代码

if student_sampler is None or step % args.sampler_refresh_steps == 0:
    student_sampler = training_client.save_weights_and_get_sampling_client()

对每个 prompt，Student 先生成 group_size 条轨迹，Teacher 再对每条非空 completion 打分：

python 复制代码

result = student_sampler.sample(
    prompt=trio.ModelInput.from_ints(prompt_ids),
    num_samples=args.group_size,
    sampling_params=sampling_params,
    return_text=False,
).result()

for sequence in result.sequences:
    completion_ids = sequence.tokens
    if not completion_ids:
        continue

    student_lps = [float(value) for value in sequence.logprobs]
    teacher_lps = completion_teacher_logprobs(
        teacher_client, prompt_ids, completion_ids
    )

    reverse_kl = np.asarray(student_lps) - np.asarray(teacher_lps)
    advantages = -args.kl_penalty_coef * reverse_kl

    datums.append(build_opd_datum(
        prompt_ids, completion_ids, student_lps, advantages
    ))
    reverse_kls.extend(reverse_kl.tolist())
    completion_token_counts.append(len(completion_ids))

每条轨迹独立得到 Teacher 信号，不进行 GRPO 式的组内中心化。

9.5 更新 Student

python 复制代码

if not datums:
    raise RuntimeError("No OPD datums were built")

fwd_bwd = training_client.forward_backward(
    datums,
    loss_fn="importance_sampling",
)
optim = training_client.optim_step(adam)

fwd_bwd_result = fwd_bwd.result()
optim.result()

10. 采样条数：batch_size 与 group_size

参数	含义
`batch_size`	这一 step 用多少个不同 prompt
`group_size`	对同一个 prompt 采样多少条 completion

代码：

python 复制代码

result = student_sampler.sample(
    prompt=trio.ModelInput.from_ints(prompt_ids),
    num_samples=args.group_size,
    ...
)

group_size=1：同一个问题只生成 1 条答案；
group_size=4：同一个问题生成 4 条答案。

同步版示例：

bash 复制代码

--batch-size 1 --group-size 1

是"1 个问题，1 条答案"，教学联调省算力。

异步版示例：

bash 复制代码

--batch-size 4 --group-size 2

是"4 个不同问题，每个问题 2 条答案"，一共最多 8 条 completion。

10.1 是"同时输出"吗

更准确地说：是对同一个 prompt 一次请求采样多条完整 completion。

不是模型在同一个 token 位置上同时输出多个 token，而是：

text 复制代码

同一个 prompt
   ├── 采样答案 1：C1, C2, C3, ...
   ├── 采样答案 2：D1, D2, D3, ...
   └── 采样答案 3：E1, E2, E3, ...

每条答案都是一条完整的自回归生成轨迹。同步版逐 prompt 调用 sample(num_samples=group_size)；异步版在 batch 内并发调用 sample_async()，所以"同时"更多是请求层面的并发/批量，不是模型内部并行吐多条 token。

10.2 多条答案各自怎么处理

每一条 completion 都独立做三件事：

Student 采样时记录这条答案的 student_lps，即旧策略 logprob；
Teacher 对这条 prompt + completion 前向打分，得到 teacher_lps；
逐 token 算：

python 复制代码

reverse_kl = student_lps - teacher_lps
advantages = -kl_penalty_coef * reverse_kl

然后每条 completion 单独构造一个 Datum。所以 group_size=2 时，同一个 prompt 会产生两个 Datum，各自带自己的逐 token advantage。OPD 不做 GRPO 那种组内中心化，每条答案直接和 Teacher 比，不和其他答案比高低。

10.3 为什么采样多条

主要为了：

扩大同一个 prompt 下 Student 可能走到的状态覆盖；
降低这个 prompt 上 reverse KL 估计的方差；
让梯度更稳。

代价是线性增加 Student 采样成本、Teacher 前向打分成本、存储和反向传播成本。所以不是越多越好，实践中靠"多 prompt × 多 step"累积样本。

10.4 group_size=1 也能跑

Monte Carlo 估计 N=1N=1N=1 也是无偏估计，只是方差大。OPD 的 advantage 是逐 token 直接算的，不需要组内比较，所以 group_size=1 也能工作。设大只是降低方差，不是算法必需。

和 GRPO 对比：GRPO 的 advantage 是组内相对的，需要同一 prompt 下多条回答算组内均值和标准差，所以通常 group_size 至少 2，常见 4、8。OPD 不一样，不需要组内比较。

11. 多条 completion 的梯度如何回传

结论：

代码层面不是"每条 completion 分别 backward 再更新"，而是把整个 batch 的所有 Datum 收集起来，一次性调用 forward_backward(datums, ...)，做一次反向传播、一次 optim_step。

代码：

python 复制代码

datums = []

for sequence in result.sequences:
    ...
    datums.append(build_opd_datum(...))

fwd_bwd = training_client.forward_backward(
    datums,
    loss_fn="importance_sampling",
)
optim = training_client.optim_step(adam)

不是：

text 复制代码

completion 1 → backward → step
completion 2 → backward → step

而是：

text 复制代码

completion 1 ┐
completion 2 ├→ 一起 forward_backward → 一次 step
completion 3 ┘

11.1 数学上等价于分别回传再相加

设第 iii 条 completion 的损失是 LiL_iLi。总损失通常是：

L=∑iLi或L=1N∑iLi L=\sum_i L_i \quad \text{或} \quad L=\frac{1}{N}\sum_i L_i L=i∑Li或L=N1i∑Li

因为梯度算子是线性的：

∇θL=∇θ∑iLi=∑i∇θLi \nabla_\theta L = \nabla_\theta \sum_i L_i = \sum_i \nabla_\theta L_i ∇θL=∇θi∑Li=i∑∇θLi

所以"把所有 LiL_iLi 加起来回传一次"和"每条 LiL_iLi 分别回传，再把梯度加起来"在数学上等价。区别主要在工程实现和数值细节，不在梯度方向本身。

11.2 为什么不是每条分别 backward 再 step

如果每条 completion 都单独 backward → optim_step → backward → optim_step：

低效：每次 backward + step 都有开销；
破坏 on-policy 假设 ：这一批 completion 都是用同一个旧策略 πθold\pi_{\theta_{\text{old}}}πθold 采出来的，应该基于同一版参数算梯度、更新一次。如果中途 step 了，后面那条 completion 的旧策略 logprob 就和当前参数不一致了，importance sampling 的语义会乱。

所以正确做法是：同一批数据，先算总损失/总梯度，再更新一次。

11.3 注意细节

loss_fn="importance_sampling" 内部可能对 batch 内所有 token 做平均、求和，或按 token 数加权。文章没有展开这个 loss 的内部归一化细节。这不影响"一次 backward"的结论，但会影响不同长度 completion 的相对权重，实际使用时最好查该 loss_fn 的文档或源码。
如果训练里有梯度裁剪，通常是在总梯度算完之后裁剪。这时"加在一起回传再裁剪"和"分别回传再相加再裁剪"可能不完全等价，因为裁剪是非线性的。
每条 Datum 里，prompt 区间的 logprobs 和 advantages 都是 0 占位，所以虽然一个 batch 里有多条 completion，梯度只来自 completion token，prompt token 不贡献损失。

12. OPD 与强化学习、离线蒸馏的区别

12.1 蒸馏 vs 强化学习

更准确的说法：

蒸馏（包括 OPD）的学习目标来自 Teacher 的分布；强化学习的学习目标来自 reward / 环境反馈。新旧模型概率差在 RL 里主要是"校正或约束"用的，不是学习目标本身。

项目	OPD / 蒸馏	强化学习
学习目标来自	Teacher 分布	reward / 环境反馈
新旧概率差的作用	直接构造逐 token advantage	重要性采样校正、clip、KL 约束
Teacher / reference 的作用	提供目标分布	通常用于约束不要跑太偏，不是目标本身
旧策略 logprob 的作用	直接进入 advantage 公式	作为重要性采样分母或约束项

OPD 的 advantage：

$$

A_t=-\beta\left(
\log\pi_{\theta_{\text{old}}}(y_t|x,y_{<t})

\log\pi_T(y_t|x,y_{<t})

\right)

这里 log⁡πT\log\pi_TlogπT 是 Teacher 给的目标概率，log⁡πθold\log\pi_{\theta_{\text{old}}}logπθold 是 Student 采样时的旧概率。

强化学习里，新旧策略的概率比通常是：

πθ(yt∣x,y<t)πθold(yt∣x,y<t) \frac{\pi_\theta(y_t|x,y_{<t})} {\pi_{\theta_{\text{old}}}(y_t|x,y_{<t})} πθold(yt∣x,y<t)πθ(yt∣x,y<t)

这个比值在 PPO 里叫 importance ratio，作用是重要性采样校正和 clip，防止新策略一步走太远。但真正告诉模型"这个动作好不好"的，是 advantage，而 advantage 来自 reward。

所以：

蒸馏：老师告诉你该往哪学。

强化学习：裁判告诉你做得好不好。

OPD：学生自己做题，老师逐字批改，但批改方式借用了强化学习的训练流程。

12.2 OPD 里仍然有"新旧模型"

OPD 里仍然有旧策略和新策略：

python 复制代码

student_lps = [float(value) for value in sequence.logprob]

这里的 student_lps 就是 Student 旧策略 πθold\pi_{\theta_{\text{old}}}πθold 在采样时给出的 logprob。Student 更新之后，参数变成 πθ\pi_\thetaπθ，下一轮又会有新的旧策略。

所以更准确的说法是：

OPD 不是"没有新旧模型"，而是"新旧关系发生在 Student 自己身上，Teacher 是额外引入的固定参照模型"。

三层关系：

text 复制代码

Teacher（固定，不更新）
        ↑ 提供逐 token 目标 logprob
Student 旧策略（采样时用的权重）
        ↓ 采样得到 completion
Student 新策略（当前要更新的权重）
        ↑ 用 importance sampling 从旧策略更新到新策略

OPD 可以看成：

PPO 式 on-policy 训练框架 + 一个固定的 Teacher 作为逐 token 目标分布。

12.3 和 GRPO 的数据来源对比

原文表 8.2：

项目	GRPO	OPD
轨迹来源	当前策略	当前 Student
反馈来源	环境或规则判题器	Teacher 概率分布
信号粒度	常见为整条轨迹 reward	逐 completion token
是否需要标准答案	通常需要可验证结果	prompt-only 数据即可
是否需要额外模型	可省去 Value Model	需要 Teacher 前向计算
主要风险	reward hacking、退化组	Teacher 不兼容、能力差距不合适

通俗说：GRPO 更像"做完一整道题，裁判给分"；OPD 更像"老师跟着学生写的每个字批注"。OPD 不需要标准答案，但非常依赖老师靠谱且和学生兼容。

13. 运行、异步版与指标观察

13.1 同步版运行

bash 复制代码

python docs/chapter8/opd/01-demo-sync.py \
    --steps 1 \
    --batch-size 1 \
    --group-size 1 \
    --max-tokens 512 \
    --sample-size 20 \
    --num-shards 1 \
    --swanlab-mode disabled

13.2 异步版运行

bash 复制代码

python docs/chapter8/opd/02-demo-async.py \
    --steps 1 \
    --batch-size 4 \
    --group-size 2 \
    --max-tokens 512 \
    --sample-size 20 \
    --num-shards 1 \
    --swanlab-mode disabled

异步 OPD 同时利用两层并发：

训练环节	同步版	异步版
创建 Student / Teacher	同步 client 接口	`create_lora_training_client_async()` 与 `create_sampling_client_async()`
Student rollout	逐 prompt 调用 `sample().result()`	batch 内并发 `sample_async()`
Teacher 打分	逐 completion 调用 `compute_logprobs().result()`	同一 prompt 内并发 `compute_logprobs_async()`
汇总 batch	顺序追加 Datum	两层 `asyncio.gather()` 后统一汇总
Student 更新	同步 future + `.result()`	async API 提交后再次 await future
程序入口	`main(args)`	`asyncio.run(main(args))`

正文只展开同步版，因为它能直接呈现"Student 采样 → Teacher 打分 → reverse KL → Student 更新"的因果顺序。异步版保留同样的 token 对齐、advantage 和优化目标，只是把等待 I/O 的时间重叠起来。

13.3 模型与指标

示例默认使用 Qwen/Qwen3.5-4B 作为 Student，Qwen/Qwen3.6-27B 作为 Teacher。实际可用模型应以当前 PyTRIO 服务返回的支持列表为准。

训练时应同步观察：

opd/reverse_kl_mean：Student 采样 token 与 Teacher 概率的平均差异；
opd/reverse_kl_std：不同 token 蒸馏信号的离散程度；
data/completion_tokens_total：本 step 完成 Teacher 打分的 token 总量；
data/completion_tokens_per_second：整步 OPD 的有效 completion token 吞吐；
trainer/*：PyTRIO 返回的训练指标；
独立任务评测与通用能力评测：确认蒸馏后的真实能力变化。

reverse KL 下降说明 Student 的采样分布正在接近 Teacher，但任务正确率仍需通过固定的 held-out 评测确认。

出现 loss 波动或能力下降时，应依次检查：token 对齐、Teacher 质量、KL 系数、学习率、sampler 刷新频率和 completion 截断比例。

14. 常见误区与检查清单

14.1 常见误区

误区	正确理解
OPD 里 Teacher 会生成答案	不会。Teacher 只评价 Student 生成的同一条 completion
Teacher 也参与反向传播	不参与。Teacher 只前向计算 logprob，参数固定
OPD 对齐每一层 hidden state	本文 OPD 对齐的是输出层 token 概率分布，不是每层隐藏变量
Student 必须由 Teacher 剪枝得到	不一定。常见做法是选同系列小模型，本文用 LoRA 微调
group_size 必须大于 1	不必须。OPD 不需要组内比较，group_size=1 也能跑
多条 completion 各自 step 一次	不是。所有 Datum 收集后一次 forward_backward、一次 optim_step
蒸馏和强化学习完全一样	不一样。蒸馏目标来自 Teacher，RL 目标来自 reward
OPD 里没有新旧模型	有。新旧关系发生在 Student 自己身上，Teacher 是额外固定参照
logprob 切片只靠长度检查就够	不够。长度检查挡不住"长度对但语义错一位"，要用短序列探针确认
两边的 chat template 可以各套各的	不可以。只由 Student tokenizer 渲染一次，Teacher 直接评价同一串 token

14.2 训练前检查清单

tokenizer 对齐：同一段文本两边 encode 结果逐元素相等；
chat template 只渲染一次：由 Student tokenizer 渲染，Teacher 不重新渲染；
add_generation_prompt=True：prompt 末尾有"轮到 assistant"的标记；
logprob 位置约定已确认：用短序列探针测试，不靠长度猜；
completion 切片长度正确 ：len(completion_logprobs) == len(completion_ids)；
没有 None：Teacher 返回的 completion logprob 没有缺失；
Teacher 参数固定 ：Teacher 不调用 forward_backward() / optim_step()；
旧策略 logprob 来自采样时 ：student_lps 是旧策略 πθold\pi_{\theta_{\text{old}}}πθold 的 logprob；
多条 completion 一次更新 ：所有 Datum 收集后一次 forward_backward、一次 optim_step；
观察 reverse KL 与独立评测：reverse KL 下降不等于任务正确率一定上升。

15. 一句话总结

OPD 不是让学生抄老师的标准答案，而是学生自己写，老师逐字批改；学生根据"我比老师更偏爱还是更不偏爱这个字"来调整每个 token 的概率。

更完整地说：

OPD 借用了强化学习的 on-policy 采样和 importance sampling 框架，但学习目标来自固定 Teacher 的逐 token 概率分布，而不是环境 reward。Student 自己仍有旧策略和新策略的区别；Teacher 是额外引入的固定参照模型。多条 completion 的损失会收集成一个 batch，一次反向传播、一次更新；数学上等价于各条损失分别求梯度后相加。tokenizer、chat template、logprob 位置和 prompt/completion 边界都必须严格对齐，否则训练信号会变成噪声。

On-Policy Distillation（OPD）学习文档

目录

1. OPD 一句话总览

2. 核心术语

3. 从离线蒸馏到 OPD

3.1 离线蒸馏

3.2 OPD 的做法

3.3 OPD 的用途

4. Reverse KL 与逐 token 训练信号

4.1 怎么理解 dtd_tdt

4.2 reverse KL 的特点

4.3 Monte Carlo 估计是什么

5. Teacher 必须评价同一条 Student 轨迹

5.1 代码：Teacher 对同一条轨迹打分

6. tokenizer 兼容与 chat template

6.1 tokenizer 兼容是什么意思

6.2 为什么必须一样

6.3 chat template 是什么

6.4 代码中的 chat template 使用

6.5 在 OPD 里为什么强调"prompt 只渲染一次"

6.6 add_generation_prompt=True 的作用

7. logprob 位置对齐与"一位偏移"

7.1 位移从哪里来

7.2 两种常见返回约定

7.3 错一位会发生什么

7.4 文章代码为什么写 all_logprobs[len(prompt_ids):]

7.5 如何保证位置对齐

8. Student 网络一般怎么实现

9. 同步版 OPD 代码流程

9.1 数据

9.2 创建 Student 与 Teacher

9.3 构造逐 token OPD Datum

9.4 生成轨迹并计算 reverse KL

9.5 更新 Student

10. 采样条数：batch_size 与 group_size

10.1 是"同时输出"吗

10.2 多条答案各自怎么处理

10.3 为什么采样多条

10.4 group_size=1 也能跑

11. 多条 completion 的梯度如何回传

11.1 数学上等价于分别回传再相加

11.2 为什么不是每条分别 backward 再 step

11.3 注意细节

12. OPD 与强化学习、离线蒸馏的区别

12.1 蒸馏 vs 强化学习

$$

12.2 OPD 里仍然有"新旧模型"

12.3 和 GRPO 的数据来源对比

13. 运行、异步版与指标观察

13.1 同步版运行

13.2 异步版运行

13.3 模型与指标

14. 常见误区与检查清单

14.1 常见误区

14.2 训练前检查清单

15. 一句话总结

7.4 文章代码为什么写 `all_logprobs[len(prompt_ids):]`