Make a 7B Model Think Like a 70B: Knowledge Distillation with TRL's DistillationTrainer

TL;DR: Knowledge distillation lets you train a small, cheap model to think like a large expensive one — TRL's DistillationTrainer does it in under 20 lines and runs on a single GPU.

What you'll be able to do after this:

• Compress a large teacher LLM into a smaller, faster student using on-policy distillation
• Run TRL's DistillationTrainer from the command line or a Python script on any prompt-response dataset
• Choose between fully on-policy training (student generates, teacher scores — best quality) and off-policy (imitate teacher's dataset outputs — fastest)

What is knowledge distillation?

A 70B model knows a lot but costs a lot to serve. Knowledge distillation (KD) trains a smaller student model to match the teacher's behavior — not just the ground-truth answer, but the teacher's full probability distribution over tokens. The student learns the teacher's uncertainty, not just its final choice.

The classic problem: train the student on fixed teacher outputs, and it sees only completions the teacher would generate — it never learns to handle sequences it produces at inference time. Generalized KD (GKD) fixes this: let the student generate its own completions during training, then use the teacher to score them. The student learns from its own mistakes.

Famous example: DeepSeek R1-Distill-Qwen-32B captured ~85% of R1's reasoning ability at 1/20th the inference cost using this family of techniques.

TRL DistillationTrainer: 20 lines to a distilled model

TRL's DistillationTrainer (in trl.experimental.distillation) implements GKD with three key speedups:

1. Generation buffer — batches on-policy student generations across gradient accumulation steps into a single vLLM call. Up to 40× faster than naively alternating between generate and train steps.
2. Teacher server support — offloads the teacher to a separate vLLM server so it doesn't compete for GPU memory with the student. Needed for 70B+ teachers.
3. Binary-encoded logprob payloads — packs teacher log-probabilities into base64 NumPy arrays (~5× smaller than JSON).

Quickstart: Qwen2.5-7B teacher → Qwen2.5-1.5B student on GSM8K

from datasets import load_dataset
from trl.experimental.distillation import DistillationConfig, DistillationTrainer

# Prompt-only dataset (student generates its own answers)
dataset = load_dataset("openai/gsm8k", "main", split="train")
dataset = dataset.map(
    lambda x: {"messages": [{"role": "user", "content": x["question"]}]},
    remove_columns=dataset.column_names,
)

config = DistillationConfig(
    output_dir="results/distill-gsm8k",
    num_train_epochs=1,
    bf16=True,
    lmbda=1.0,   # fully on-policy: student generates, teacher scores
    beta=1.0,    # reverse KL — student learns to avoid its own bad modes
    teacher_model_init_kwargs={"dtype": "bfloat16"},
)

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

Or from the command line:

python examples/scripts/distillation.py \
    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
    --teacher_model_name_or_path Qwen/Qwen2.5-7B-Instruct \
    --dataset_name trl-lib/chatbot_arena_completions \
    --lmbda 1.0 --beta 1.0 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --output_dir distilled-model

Key parameters

Param	What it controls
lmbda=1.0	Fully on-policy (student generates all completions — highest quality)
lmbda=0.0	Fully off-policy (imitate teacher's dataset outputs only — fastest)
beta=1.0	Reverse KL loss (student avoids its own confabulation modes)
beta=0.0	Forward KL loss (match every mode of the teacher, including rare ones)

For teachers that don't fit on training GPUs (100B+): set use_teacher_server=True with teacher_model_server_url pointing to a vLLM server hosting the teacher, and loss_top_k=1.

Sources: TRL DistillationTrainer docs · TRL distillation example script · Knowledge Distillation guide (2026)