Add AsyncDistillationTrainer - #6705
Conversation
Async on-policy distillation: a background rollout worker generates the student's own completions and scores them against a teacher served over HTTP, decoupling generation from training the way AsyncGRPOTrainer does for GRPO. Also supports MOPD (multi-teacher on-policy distillation) via per-sample teacher_id routing across several teacher_server_urls entries.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a09d0bfb5e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
- Remove the teacher vocab_size check: it called GET /v1/models on the teacher's server, but trl vllm-serve doesn't expose that route (it's the standard vLLM OpenAI server's endpoint, not this custom FastAPI app). No way to fix it without adding new shared server infra, so reverting rather than shipping a check that always fails. - Fix the paper citation: it was misattributed to the MiniLLM authors instead of Agarwal et al. (the actual 2306.13649 authors), which would have shown up as a false citation on every model card. - Note that async execution and MOPD aren't described in 2306.13649 (a synchronous, single-teacher paper); add the actual MOPD paper citation and a paper_index.md entry for it. - Document the TokenBudgetBatcher/max_steps interaction as a known, pre-existing limitation shared with AsyncGRPOTrainer rather than fixing it in isolation here. - Drop a dangling reference to a planning doc that doesn't exist in this repo.
MOPD is Stage 3 of a 3-stage pipeline (SFT, per-domain RL experts, then fusion) — AsyncDistillationTrainer only implements the fusion stage, with pre-trained teachers. The paper's Stage 3 also uses reverse KL (beta=1.0), not this trainer's default beta=0.0.
…mory Project hidden states through lm_head in fixed-size chunks (checkpointed) instead of materializing the full [seq_len, vocab_size] logits tensor at once, mirroring DistillationTrainer's own chunked JSD path. Also fixes a missing autocast around the bypassed base_model forward call, and a stale test fixture bug (tokenizer.vocab_size excludes added special tokens like eos_token_id, causing out-of-bounds indices in synthetic teacher data).
Calling unwrapped_model.base_model(...) directly to bypass lm_head skips the wrapper's forward, so DDP's prepare_for_backward and FSDP's parameter materialization never run for the loss. Mirrors DistillationTrainer's identical need for its own chunked JSD path.
jsd/{teacher_id} collided with the flat jsd metric name in dashboards
that treat / as a grouping separator; teacher_entropy/{id} already used
the teacher_ prefix, so jsd/{id} was the odd one out.
| if len(self.teacher_server_urls) == 1: | ||
| ((teacher_id, url),) = self.teacher_server_urls.items() | ||
| return teacher_id, url | ||
| teacher_id = row.get("teacher_id") |
There was a problem hiding this comment.
If I understand this couples the mapping teacher_id->vllm_server_url to the dataset having a specific column teacher_id in the format ?
I wonder if we can check/enforce this somehow at runtime like checking that the dataset has this columns and the unique values map correctly to the teacher 🤔 ? wdyt ?
Not a strong requirement
There was a problem hiding this comment.
Yep, with more than one teacher each row needs a teacher_id matching a key in teacher_server_urls. I kept it the same as AsyncGRPO, which routes multi-env runs off an environment column and checks it per row in the worker rather than validating the dataset upfront. A missing or unknown id raises instead of silently picking a teacher. I did align the error with the GRPO one (ValueError, same wording), so it reads the same in both.
| "response_format": "json", | ||
| } | ||
| output = await _retry_on_http_error( | ||
| lambda: self._post(teacher_server_url, "/get_sequence_logprobs/", payload, self.request_timeout), |
There was a problem hiding this comment.
Q: I was unaware of /get_sequence_logprobs endpoint. Is this in stock vllm ? :o
There was a problem hiding this comment.
Nope, that one is ours, not stock vllm. It only exists on trl vllm-serve (trl/scripts/vllm_serve.py), wrapping vllm prompt_logprobs for teacher forcing. Same endpoint sdft/sdpo/server_distillation already use.
AmineDiro
left a comment
AmineDiro
left a comment
There was a problem hiding this comment.
Implementation looks good to me. Well done @kashif 👏🏼
The only design question is about separating the RolloutWorker in a separate process. We've done this in AsyncGRPOTrainer in PR mainly because recursive_parse / accuracy_reward were locking the GIL. In this AsyncDistillationTrainer I dont see anything similar to that so maybe a simple threaded Worker is the simpler + remove serialization/deserialization cost ?
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit d4060f3. Configure here.
2 participants
Async on-policy distillation, architected like
AsyncGRPOTrainer: a background rollout worker generates the student's own completions and scores them against a teacher served over HTTP, so generation and training overlap instead of alternating. The teacher is never loaded locally, just a vLLM server URL.Also supports MOPD (multi-teacher on-policy distillation) — pass more than one entry in
teacher_server_urlsand route each sample to a teacher via ateacher_idcolumn (e.g. a math teacher and a code teacher, each served independently).This went through several rounds of GPU validation in a staging fork before landing here, and the multi-teacher routing design was cross-checked against how NeMo RL, verl, and miles do the same kind of per-sample teacher routing.
Adds:
trl/experimental/async_distillation/(config, trainer, rollout worker, vLLM client, weight transfer)tests/experimental/test_async_distillation_trainer.pyexamples/scripts/async_distillation.py(single-teacher) andasync_distillation_mopd.py(math+code MOPD demo)docs/source/async_distillation_trainer.md+ toctree entry, one-line cross-ref fromdistillation_trainer.mdmake precommitpasses. Test suite passes (26 passed / 8 GPU-gated skipped without a GPU); also ran the GPU-gated tests viasrun --gres=gpu:1— one pre-existing failure reproduces there, but it's unrelated to this PR (it also reproduces on unmodifiedasync_grpotests on the same node, looks like a tiny-test-fixture/environment issue rather than anything this PR changes).Happy to open a tracking issue first if that's the preferred process for a new experimental trainer — let me know.
Note
Medium Risk
Large new experimental training path (async queues, NCCL weight sync, distributed FSDP2-only) with complex loss math; well-tested but operationally sensitive to vLLM/transformers versions and multi-process failures.
Overview
Introduces
AsyncDistillationTrainerundertrl/experimental/async_distillation/, modeled onAsyncGRPOTrainer: a spawned rollout worker generates student completions via vLLM while the main loop trains on queued samples, overlapping generation and updates.The teacher is not loaded in-process—only
trl vllm-serveURLs and/get_sequence_logprobs/sparse top-k scoring. Loss is generalized JSD (same family as sync distillation / server distillation), with chunkedlm_headprojection and beta-dependent support narrowing. Student weights sync to vLLM over NCCL on a schedule;max_stalenessdrops outdated rollouts.MOPD is supported via multiple
teacher_server_urlsand a per-rowteacher_idcolumn (strict routing, no silent fallback). AddsAsyncDistillationConfig, rollout worker, vLLM client, weight transfer client, 646-line test suite, single- and multi-teacher example scripts, dedicated docs + toctree entry, paper index MOPD blurb, and a cross-link from the sync distillation doc.Reviewed by Cursor Bugbot for commit d4060f3. Bugbot is set up for automated code reviews on this repo. Configure here.