Distillation Trainer

Overview

The Distillation Trainer implements on-policy knowledge distillation as described in On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes by Rishabh Agarwal, Nino Vieillard, Yongchao Zhou, Piotr Stanczyk, Sabela Ramos, Matthieu Geist, and Olivier Bachem.

The abstract from the paper is the following:

Knowledge distillation (KD) is widely used for compressing a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, we introduce Generalized Knowledge Distillation (GKD). Instead of solely relying on a fixed set of output sequences, GKD trains the student on its self-generated output sequences by leveraging feedback from the teacher on such sequences. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which can be useful when the student lacks the expressivity to mimic the teacher's distribution.

The DistillationTrainer trains a smaller student model to match a teacher's next-token distribution on the student's own on-policy generations. It generates the student's completions on-policy (optionally vLLM-powered) and matches the teacher over its full next-token distribution with a memory-efficient chunked Jensen-Shannon divergence loss, so the teacher's dense distribution is never materialized in full.

This trainer was contributed by Carlos Miguel Patiño.

For an async counterpart that decouples generation from training and scores against a teacher served over HTTP instead of a local forward pass, see AsyncDistillationTrainer.

Quick start

This example demonstrates how to train a model using the distillation method. We distill a Qwen 2.5 0.5B Instruct model from a Qwen 2.5 1.5B Instruct teacher on the prompts from the UltraFeedback prompt dataset. You can view the data in the dataset here:

Below is the script to train the model.

python

# train_distillation.py
from datasets import load_dataset
from trl import DistillationTrainer

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

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

Execute the script using the following command:

bash

accelerate launch train_distillation.py

Looking deeper into the distillation method

On-policy knowledge distillation trains a student to reproduce a teacher's next-token distribution on completions the student generates itself, rather than on a fixed set of teacher outputs. Learning from its own generations lets the student correct its own mistakes, which generally outperforms off-policy distillation. This section breaks down how it works in practice, covering the two key steps: generating completions and computing the loss.

Generating completions

At each training step, the student generates a batch of completions for the sampled prompts.

Computing the loss

The loss is the generalized Jensen-Shannon divergence (JSD) between the student distribution $ p_S $ and the teacher distribution $ p_T $ over the generated completion tokens, interpolated by beta and defined as:

$$ \mathcal{L}\beta = \beta , \mathbb{D}{\mathrm{KL}}!\left[ p_T | p_M \right] + (1 - \beta) , \mathbb{D}_{\mathrm{KL}}!\left[ p_S | p_M \right], \qquad p_M = (1 - \beta) , p_S + \beta , p_T, $$

where $ p_M $ is the $ \beta $ -mixture of the two distributions. The endpoints reduce to the pure divergences: beta=0.0 gives the forward KL $ \mathbb{D}{\mathrm{KL}}!\left[ p_T | p_S \right] $ and beta=1.0 the reverse KL $ \mathbb{D}{\mathrm{KL}}!\left[ p_S | p_T \right] $.

In practice, the projection to vocabulary logits and the divergence are computed in chunks, so peak activation memory does not scale with the full vocabulary × sequence-length logits tensor. See Reducing Memory Usage.

Expected dataset type

The dataset should be formatted as a conversational prompt-only dataset. The student generates its own completions on-policy, so only the prompt is needed:

python

{"prompt": [{"role": "user", "content": "What color is the sky?"}]}

Logged metrics

While training and evaluating, we record the following metrics:

num_tokens: The total number of tokens processed so far, including both prompts and completions. When using tools, only non-tool tokens are counted.
step_time: The average time (in seconds) taken per training step (including generation).
completions/mean_length: The average length of generated completions. When using tools, only non-tool tokens are counted.
completions/min_length: The minimum length of generated completions. When using tools, only non-tool tokens are counted.
completions/max_length: The maximum length of generated completions. When using tools, only non-tool tokens are counted.
completions/mean_terminated_length: The average length of generated completions that terminate with EOS. When using tools, only non-tool tokens are counted.
completions/min_terminated_length: The minimum length of generated completions that terminate with EOS. When using tools, only non-tool tokens are counted.
completions/max_terminated_length: The maximum length of generated completions that terminate with EOS. When using tools, only non-tool tokens are counted.
completions/clipped_ratio: The ratio of truncated (clipped) completions.
tools/call_frequency: The average number of tool calls per completion in the generation batch. Logged only when tools are provided.
tools/failure_frequency: The fraction of tool calls that failed (the tool was not found, raised an exception, or the call type is unsupported). It is 0.0 when no tool was called. Logged only when tools are provided.
entropy: Average entropy of token predictions across generated completions (in nats). Not logged on the Liger fast path.

Customization

Speed up training with vLLM-powered generation

Generation is often the main bottleneck when training with on-policy methods. To accelerate generation, you can use vLLM, a high-throughput, low-latency inference engine for LLMs. To enable it, first install the package with

shell

pip install trl[vllm]

We support two ways of using vLLM during training: colocate mode and server mode.

Option 1: Colocate mode

In this mode, vLLM runs inside the trainer process and shares GPU memory with the training model. This avoids launching a separate server and can improve GPU utilization, but may lead to memory contention on the training GPUs. This is the default mode.

python

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,  # vllm_mode="colocate" by default
)

Option 2: Server mode

In this mode, vLLM runs in a separate process (and using separate GPUs) and communicates with the trainer via HTTP. This is ideal if you have dedicated GPUs for inference.

Start the vLLM server:

bash

VLLM_SERVER_DEV_MODE=1 vllm serve <model_name> \
    --weight-transfer-config '{"backend": "nccl"}' \
    --logprobs-mode processed_logprobs \
    --max-logprobs -1

Enable server mode in your training script:

python

from trl import DistillationConfig

training_args = DistillationConfig(
    ...,
    use_vllm=True,
    vllm_mode="server",
)

[!WARNING] Make sure that the server is using different GPUs than the trainer, otherwise you may run into NCCL errors. You can specify the GPUs to use with the CUDA_VISIBLE_DEVICES environment variable.

[!TIP] Depending on the model size and the overall GPU memory requirements for training, you may need to adjust the vllm_gpu_memory_utilization parameter in [DistillationConfig] to avoid underutilization or out-of-memory errors.

For more information, see Speeding up training with vLLM.

Train adapters with PEFT

We support tight integration with the 🤗 PEFT library, letting you train adapters and share them on the Hub rather than training the whole student.

python

from datasets import load_dataset
from trl import DistillationTrainer
from peft import LoraConfig

dataset = load_dataset("trl-lib/ultrafeedback-prompt", split="train")

trainer = DistillationTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    teacher_model="Qwen/Qwen2.5-1.5B-Instruct",
    train_dataset=dataset,
    peft_config=LoraConfig(),
)
trainer.train()

[!WARNING] The distillation loss reads lm_head.weight directly and runs the student backbone without going through PeftModel.forward(). Adapters on lm_head (via target_modules) and prompt-learning methods (PromptTuning, PrefixTuning, P-Tuning) are therefore rejected, since they would be silently ignored. To train the head, use modules_to_save=["lm_head"] instead.

Train with Liger Kernel

Liger Kernel is a collection of Triton kernels for LLM training that boosts multi-GPU throughput, cuts memory use, and works seamlessly with tools like FlashAttention, PyTorch FSDP, and DeepSpeed. For more information, see Liger Kernel Integration.

Set use_liger_kernel=True in the [DistillationConfig] to compute the JSD with the fused Liger kernel instead of the chunked path.

[!WARNING] The fused Liger kernel cannot apply per-model logit_scale (e.g. Cohere) or final_logit_softcapping (e.g. Gemma), so it is rejected for models that set them — use the default chunked path for those.

Agent Training

[DistillationTrainer] supports agent training: the student calls tools during generation and is distilled on the resulting trajectory. Tool-result tokens are masked out of the loss, so the student is only trained on the tokens it generated itself.

Tools

The tools argument expects a list of Python functions that define the tools available to the agent:

python

from trl import DistillationTrainer

trainer = DistillationTrainer(
    tools=[tool1, tool2],
    ...,
)

Each tool must be a standard Python function with type-hinted arguments and return types, along with a Google-style docstring describing its purpose, arguments, and return value. For more details, see the Passing tools guide.

[!TIP] The tool call loop requires the chat template to be prefix-preserving (appending a tool message must not change how earlier messages are rendered). For known model families (e.g. Qwen3, DeepSeek-V3), TRL automatically swaps in a patched training template when tools are enabled. See Chat Templates for the full list.

Use max_tool_calling_iterations in the [DistillationConfig] to cap the number of tool-calling turns. By default there is no limit, and generation stops when the student produces a response turn with no tool calls.

Example:

python

from trl import DistillationTrainer

def multiply(a: int, b: int) -> int:
    """
    Multiplies two integers.

    Args:
        a: The first integer.
        b: The second integer.

    Returns:
        The product of the two integers.
    """
    return a * b

trainer = DistillationTrainer(
    tools=[multiply],
    ...,
)

[!WARNING] Async tools are not supported yet — pass synchronous functions.

Multimodal Tool Responses

Tools can return images alongside text by returning a list of content blocks. This is useful for VLM agent training where the tool provides visual feedback (e.g., screenshots, plots, camera captures).

python

from PIL import Image

def take_screenshot() -> list:
    """
    Takes a screenshot of the current screen.

    Returns:
        The screenshot image with a description.
    """
    img = Image.open("screenshot.png")
    return [{"type": "image", "image": img}, {"type": "text", "text": "Here is the screenshot."}]

The returned images are automatically injected into the conversation and passed to the VLM for subsequent generation turns.

Training Vision Language Models

[DistillationTrainer] supports distilling Vision-Language Models (VLMs) on multimodal datasets containing both text and images. Pass a VLM as both the student and the teacher, and provide a prompt-only dataset with either an image column (single image per sample) or an images column (list of images per sample). For more information on the expected dataset structure, see the Dataset Format — Vision datasets section.

Tested with:

Gemma 3 — e.g., google/gemma-3-4b-it
LLaVA-NeXT — e.g., llava-hf/llava-v1.6-mistral-7b-hf
Qwen2-VL — e.g., Qwen/Qwen2-VL-2B-Instruct
Qwen2.5-VL — e.g., Qwen/Qwen2.5-VL-3B-Instruct

[!TIP] Compatibility with all VLMs is not guaranteed. If you believe a model should be supported, feel free to open an issue on GitHub — or better yet, submit a pull request with the required changes.

Command line interface

Use the trl distillation CLI to launch distillation training from the command line. It supports full training and LoRA via the standard ModelConfig flags.

bash

# Full training:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-5 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1

bash

# LoRA:
trl distillation \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
    --dataset_name trl-lib/ultrafeedback-prompt \
    --learning_rate 2e-4 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model \
    --num_train_epochs 1 \
    --use_peft \
    --lora_r 64 \
    --lora_alpha 16

DistillationTrainer

[[autodoc]] DistillationTrainer - train - save_model - push_to_hub

DistillationConfig

[[autodoc]] DistillationConfig