Distilling a Large Model into a Small One: A PyTorch Playbook

Every team we work with eventually hits the same wall: the model that scores best in evaluation is too expensive to serve. A 70B teacher answers your support tickets beautifully at $0.60 a conversation and 4 seconds to first token. Finance wants $0.02. Product wants 400ms. Nobody wants the accuracy drop that comes from simply swapping in a smaller checkpoint.

Distillation is how you close that gap. Instead of hoping an off-the-shelf small model is good enough, you train it on the behaviour of your large one, on your own task distribution. Done properly it routinely recovers 90-98% of teacher quality on a narrow task at a tenth of the cost. Done carelessly it produces a small model that is confidently wrong in exactly the places the teacher was uncertain.

This is the playbook we use on client engagements, in PyTorch, end to end.

Pick the right kind of distillation first

"Distillation" covers three quite different techniques. Choosing the wrong one wastes a month.

Approach	What the student learns from	Needs teacher weights?	Best for
Response (hard-label) distillation	Teacher's generated outputs, treated as ground truth	No — API access is enough	LLM task specialisation, when the teacher is a closed model
Logit (soft-label) distillation	Full output distribution, via KL divergence	Yes	Classifiers, rankers, encoders, same-tokenizer LLM pairs
Feature / hidden-state distillation	Intermediate activations or attention maps	Yes, plus architectural alignment	Compressing an encoder into a shallower one of the same family

Response distillation is the workhorse for generative tasks and the only option when the teacher is behind an API. Logit distillation is strictly more informative — the soft distribution carries the teacher's uncertainty, the "dark knowledge" that makes distillation better than plain supervised fine-tuning — but it requires the same tokenizer and vocabulary. Feature distillation gives the biggest compression on encoders and the biggest headache in engineering.

A fourth option is worth naming so you can rule it out: pruning plus continued pretraining. It keeps the teacher's weights and cuts layers or width. It is a good complement to distillation (prune, then distil the pruned model back up) and a poor substitute for it.

Step 1: build the transfer set, not the training set

The most common failure we see is a distillation dataset that looks nothing like production traffic. The student only learns the teacher's behaviour on inputs it actually sees.

Rank your sources in this order:

1. Real production inputs. Logged, deduplicated, PII-scrubbed. This is worth ten times any synthetic corpus.
2. Teacher-generated variations of real inputs — paraphrases, edge-case mutations, harder versions.
3. Fully synthetic prompts to cover intents you know exist but have little traffic for.

Aim for coverage rather than volume. Ten thousand well-spread examples beat a million clustered around three intents. Cluster your embeddings and check the histogram before you spend GPU hours; we have seen 400k-row transfer sets where 80% of rows were variations of "where is my order".

For response distillation, generate with the teacher offline, in bulk, with vLLM if the teacher is open-weights:

from vllm import LLM, SamplingParams

teacher = LLM(model="meta-llama/Llama-3.3-70B-Instruct",
              tensor_parallel_size=4, max_model_len=8192)

params = SamplingParams(temperature=0.7, top_p=0.95, n=4, max_tokens=1024)
outputs = teacher.generate(prompts, params)

Note n=4. Sampling several completions per prompt and keeping the best — by a verifier, a rubric grader or majority vote on the final answer — is rejection sampling distillation, and it is the single highest-leverage trick in the response-distillation toolbox. You are not teaching the student the teacher's average behaviour; you are teaching it the teacher's good behaviour.

Filter hard. Drop completions that fail a schema check, contradict a retrieved source, or that your grader scores below threshold. A 30% rejection rate is normal and healthy.

Step 2: pick a student that can actually be taught

Capability ceilings are real. A 1B student will not learn multi-step reasoning from a 70B teacher no matter how good your data is. Rough guidance from our engagements:

• Extraction, classification, routing, structured output: 0.5B–3B students recover teacher quality almost fully.
• Summarisation, RAG answering, style-constrained generation: 3B–8B.
• Multi-hop reasoning, code generation, agentic tool use: 8B–14B, and expect a real gap.

Also prefer a student from a family with the same tokenizer as the teacher when you can — it keeps logit distillation on the table and avoids a class of subtle detokenization bugs.

Step 3: the loss

For classifiers and same-vocabulary LLM pairs, the classic combined objective still works. Soft targets carry most of the signal; temperature spreads the teacher's distribution so the student sees relative probabilities among wrong answers, and the T**2 factor keeps gradient magnitudes comparable as T changes.

import torch
import torch.nn.functional as F

def distill_loss(student_logits, teacher_logits, labels, T=2.0, alpha=0.9):
    soft = F.kl_div(
        F.log_softmax(student_logits / T, dim=-1),
        F.log_softmax(teacher_logits / T, dim=-1),
        reduction="batchmean",
        log_target=True,
    ) * (T * T)
    hard = F.cross_entropy(student_logits, labels)
    return alpha * soft + (1.0 - alpha) * hard

Two practical notes. Use log_target=True and pass log-probabilities on both sides — computing softmax then log separately is a well-worn source of NaNs. And alpha near 0.9 is a sane default: the teacher distribution is a richer supervisory signal than the one-hot label, so weight it accordingly.

For sequence models, apply the same loss token-wise over the response span only, masking the prompt:

mask = (labels != -100)
kl = F.kl_div(s_logprobs, t_logprobs, reduction="none", log_target=True).sum(-1)
loss = (kl * mask).sum() / mask.sum() * (T * T)

If the vocabularies differ, you are back to hard-label training and the loss is plain cross-entropy on the teacher's text. That is fine — most production LLM distillation is exactly this — but it is why the transfer set quality matters so much more in that setting.

Precompute or co-run the teacher?

Precomputing teacher logits for a 128k-vocabulary model at 8k context costs about 2 GB per sequence in BF16. Do not store full distributions. Two workable options:

• Top-k logits. Store the top 64–128 logits and indices per token, renormalise, and treat the tail as uniform. Roughly 99% storage reduction with negligible quality loss.
• Teacher in the loop. Keep the frozen teacher on separate GPUs under torch.no_grad() and torch.autocast("cuda", dtype=torch.bfloat16). Simpler, but your training job now needs teacher-sized memory and every epoch pays the forward cost again.

We default to top-k on disk for anything running more than two epochs.

Step 4: train it

Nothing exotic. Full fine-tuning of the student beats LoRA here — you are trying to move the student's behaviour substantially, not nudge it — but LoRA is a reasonable first pass when you want a same-day signal. If you have not set up a multi-GPU student run before, our FSDP2 guide covers the sharding side, and QLoRA on a single GPU covers the cheap end.

Details that matter more than they should:

• Learning rate: 1e-5 to 2e-5 for full fine-tuning of a 3B–8B student, cosine decay, 3% warmup. Distillation tolerates a slightly higher LR than SFT because the targets are smoother.
• Epochs: 2–3 on a well-filtered transfer set. Past that the student memorises teacher quirks.
• Sequence packing: on. Wasted padding is 20–40% of your compute on short-answer tasks.
• BF16 autocast, fused AdamW, torch.compile on the student. See torch.compile in practice; a compiled student plus a compiled teacher forward is usually a 1.3–1.6x end-to-end win on this workload.
• Keep a frozen eval slice from real traffic that never enters the transfer set.

Step 5: evaluate against the teacher, not against a benchmark

The question is never "is the student good", it is "where does the student diverge from the teacher, and does it matter". Build a comparison harness — the same discipline as our LLM evaluation harness post — and report:

1. Agreement rate with the teacher on held-out real inputs (exact match for structured tasks, graded preference for generative ones).
2. Per-segment agreement. Slice by intent, language, input length, customer tier. Distillation failures are almost never uniform; they concentrate in the long tail.
3. Calibration. Distilled students are often overconfident. If you route low-confidence cases to the teacher, a miscalibrated student silently breaks the router.
4. Cost and latency, measured, not estimated: tokens/sec at your real batch size on your real hardware.

A useful acceptance shape: "student handles 92% of traffic at ≥98% teacher agreement; the remaining 8% escalate to the teacher." That hybrid is usually where the economics land, and it degrades gracefully.

Step 6: quantize afterwards, and re-measure

Distillation and quantization compose. Distil first, then quantize the student — an 8B student at INT8 or FP8 is a very different cost line from a 70B teacher at BF16. Our precision guide covers the format choice; the only distillation-specific warning is that a distilled student sits closer to its decision boundaries than a generically fine-tuned one, so post-quantization accuracy checks are mandatory rather than a formality. Re-run the agreement harness after quantizing, every time.

Then serve it properly: vLLM for generative students, or torch.export and AOTInductor for encoders and rankers where Python-free, fixed-shape serving wins.

Failure modes we keep meeting

Distilling a teacher that was never good enough. If the teacher is 82% accurate on your task, the student will be worse. Fix the teacher — better prompting, retrieval, or fine-tuning — before you compress it.

Transfer set leakage into eval. Teacher-generated paraphrases of eval inputs land in training and your agreement numbers become fiction. Split by source input before generating.

Ignoring the refusal and safety surface. Students inherit capability much faster than they inherit refusals. Include adversarial and out-of-scope inputs in the transfer set explicitly, and evaluate them separately.

One-shot distillation with no refresh loop. Traffic drifts, the teacher gets upgraded, and the frozen student quietly diverges. Budget a quarterly re-distillation with fresh logged inputs; it is a few GPU-hours, not a project.

Compressing before profiling. Sometimes the expensive thing is not the model. If p99 is dominated by retrieval or a serialisation hop, a smaller model saves you nothing — profile first.

A sensible first project

Take one high-volume, narrow task. Log 20k real inputs. Generate four teacher completions each, filter to the best one, hold out 2k inputs as a frozen eval slice. Fine-tune a 3B student for two epochs. Measure agreement per segment. If it clears 95% overall, ship it behind a confidence-based router with teacher fallback and watch it for two weeks.

That is a two-to-three week engagement, not a research programme, and it usually pays for itself in the first month of inference bills.

We do this work with client teams end to end — transfer set design, rejection-sampling pipelines, student selection, the training run, the agreement harness and the serving swap. See our LLM fine-tuning and LLM inference and serving optimization services, or contact us with your teacher model, your traffic volume and your target cost per request, and we will tell you whether distillation is the right lever.