Async Distillation

[!IMPORTANT] This trainer requires vllm>=0.22.0 and transformers>=5.2.0. For distributed training, only FSDP2 is supported (DeepSpeed ZeRO is not).

Currently, vllm and transformers have conflicting dependency constraints. To work around this, install vLLM first and then force-install transformers:

bashCopy

pip install 'vllm>=0.22.0'
pip install 'transformers>=5.2.0' --no-deps

Overview

[AsyncDistillationTrainer] is the async counterpart to [~trl.experimental.distillation.DistillationTrainer], architected like [AsyncGRPOTrainer]: a background rollout worker generates the student's own on-policy completions and scores them against a teacher, while training proceeds concurrently instead of alternating between generation and gradient updates. Unlike the synchronous trainer, the teacher is never loaded locally — only a vLLM server URL is needed, so the teacher can run on entirely separate hardware from the student and trainer, or even be a much larger model than would otherwise fit alongside the student.

compute_loss minimizes a generalized Jensen-Shannon Divergence between the student's and teacher's per-position token distributions (beta=0.0 is forward KL, beta=1.0 is reverse KL, values in between interpolate), the same objective [~trl.experimental.distillation.DistillationTrainer] and [~trl.experimental.server_distillation.ServerDistillationTrainer] use. Training is always on-policy: the student generates every completion it trains on.

Multi-teacher on-policy distillation (MOPD)

MOPD is not part of the 2306.13649 paper this trainer's core objective is based on (which describes a synchronous, single-teacher setting); it's a separate method, described in MOPD: Multi-Teacher On-Policy Distillation for Capability Integration in LLM Post-Training. There, MOPD is the third of three stages — general SFT, then independent per-domain RL training of one expert per domain, then MOPD fuses those frozen experts into a single student. AsyncDistillationTrainer implements that third, fusion stage only: the per-domain expert teachers must already exist (e.g. trained separately with [GRPOTrainer]/[RLOOTrainer]) and be served over HTTP before you point teacher_server_urls at them. The paper's own Stage 3 uses reverse KL (beta=1.0), not this trainer's default beta=0.0 (forward KL).

teacher_server_urls accepts more than one entry. With a single entry, every sample is scored by that one teacher (plain on-policy distillation). With multiple entries, each training sample's teacher_id column selects which teacher scores it — for example, routing math prompts to a math-specialist teacher and code prompts to a code-specialist teacher, each served independently. Each sample is dispatched to its one matching teacher, never averaged or ensembled across teachers. A sample with a missing or unmapped teacher_id raises rather than silently falling back to the wrong teacher. See examples/async_distillation_math/async_distillation_mopd.py for a runnable two-teacher example.

[!WARNING] Every teacher must share the student's tokenizer. Completions travel to the teachers as raw token ids, and the candidate ids a teacher reports back index the student's own vocabulary directly in compute_loss. Teachers from the same model family as the student (as in the MOPD example, where a Qwen2.5 student is fused from Qwen2.5 and Qwen2.5-Coder experts) satisfy this; a teacher with a different vocabulary trains the student against the wrong tokens, silently unless its vocabulary is larger than the student's.

How it differs from [~trl.experimental.distillation.DistillationTrainer]

In [~trl.experimental.distillation.DistillationTrainer], the teacher is a locally loaded model: generation, teacher forward pass, and the gradient update all happen sequentially in the same process. [AsyncDistillationTrainer] separates these concerns the same way [AsyncGRPOTrainer] separates GRPO's rollout from its update:

• Rollout worker (background process) — generates completions from the student's vLLM server, sends the full sequence back to the resolved teacher's /v1/completions with prompt_logprobs for teacher-forced scoring (no new tokens generated by the teacher), and pushes ready-to-train samples into a queue.
• Training loop (main process) — pulls samples from the queue, computes the generalized-JSD loss, and updates the student's weights.

Because the teacher is scored over HTTP rather than a local forward pass, only a sparse, top-k slice of its distribution is ever transmitted (teacher_top_k), not the full vocabulary — see [AsyncDistillationConfig]'s beta and teacher_top_k documentation for exactly which candidates the wire protocol guarantees a teacher logprob for at each beta regime.

After every weight_sync_steps training steps, the updated student weights are transferred to its vLLM server via NCCL. As with [AsyncGRPOTrainer], generation runs ahead of training, so samples may reflect a slightly stale policy; max_staleness controls how many weight updates a sample can lag behind before being discarded.

Quick start

# train_async_distillation.py
from datasets import load_dataset
from trl.experimental.async_distillation import AsyncDistillationTrainer

dataset = load_dataset("trl-lib/DeepMath-103K", split="train")

trainer = AsyncDistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    train_dataset=dataset,
)
trainer.train()

The teacher server, the student's vLLM server, and the trainer must run on separate GPUs. Both are plain vllm serve instances, but they need different flags: the teacher is only ever scored (--logprobs-mode processed_logprobs so teacher_temperature reaches its returned logprobs, --max-logprobs -1 so teacher_top_k can exceed vLLM's default cap of 20), while the student is only ever generated from and needs NCCL weight transfer enabled so the trainer can push updated weights into it:

# Terminal 1: teacher server on GPU 0 (static, never updated)
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen2.5-1.5B-Instruct \
    --port 8001 \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

# Terminal 2: student's vLLM server on GPU 1 (dev mode + NCCL weight transfer are required)
CUDA_VISIBLE_DEVICES=1 VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen2.5-0.5B-Instruct \
    --port 8000 \
    --weight-transfer-config '{"backend":"nccl"}'

# Terminal 3: training on GPU 2
CUDA_VISIBLE_DEVICES=2 accelerate launch train_async_distillation.py

Logged metrics

A rollout passes through several stages before it becomes a gradient. Let's first take a look at the different stages, from a prompt to our final batch composition:

Dataset row = PROMPT (message list + optional `teacher_id`)
 └─ ROLLOUT            (1 prompt -> 1 student completion -> 1 teacher scoring call)
     ├─ generate       (student's vLLM /v1/completions, sampled)
     └─ score          (the routed teacher's /v1/completions with `prompt_logprobs`, teacher-forced)
         └─ SAMPLE     (prompt + completion + the teacher's sparse per-position distribution)

════════════ Trainer process boundary: `rollout_buffer` (mp.Queue) ════════════

         └─ SAMPLE (Pulled 1 at a time; dropped if staleness > `max_staleness`)
             └─ ROW            (Planner assigns it to one of `dp` rows, Σ Lᵢ²-balanced)
                 └─ MICRO-BATCH (`dp` rows, one per rank)
                     └─ PACKED ROW (1 concat sequence, `position_ids` reset per sample)
                         └─ FORWARD (`compute_loss`, bs=1, inter-rank padding stripped)
             └─ OPTIMIZER STEP (from `grad_accum` micro-batches)

Reading it top to bottom:

• a rollout is one prompt generated once and scored once. There is no group: distillation has no advantage baseline to compute across generations, so a prompt is not repeated and one rollout yields exactly one training sample.
• a sample is that completion plus the teacher's top-teacher_top_k candidates per completion position. This is what crosses the process boundary.
• past the boundary the trainer plans samples into rows (one per DP rank), packs each row into one concatenated sequence, and accumulates grad_accum micro-batches into a single optimizer step.

So we split our metrics based on which stage (or entity) they are about:

Everything the loss itself measures is logged unprefixed (jsd, entropy, teacher_entropy), with a teacher_*/<teacher_id> breakdown under MOPD — see the divergence below.

Let's also define some terms to make token metrics unambiguous:

• generated tokens are what the student produced (completions/*).
• forwarded tokens are every token the forward pass processes: the prompt and the generated tokens.
• trained tokens are the subset the loss is taken over where completion_mask == 1. Trained ≠ generated: a completion position the teacher scored no candidate for is masked out of the divergence but still forwarded.

How the batch numbers compose

A row is what one DP rank forwards in one micro-batch: several samples concatenated into a single sequence, with position_ids restarting at each sample boundary. The planner decides which samples land in which row (balancing Σ Lᵢ² so no rank straggles), and a data collator does the concatenating.

In all our metrics a "step" always means one FULL optimizer step, never a micro-batch. gradient_accumulation_steps micro-batches make one step, and every metric named _per_step is per optimizer step, matching global_step, logging_steps and the rest of the [Trainer] vocabulary. Where a metric is per micro-batch it says so (batch/microbatches_per_step) or it is a per-row quantity (batch/row_*, batch/samples_per_row).

One step therefore holds a fixed number of row-slots:

row-slots per step = gradient_accumulation_steps x world_size

which makes the batch metrics compose into a ladder you can check against each other:

sample                                     sample/forwarded_tokens_mean
  └─ packed into a ROW                     batch/samples_per_row, batch/row_tokens_mean
      └─ one row per DP rank               = one micro-batch
          └─ gradient_accumulation_steps micro-batches = one STEP
                                           batch/samples_per_step, batch/forwarded_tokens_per_step

So batch/samples_per_step ≈ row-slots x batch/samples_per_row, and likewise for tokens. The two sides agree only up to the variation between micro-batches inside the step — the per-step metrics are sums, the per-row ones are means — so expect a fraction of a percent, not an exact match.

batch/row_fill_frac is the one to watch when samples are long: a 10k-token sample tiles a 32k budget badly (three fit, four never do, so the packer often gets two and the row runs ~77% full), while 1k-token samples tile it almost perfectly. That is quantization, not a bug, and token_budget is the lever — bearing in mind that attention is O(L²) per sequence, so a fuller row of long sequences does not cost linearly more memory.

The rollout queue

The worker pushes scored samples into a queue and the trainer pulls from it. Four metrics describe that one queue, and they answer different questions:

The last two metrics (perf/rollout_wait_s and rollout/backpressure_s) are mirror images and are never both large! Reading them together with the queue size tells you which side is the bottleneck:

• queue near empty, perf/rollout_wait_s high → generation-bound. The trainer is starving; look at rollout/generated_tok_s, rollout/inflight and rollout/score_s.
• queue near full, rollout/backpressure_s high → trainer-bound. Generation is throttled and its output is aging in the queue, so watch sample/staleness_mean climb.
• both near zero → balanced !

Completions

A completion is what the student's vLLM server generated for one prompt.

Rollouts

A rollout is one prompt taken all the way through: generated by the student, then teacher-forced through its teacher. Because the teacher call is on the critical path of every rollout (unlike GRPO, where scoring is a separate loop over completed groups), a slow teacher shows up directly as rollout/duration_s.

Samples arriving from the queue

Each sample is one RolloutSample: a prompt, the student's completion, and the teacher's sparse scoring of it. The worker pushes them into rollout_buffer; everything below is measured on the trainer side, as it pulls them for training.

Batches

One micro-batch is world_size rows (remember packing flattens into 1 sequence), so one per DP rank, so that rank i forwards row i. gradient_accumulation_steps of those make one optimizer step. A step then covers gradient_accumulation_steps × world_size rows in total. Every _per_step metric below is a sum over the whole step and across every rank; the batch/row_* ones are means over the rows.

The divergence

What the objective itself measures, averaged over the trained tokens of the window. These are the numbers to watch for learning, as opposed to for throughput.

There is no per-teacher entropy: the student's entropy is a property of its own policy, not of which teacher scored the sample, so the blended metric already covers it.

Performance

Throughput and MFU are each reported twice, over the same optimizer step, differing only in what they divide by. The suffix names the denominator:

• _fwd_bwd divides by perf/fwd_bwd_s — the compute alone. How efficiently does the trainer run when it has data? If it is low, the trainer is the problem.
• _wall_clock divides by perf/step_s — the whole step, including the time spent waiting for rollouts. What fraction of the allocation actually became training? If this is far below the _fwd_bwd one, generation is probably the bottleneck.

The gap between them is perf/rollout_wait_s plus the optimizer and weight-sync time. But it's useful to look at both: looking only at _fwd_bwd hides the GPU-hours spent generating rollouts and scoring them, and quoting only _wall_clock could blame the trainer for the generator's or the teacher's latency.

Design philosophy

This trainer is intentionally kept minimal and is not meant to grow into a general-purpose solution. If you need a feature that is not supported, we recommend cloning the repository and adapting the trainer to your needs directly. New features will only be considered when there is significant community demand.

AsyncDistillationConfig

[[autodoc]] trl.experimental.async_distillation.AsyncDistillationConfig

AsyncDistillationTrainer

[[autodoc]] trl.experimental.async_distillation.AsyncDistillationTrainer