Marginalia 旁注旁观
中文

Marginalia

中文

Multi-teacher distillation is a token budget

TRL's multi-teacher on-policy distillation example ran 60 steps. Instruction following held 20 percent of the prompts and got 1 to 3 percent of the trained tokens, and it was the only domain that barely moved. Open-MOPD calls this token budget misallocation, and fixing it raised the recoverable headroom from 35.6 to 83.4 percent.

AI 10 min LB 中文

中文版是这篇文章的默认版本。 阅读中文版 →
This English page stays available at this URL; the Chinese version is the one we maintain first.

TRL’s multi-teacher on-policy distillation example ran 60 steps. Instruction following held 20 percent of the prompts and got 1 to 3 percent of the trained tokens, and it was the only domain that barely moved. Open-MOPD calls this token budget misallocation, and fixing it raised the recoverable headroom from 35.6 to 83.4 percent.

The film runs two minutes, voiced in English, with the captions burned into the picture.

1. After sixty steps, only one curve was flat#

In November 2026 Sergio Paniego, at Hugging Face, wired the three experts released with Open-MOPD into TRL’s AsyncDistillationTrainer. The student is the SmolLM3-3B SFT checkpoint. The three teachers were each trained with GRPO on one domain, math, code and instruction following. The run is 60 steps, about one percent of the training the paper did.

Two things show up in the logs. teacher_token_frac puts instruction following at 1 to 3 percent of the trained tokens per step while it holds 20 percent of the prompts. teacher_jsd falls steadily for the math and code teachers and stays almost flat for the instruction-following one.

The evaluation follows the same shape.

ModelAIME24AIME25IFEval
MixSFT starting point17.519.267.5
60-step MOPD24.222.573.0
Paper’s final model19.225.075.8

The AIME24 column needs a note. The 60-step run scores 24.2 against the paper’s final 19.2, which reads like a short run beating a long one. Four samples per problem put a few points well inside the noise, so that cell supports no conclusion. The stable signal is IFEval, 67.5 at the start, 73.0 after 60 steps, still 2.8 points short of the paper’s model, and that is the domain with the fewest tokens.

2. Three stages, and the third only scores#

Multi-teacher on-policy distillation (MOPD) is a three-stage post-training recipe, named in January 2026 by Xiaomi in the technical report for MiMo-V2-Flash.

Stage one is general SFT, where one model learns the basics of every domain on mixed data, and that checkpoint is the starting point for everything after. Stage two trains one RL expert per domain, one for math, one for code, one for instruction following, and none of them has to care about the others. Stage three fuses them. The student answers prompts from every domain, the expert of each prompt’s domain scores the answer token by token, and the student moves toward that teacher’s distribution.

Three properties explain why it works.

PropertyConsequence
The student generates its own samplesThe teacher grades text the student actually wrote. Plain SFT on teacher outputs only trains the student on text the teacher wrote, while at inference time it continues its own text, and one early mistake lands it in a distribution it never saw
Every token carries a signalRL gives one reward per answer, this gives the teacher’s distribution at every position, and the loss is the divergence between student and teacher over the student’s own tokens
Each prompt goes to exactly one teacherA math prompt is scored by the math expert, with no averaging and no ensemble

Teachers usually match the student in size and share its starting checkpoint. What makes them useful is one pass of RL on a single domain. The recipe buys a schedule, splitting multi-domain training into training domains separately and merging them afterwards.

The industry has been on this for a while. Kimi K3 uses nine experts, split by domain and by reasoning effort. Nemotron 3 Ultra uses more than ten teachers across two rounds of MOPD, retraining the teachers from the merged model before the second round. DeepSeek-V4 replaced the mixed RL stage of V3.2 with more than ten domain teachers. Mistral Large 4 splits differently, three teachers by reasoning effort, high, medium and low, with domain no longer the axis.

Figure 1 — The three stages and the routing in the fusion loop. The student generates, the teacher only reads, and each prompt goes to one expert
tap to enlarge

3. The paper blames the budget#

Open-MOPD (arXiv 2608.19098, August 19) does something rare. It takes routing out of the picture first. The benchmark sits on SmolLM3-3B-Base with oracle routing, so every prompt reaches the right expert, and whatever gap remains comes from fusion itself.

The result is a pair of numbers. Against a domain-routed oracle ensemble, standard MOPD recovers 35.6 percent of the available headroom. With the paper’s three fixes it recovers 83.4 percent. Same student and same teachers, and the 47.8 points between them sit entirely in the training procedure.

The paper names the cause plainly. The failure comes from misallocation of the token-level optimization budget, not from gradient conflict. The authors separate three drivers. Sequence lengths differ by orders of magnitude between domains, thousands of tokens for a math answer and a few hundred for an instruction-following one. Domains converge at different rates, so the fast ones take the gradient first. Rewards go stale in an asynchronous pipeline, where an old sample meets a newer policy.

The three fixes are token-share balancing, gap-aware dynamic budget allocation and a student reward refresh. The first two handle the same thing. A domain’s share of the loss is its share of the tokens, and that share can be set by hand.

4. Averaging over tokens means voting by answer length#

The shape of the loss explains it. The divergence is computed at every token and averaged over tokens, so a domain’s say equals its trained tokens divided by all trained tokens.

Prompt share and token share can therefore differ by an order of magnitude. The example data splits math, code and instruction following 40/40/20, so instruction following is a fifth of the questions and one to three percent of the text. A math or code answer is a long reasoning chain, thousands of tokens. An instruction-following answer is often “write two paragraphs separated by asterisks”, a few hundred. Math and code answers keep getting longer over those 60 steps, and the instruction-following share keeps falling.

Figure 2 — Prompt share against trained-token share inside one batch. Instruction following holds a fifth of the questions and produces one to three percent of the text
tap to enlarge

The TRL docs put this in the metric description. teacher_token_frac/<id> records the share of scored tokens each teacher took, and the doc says a teacher starved of rows still reports a healthy teacher_jsd/<id>. Routing skew stays invisible without that metric. Watching only the absolute loss, the instruction-following domain looks like it is learning the whole time, while each step gives it the gradient of a few dozen samples.

5. Three defaults inside TRL#

Running MOPD in TRL, three values in the config decide how the budget gets recorded.

beta defaults to 0.0, forward KL, mean-seeking. Stage three of the MOPD paper uses 1.0, reverse KL, mode-seeking. The docs give that difference its own paragraph. Copy the official example and delete the beta line, and the run optimizes a different objective while the logs keep falling.

teacher_top_k defaults to 8, and it only shapes anything at beta=0.0. At a nonzero beta the support the teacher sends back narrows to two candidates, the teacher’s own top-1 and the token the completion actually has, because those two identities are the only ones the HTTP protocol guarantees a logprob for. The example sets teacher_top_k=1 to keep each teacher’s response small, and under that beta the value no longer touches the loss.

Figure 3 — beta decides how many candidates come back from the teacher. Forward KL uses the full top-k plus a tail bucket, reverse KL uses two
tap to enlarge

The third item is not a default. It is a hard constraint. Every teacher must share the student’s tokenizer. Completions travel to the teacher as raw token ids, and the candidate ids coming back index the student’s own vocabulary directly. A teacher with a different vocabulary trains the student against the wrong tokens, and it only raises an error when its vocabulary is the larger one.

# examples/async_distillation_mopd/async_distillation_mopd.py
config = AsyncDistillationConfig(
    teacher_server_urls={
        "math": "http://localhost:8001",
        "code": "http://localhost:8002",
        "if":   "http://localhost:8003",
    },
    beta=1.0,         # reverse KL, as in stage three of the paper; the default is 0.0
    teacher_top_k=1,  # with a nonzero beta the support is already two candidates
    max_completion_length=16384,
    vllm_server_timeout=3600,
)

6. Five processes on four GPUs#

The fusion stage runs five things at once. The student’s own vLLM server generates, the teacher servers only score, and the trainer process pulls from a queue in between. The example spreads those five across four H200s.

GPURunsWhy there
0The student’s vLLM serverIt receives new weights every weight_sync_steps
1The math and instruction-following teachersTwo small models on one card, started one after the other
2The code teacherCode prompts are the longest, so it runs alone
3The trainerForward, backward and the optimizer

Figure 4 — Five processes on four cards. The three teachers never receive weights; only the student&rsquo;s vLLM server is updated by NCCL
tap to enlarge

The asynchronous pipeline hides its cost in two defaults. weight_sync_steps is 1, so every step pushes the student’s new weights into vLLM. max_staleness is 4, so a sample that waits through more than four weight updates before it is trained on gets discarded. Together they set how deep the queue runs and how much stale data reaches the gradient.

The LoRA path needs its own arithmetic. The default weight sync merges the adapter into the base model and ships the whole thing, about 15 GB per sync for a 7B model. Start the student server with --enable-lora and it sends only the adapter, about 160 MB at rank 32. The saved bandwidth has a price, since vLLM generates 1.33 to 1.39 times slower per token when it serves a LoRA adapter, and generation usually dominates wall-clock time.

7. Numbers and their sources#

MetricValueSource
Headroom recovery, standard MOPD35.6 percentPaper abstract (arXiv 2608.19098)
Headroom recovery, Open-MOPD fixes83.4 percentSame
Instruction-following prompt share / token share20 percent / 1 to 3 percentHugging Face blog, 60-step training logs
MixSFT AIME24 / AIME25 / IFEval17.5 / 19.2 / 67.5Same, four samples per AIME problem, strict prompt-level IFEval
60-step MOPD, same three24.2 / 22.5 / 73.0Same
Paper’s final model, same three19.2 / 25.0 / 75.8Same
beta default / value used in stage three0.0 / 1.0TRL docs
teacher_top_k default8TRL docs
What teacher_top_k above 20 needs--max-logprobs -1 on the teacher serverTRL docs
Sparse support in neighbouring frameworksmiles 16, EasyOPD 64TRL docs
LoRA weight sync, 7B, rank 3215 GB / 160 MBTRL docs
LoRA serving cost in generation1.33 to 1.39 times slowerTRL docs

The two columns come from different places. The 35.6 and 83.4 percent figures are what the paper reports, with no third-party reproduction. The 60-step numbers are one run by the blog’s author, 60 steps long, and they support no claim about a short run beating a long one.

8. What a reader can check#

All four checks fit in an hour and need no extra hardware.

CheckWhereWhat to look at
Routing skew is realteacher_token_frac/<id> in the training logsWhether instruction following sits far below its prompt share
The default changes the objectiveThe beta entry in the AsyncDistillationConfig docsThe default 0.0 is forward KL, stage three of the paper uses 1.0
top-k does nothing at a nonzero betaThe support description in the same entryThe sentence where the support narrows to two candidates
A tokenizer mismatch fails silentlyThe note under teacher_server_urlsNo error when the teacher’s vocabulary is the smaller one

9. Two questions I cannot answer#

The first is how the budget should be counted when teachers outgrow the student. In the paper and the example, teachers match the student in size and start from the same checkpoint, so token share and capacity share are roughly interchangeable. With a large teacher and a small student, one token carries a different amount of information, and nobody has measured whether averaging over tokens stays fair.

The second is routing. Open-MOPD removes the ambiguity with oracle routing before it attributes the gap to the budget. There is no oracle in production. A classifier decides which domain a prompt belongs to, a misrouted prompt is charged to another teacher’s account, and the paper does not say how that error mixes into token share.

References#

Search 40 essays

Browse all writing