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.

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.
step_time: The average time (in seconds) taken per training step (including generation).
completions/mean_length: The average length of generated completions.
completions/min_length: The minimum length of generated completions.
completions/max_length: The maximum length of generated completions.
completions/mean_terminated_length: The average length of generated completions that terminate with EOS.
completions/min_terminated_length: The minimum length of generated completions that terminate with EOS.
completions/max_terminated_length: The maximum length of generated completions that terminate with EOS.
completions/clipped_ratio: The ratio of truncated (clipped) completions.
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
```
trl vllm-serve --model <model_name>
```

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.

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.

Example script

Use examples/scripts/distillation.py to launch distillation training from the command line. The script supports full training and LoRA via the standard ModelConfig flags.

bash

# Full training:
python examples/scripts/distillation.py \
    --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:
python examples/scripts/distillation.py \
    --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