Lab: Logit Distillation with TRL's Distillation Trainers

Validated on: written from the documentation cited above; not yet validated on hardware on any track. The per-track versions, memory peaks and wall-clock times will be recorded here when the validation pass is done.

Objective

By the end of this lab you will have trained a student against a teacher’s probability distribution rather than against its text, and you will be able to say whether that was worth the memory it cost.

Concretely: a tokeniser compatibility report for your pair, a memory budget written down before the run, a training run with TRL’s DistillationTrainer, a divergence loss you can read and explain, and a comparison against the sequence-level student from the previous lab on the same tasks, with what each route cost taken from the run log.

The result may well be that the sequence-level student is as good and cost a fifth as much. That is a legitimate outcome of this lab and a useful thing to know about your own hardware.

⌂Single-machine path

Single-machine, and unlike the previous lab there is no way to split it: the trainer loads the teacher into the same process as the student. TRL does ship an asynchronous trainer that scores against a teacher served over HTTP, which would allow a split, but its documentation states that the teacher server, the student’s vLLM server and the trainer must run on separate GPUs, which is outside every track in this course.

Requirements

The Part 11 training environment with TRL 1.12.0 · verified 2026-09-08, transformers 5.16.1 · verified 2026-09-08 and PEFT 0.20.0 · verified 2026-09-08, the seed prompt file from the previous lab, your Part 10 task file, and 60 minutes of which about 30 are unattended.

This is the one page in this part with a real memory floor, because two models are resident.

Memory	Teacher	Student	Notes
32 GB and above	Qwen3-8B	Qwen3-1.7B	Comfortable; the headline pair
24 GB	Qwen3-8B	Qwen3-1.7B	Under a gigabyte of headroom; see the budget below
12 to 16 GB	Qwen3-4B	Qwen3-0.6B	The reduced path; not the primary path this page is written for

Both pairs are same-family and Apache-2.0 according to their model cards, read on 2026-09-09.

Track S — NVIDIA DGX Spark

The primary track for this lab, along with N. The 128 GB pool makes the headline pair comfortable and leaves room to raise --max-completion-length and see what it costs. Work inside the NGC PyTorch container from Part 11 with your course directory mounted.

TRL’s distillation trainer can use vLLM to generate the student’s on-policy completions, which is the slow part of every step. That is available on this track and is the first thing to try if the run is slower than you can stand; it is not required, and the lab is written without it so that the same commands work on all four tracks.

Track X — AMD Ryzen AI Max+ 395Partial

This lab trains rather than serves, so it needs the ROCm build of PyTorch. The ROCm 10.0.0 compatibility matrix dated 2026-08-14 lists gfx1151 without a support-tier qualifier and AMD's PyTorch install page read on 2026-09-09 does not mention the chip. On the CPU, an on-policy trainer that generates at every step is impractically slow rather than merely slow.

Where the ROCm wheels work, the 128 GB pool makes this lab comfortable and everything below applies unchanged. Confirm with python -c 'import torch; print(torch.cuda.is_available())' first.

Where they do not, this is the one page in the part with no honest CPU path. The previous lab’s CPU route works because supervised fine-tuning does no generation; this trainer generates a completion at every step, and on a CPU that turns a 30-minute run into an overnight one. Read the lab, run the tokeniser check in task 1, which needs no accelerator, and take the comparison in task 6 from the sequence-level student you already have.

Track M — Apple siliconPartial

mlx-lm has no distillation trainer at the time of writing, so this runs under PyTorch's MPS backend in float32, which doubles the memory of both models and rules out the headline pair.

Use the reduced pair, Qwen3-4B teacher and Qwen3-0.6B student, and expect float32 rather than bfloat16. In float32 those two models are about 16 GB and 2.4 GB of weights respectively, so a 128 GB Mac has room and a 16 GB Mac does not.

The script detects the MPS backend and prints the pair it recommends. Everything else on this page, including the loss discussion and the comparison, applies unchanged.

Track N — NVIDIA desktop or laptop

The tier table above is written for this track. At 24 GB the headline pair fits with under a gigabyte to spare, so close everything else on the machine, and if the card is also driving your displays, use the reduced pair instead. There is no shame in it: the reduced pair demonstrates every idea on this page and finishes sooner.

On Windows, work inside WSL2, and remember that the virtual machine’s memory limit is what the training process sees.

Tasks

Working directory and terminal roles

Prepare the course execution workspace once before this procedure. It includes this part's scripts, data and shared Python helpers. In the client or training terminal, select this directory:

RunnableAll tracks

select this part’s execution directory

export LABS_ROOT="${LABS_ROOT:-$HOME/llm-course/labs}"
export LAB_DIR="$LABS_ROOT/part-15-distillation"
cd "$LAB_DIR"
pwd
test -f "train-logit-distil.py"

Expected result: pwd ends in part-15-distillation and the file check returns successfully. If it does not, finish workspace preparation before continuing. Activate the environment in the requirements for your track. Bare script and data filenames below are relative to this directory; paths to earlier experiments must point at the artefacts you actually retained.

Keep each foreground server in a separate terminal and send requests from this terminal. Reapply lesson-specific environment variables in each new shell. Stop at the first failed checkpoint and retain its output; the execution guide explains how to distinguish missing files, endpoint failures and capacity problems.

1. Check the tokenisers before anything else

This is the check that distinguishes a lab from a wasted evening. Logit distillation applies the teacher’s probability for token id n to the student’s vocabulary at index n. If the two vocabularies differ, the loss is well defined, the training runs, the loss falls, and the student learns the wrong thing. TRL’s own documentation warns that a teacher with a different vocabulary trains against the wrong tokens, silently when its vocabulary is no larger than the student’s.

RunnableAll tracks

train-logit-distil.py

"""Distil a student against a teacher's next-token distribution with TRL.

Purpose: the logit lab's training run. Holds a teacher and a student in memory at
    once and trains the student to match the teacher's distribution rather than its
    text. Two trainers are selectable: TRL's DistillationTrainer, which release
    1.12.0's documentation lists in the stable API and which is on-policy (the
    student generates the completions it is trained on), and the experimental
    GKDTrainer, whose lmbda setting lets you dial the same run from fully
    off-policy to fully on-policy so the two ends can be compared.
Platform: spark, nvidia (CUDA) primary; strix on ROCm, dated on the page; mac under
    PyTorch MPS in float32 with the small pair only. Two models are resident, so
    this is the one page in the part with a real memory floor.
Minimum memory: 24 GB for the 8B-to-1.7B pair; 12 to 16 GB for the 4B-to-0.6B pair
Assumes: torch, transformers, trl and peft installed in the active environment; a
    prompt file in JSON Lines with a "prompt" field on every line, as written by
    make-seed-prompts.py or derived from the sequence-level lab's training data;
    distillog.py sits next to this file. The teacher and the student must share a
    tokeniser: this script checks and refuses to run if they do not.

Usage: python3 train-logit-distil.py --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --prompts seeds/prompts.jsonl --output-dir runs/logit-qwen3-1.7b \\
           --labbook labbook.md
       python3 train-logit-distil.py --teacher Qwen/Qwen3-4B --student Qwen/Qwen3-0.6B \\
           --prompts seeds/prompts.jsonl --output-dir runs/logit-qwen3-0.6b \\
           --max-completion-length 256 --batch-size 1 --grad-accum 8
       python3 train-logit-distil.py --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --prompts seeds/prompts.jsonl --trainer gkd --lmbda 0.0 --beta 0.0 \\
           --output-dir runs/gkd-offpolicy
       python3 train-logit-distil.py --teacher Qwen/Qwen3-8B --student Qwen/Qwen3-1.7B \\
           --check-tokenisers-only
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer

import distillog

DEFAULT_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

# The strings the vocabulary check compares. Anything with special tokens, digits,
# whitespace and a non-Latin script in it; two tokenisers that agree on all of
# these agree on everything this lab does.
PROBE_STRINGS = (
    "<|im_start|>user\nHello, world!<|im_end|>\n",
    "<think>\n2 + 2 = 4\n</think>\nAnswer: 4",
    "def nth_word(text, n):\n    return text.split()[n - 1]\n",
    "quantisation, tokeniser, générateur, 分词器",
)


def compare_tokenisers(teacher_id: str, student_id: str) -> dict:
    """Encode the probes with both tokenisers and report every disagreement.

    The distillation loss reads the teacher's probability for a token id and
    applies it to the student's vocabulary at the same index. If the two
    vocabularies differ, the loss is well defined, runs without error, and trains
    the student towards the wrong tokens. TRL's own documentation warns that a
    teacher with a different vocabulary trains against the wrong tokens, silently
    when its vocabulary is no larger than the student's, so this check runs first.
    """
    teacher_tok = AutoTokenizer.from_pretrained(teacher_id)
    student_tok = AutoTokenizer.from_pretrained(student_id)
    report = {
        "teacher": teacher_id,
        "student": student_id,
        "teacher_class": type(teacher_tok).__name__,
        "student_class": type(student_tok).__name__,
        "teacher_vocab_size": len(teacher_tok),
        "student_vocab_size": len(student_tok),
        "mismatched_probes": [],
        "special_token_differences": [],
    }
    for probe in PROBE_STRINGS:
        a = teacher_tok(probe, add_special_tokens=False)["input_ids"]
        b = student_tok(probe, add_special_tokens=False)["input_ids"]
        if a != b:
            report["mismatched_probes"].append({
                "probe": probe[:60], "teacher_ids": a[:20], "student_ids": b[:20],
            })
    teacher_specials = dict(sorted(teacher_tok.get_added_vocab().items()))
    student_specials = dict(sorted(student_tok.get_added_vocab().items()))
    for token, tid in teacher_specials.items():
        if student_specials.get(token) != tid:
            report["special_token_differences"].append(
                {"token": token, "teacher_id": tid, "student_id": student_specials.get(token)}
            )
    report["compatible"] = (
        report["teacher_vocab_size"] == report["student_vocab_size"]
        and not report["mismatched_probes"]
        and not report["special_token_differences"]
    )
    return report


def print_tokeniser_report(report: dict) -> None:
    print(f"teacher tokeniser: {report['teacher_class']}, {report['teacher_vocab_size']} entries")
    print(f"student tokeniser: {report['student_class']}, {report['student_vocab_size']} entries")
    if report["compatible"]:
        print("compatible: the same ids for every probe, and the same special-token ids")
        return
    print("NOT compatible:")
    if report["teacher_vocab_size"] != report["student_vocab_size"]:
        print("  vocabulary sizes differ")
    for item in report["mismatched_probes"]:
        print(f"  probe {item['probe']!r}")
        print(f"    teacher {item['teacher_ids']}")
        print(f"    student {item['student_ids']}")
    for item in report["special_token_differences"][:10]:
        print(f"  special token {item['token']}: teacher {item['teacher_id']}, student {item['student_id']}")


def pick_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return "mps"
    return "cpu"


def as_prompt_only(path: str) -> str:
    """Accept either a prompt file or the sequence-level lab's training file.

    DistillationTrainer wants a prompt-only dataset, because the student writes
    the completions itself. Passing it a file that also carries the teacher's
    completions would be silently ignoring the column that took the longest to
    make, so the script says which shape it found.
    """
    first = None
    with open(path, encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                first = json.loads(line)
                break
    if first is None:
        raise SystemExit(f"{path} is empty")
    if "prompt" not in first:
        raise SystemExit(f"{path} has no 'prompt' field; make-seed-prompts.py writes one")
    if "completion" in first:
        print(f"note: {path} also has a 'completion' column. This trainer generates the "
              "student's own completions, so that column is not used here. It is what "
              "train-student.py trained on, which is exactly the comparison this lab makes.")
    return path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--teacher", default="Qwen/Qwen3-8B")
    parser.add_argument("--student", default="Qwen/Qwen3-1.7B")
    parser.add_argument("--prompts", default="seeds/prompts.jsonl")
    parser.add_argument("--output-dir", default="runs/logit-student")
    parser.add_argument("--trainer", choices=["distillation", "gkd"], default="distillation",
                        help="distillation is the stable trainer; gkd is experimental")
    parser.add_argument("--beta", type=float, default=None,
                        help="generalised JSD interpolation. 0.0 is forward KL, 1.0 reverse KL. "
                             "Defaults to the trainer's own default when omitted.")
    parser.add_argument("--lmbda", type=float, default=0.5,
                        help="GKD only: fraction of on-policy student data. 0.0 is fully off-policy.")
    parser.add_argument("--seq-kd", action="store_true",
                        help="GKD only: train on teacher-generated output, i.e. sequence-level KD")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--epochs", type=float, default=1.0)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=8)
    parser.add_argument("--lr", type=float, default=1e-4, help="adapter rate; the trainer's own default is 1e-6 for a full student")
    parser.add_argument("--max-completion-length", type=int, default=256)
    parser.add_argument("--rank", type=int, default=16)
    parser.add_argument("--alpha", type=int, default=32)
    parser.add_argument("--target-modules", nargs="+", default=DEFAULT_TARGETS)
    parser.add_argument("--full-student", action="store_true",
                        help="train every student weight instead of a LoRA adapter; needs far more memory")
    parser.add_argument("--limit", type=int, default=None, help="use only this many prompts")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--check-tokenisers-only", action="store_true",
                        help="run the vocabulary check and exit")
    parser.add_argument("--allow-tokeniser-mismatch", action="store_true",
                        help="run anyway. Only for reproducing the fault in this part's challenge.")
    args = parser.parse_args()

    report = compare_tokenisers(args.teacher, args.student)
    print_tokeniser_report(report)
    if args.check_tokenisers_only:
        return
    if not report["compatible"] and not args.allow_tokeniser_mismatch:
        raise SystemExit(
            "teacher and student do not share a tokeniser. Logit distillation matches "
            "probabilities index by index, so this run would train the student towards "
            "the wrong tokens without failing. Pick a same-family pair, or use "
            "sequence-level distillation, which needs only text. Pass "
            "--allow-tokeniser-mismatch to reproduce the fault deliberately."
        )

    device = pick_device()
    bf16 = device == "cuda" and torch.cuda.is_bf16_supported()
    print(f"device: {device}   precision: {'bfloat16' if bf16 else 'float32'}")
    if device == "mps":
        print("note: on Track M this runs in float32 under PyTorch's MPS backend. Use the "
              "4B-to-0.6B pair; the 8B teacher will not fit beside a student in float32.")

    prompts_path = as_prompt_only(args.prompts)
    dataset = load_dataset("json", data_files={"train": prompts_path})["train"]
    if args.limit:
        dataset = dataset.select(range(min(args.limit, len(dataset))))
    dataset = dataset.map(
        lambda row: {"prompt": [{"role": "user", "content": row["prompt"]}]},
        remove_columns=[c for c in dataset.column_names if c != "prompt"],
    )
    print(f"prompts: {len(dataset)}")

    peft_config = None if args.full_student else LoraConfig(
        r=args.rank, lora_alpha=args.alpha, lora_dropout=0.05,
        target_modules=args.target_modules, bias="none", task_type="CAUSAL_LM",
    )

    common = dict(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        temperature=args.temperature,
        bf16=bf16,
        logging_steps=5,
        report_to="none",
        seed=args.seed,
        save_strategy="epoch",
        save_total_limit=1,
    )

    if args.trainer == "distillation":
        from trl import DistillationConfig, DistillationTrainer

        config = DistillationConfig(
            **common,
            max_completion_length=args.max_completion_length,
            **({"beta": args.beta} if args.beta is not None else {}),
        )
        trainer = DistillationTrainer(
            model=args.student,
            teacher_model=args.teacher,
            args=config,
            train_dataset=dataset,
            peft_config=peft_config,
        )
        method = "trl.DistillationTrainer (on-policy, generalised JSD)"
    else:
        from trl.experimental.gkd import GKDConfig, GKDTrainer

        config = GKDConfig(
            **common,
            max_new_tokens=args.max_completion_length,
            lmbda=args.lmbda,
            seq_kd=args.seq_kd,
            **({"beta": args.beta} if args.beta is not None else {}),
        )
        trainer = GKDTrainer(
            model=args.student,
            teacher_model=args.teacher,
            args=config,
            train_dataset=dataset,
            peft_config=peft_config,
        )
        method = f"trl.experimental.gkd.GKDTrainer (lmbda={args.lmbda}, seq_kd={args.seq_kd})"

    print(f"trainer: {method}")
    if peft_config is not None and hasattr(trainer.model, "print_trainable_parameters"):
        trainer.model.print_trainable_parameters()

    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()
    started = time.time()
    watts_start = distillog.sample_power()
    trainer.train()
    elapsed = time.time() - started
    watts_end = distillog.sample_power()
    peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) if device == "cuda" else None

    trainer.save_model(args.output_dir)
    AutoTokenizer.from_pretrained(args.student).save_pretrained(args.output_dir)

    history = trainer.state.log_history
    losses = [row["loss"] for row in history if "loss" in row]
    summary = {
        "first_loss": round(losses[0], 4) if losses else None,
        "final_loss": round(losses[-1], 4) if losses else None,
        "steps": len(losses),
        "peak_cuda_memory_gb": peak_gb,
    }
    print(json.dumps(summary, indent=2))
    print(f"student saved to {args.output_dir}")

    if args.labbook:
        readings = [w for w in (watts_start, watts_end) if w is not None]
        rec = distillog.record(
            labbook=args.labbook,
            lab="part-15/train-logit-distil",
            stage="train-logit",
            teacher={"id": args.teacher, "resident": True,
                     "tokeniser_compatible": report["compatible"]},
            student={"id": args.student, "adapter": args.output_dir,
                     "method": "lora" if peft_config else "full"},
            dataset={"prompts": prompts_path, "sha256": distillog.file_sha256(prompts_path),
                     "used": len(dataset)},
            hyperparameters={
                "trainer": method, "beta": args.beta, "lmbda": args.lmbda if args.trainer == "gkd" else None,
                "seq_kd": args.seq_kd if args.trainer == "gkd" else None,
                "temperature": args.temperature,
                "max_completion_length": args.max_completion_length,
                "epochs": args.epochs, "batch_size": args.batch_size,
                "grad_accum": args.grad_accum, "learning_rate": args.lr,
                "precision": "bfloat16" if bf16 else "float32",
                "rank": args.rank if peft_config else None,
            },
            seed=args.seed,
            cost=distillog.build_cost(
                seconds=elapsed,
                mean_watts=(sum(readings) / len(readings)) if readings else None,
            ),
            losses=summary,
            config_path=__file__,
            notes=None,
        )
        print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()

Download train-logit-distil.py337 lines

RunnableAll tracks

compare two vocabularies and stop

python train-logit-distil.py \
    --teacher Qwen/Qwen3-8B \
    --student Qwen/Qwen3-1.7B \
    --check-tokenisers-only

Output — what you should see

teacher tokeniser: Qwen2Tokenizer, 151936 entries
student tokeniser: Qwen2Tokenizer, 151936 entries
compatible: the same ids for every probe, and the same special-token ids

The check encodes four probe strings containing chat markers, thinking tags, code, digits and non-Latin script, and compares the id sequences, then compares every added special token by id. Both Qwen3 sizes declare the Qwen2Tokenizer class and the same ids for <|im_start|>, <|im_end|>, <think> and </think>, which their tokenizer_config.json files show directly.

▷Exercise

Run the same check with a deliberately mismatched pair before continuing, so you know what a failure looks like: --teacher Qwen/Qwen3-8B --student meta-llama/Llama-3.1-8B-Instruct --check-tokenisers-only. Note that this is a check, not a training run, so it downloads only the tokenisers. The script refuses to train a mismatched pair unless you pass --allow-tokeniser-mismatch, which exists so that this part’s challenge can reproduce the fault on purpose.

2. Budget the memory for two models

The headline pair on a 24 GB machine: an 8B teacher and a 1.7B student with a rank-16 adapter

Teacher weights, BF16, frozen

16.4 GB

Student weights, BF16, frozen

3.4 GB

Adapter, gradients and optimiser states

0.3 GB

Activations and chunked divergence

1 GB

Student generation cache

0.1 GB

Reserved for the operating system

2 GB

Free

0.8 GB

Total

24 GB

Estimate from arithmetic, not a measurement. Weights are the BF16 sizes in src/data/models.json. The adapter figure is 17.4 million trainable parameters at 16 bytes each, computed from Qwen3-1.7B's config.json: 28 layers, hidden size 2,048, intermediate size 6,144, rank 16 on the seven projections. Activations assume batch 1 at a 256-token completion with TRL's chunked divergence, which its documentation says keeps peak activation memory off the full vocabulary times sequence-length logits tensor. The generation cache is the student's own, at its published 114,688 bytes per token.

Under a gigabyte free on a 24 GB machine. That is the honest picture, and it is why the tier table says 32 GB for comfort. A browser, a desktop compositor or a second Python process will take that gigabyte.

The reduced pair changes the first two bars to 8.0 GB and 1.2 GB and the adapter to 10.1 million parameters, about 0.16 GB, which is the same figure Part 11’s first training run reported for Qwen3-0.6B. That budget fits a 16 GB machine with room to work in.

▤Memory budget

The teacher’s weights are the largest bar and there is no adapter trick that shrinks them: it is frozen, in evaluation mode, and still fully resident because every step needs its forward pass. Quantising the teacher would change the distribution the student is being trained to match, which is a different experiment rather than a memory optimisation.

3. Understand what the trainer will do before you run it

TRL TRL 1.12.0 · verified 2026-09-08 lists four trainers under knowledge distillation, and its release note says the DistillationTrainer “graduates to the stable API — on-policy knowledge distillation that matches a teacher’s full next-token distribution with a memory-efficient chunked JSD loss and vLLM-powered generation”. Read on 2026-09-09, GKDTrainer, MiniLLMTrainer and AsyncDistillationTrainer are marked experimental and live under trl.experimental.

Three things about the stable trainer surprise people, and all three are documented.

It is on-policy. The dataset it takes is prompt-only: “the student generates its own completions on-policy, so only the prompt is needed”. You are not feeding it the teacher answers from the previous lab. It generates the student’s own completions and asks the teacher what it would have said at each position.

The default divergence is reverse KL. beta defaults to 1.0. The documentation is explicit that unlike GRPO’s beta, which penalises divergence from a reference model, “here it selects the divergence itself; there is no reference-model KL penalty”. Zero is forward KL, one is reverse KL, and a half is the symmetric Jensen-Shannon divergence.

The default learning rate is 1e-6. That is a full-model rate. This lab trains an adapter, and the script’s default is 1e-4, for the reasons Part 13’s LoRA lesson gives.

Pseudocode — not a real command

for each step:
    student generates a completion for each prompt in the batch
    teacher and student each produce a distribution at every completion position
    loss = generalised Jensen-Shannon divergence between them, interpolated by beta
    backward pass updates the adapter only

4. Train

RunnableAll tracks

logit distillation with the stable trainer

python train-logit-distil.py \
    --teacher Qwen/Qwen3-8B \
    --student Qwen/Qwen3-1.7B \
    --prompts seeds/prompts.jsonl \
    --output-dir runs/logit-qwen3-1.7b \
    --max-completion-length 256 \
    --batch-size 1 \
    --grad-accum 8 \
    --epochs 1 \
    --labbook labbook.md

On the reduced pair, substitute --teacher Qwen/Qwen3-4B --student Qwen/Qwen3-0.6B and --output-dir runs/logit-qwen3-0.6b.

Output — what you should see

teacher tokeniser: Qwen2Tokenizer, 151936 entries
student tokeniser: Qwen2Tokenizer, 151936 entries
compatible: the same ids for every probe, and the same special-token ids
device: cuda   precision: bfloat16
prompts: 600
trainer: trl.DistillationTrainer (on-policy, generalised JSD)
trainable params: 17,432,576 || all params: 1,7xx,xxx,xxx || trainable%: 1.0xxx

The script reuses the seed prompts from the previous lab and says so if the file also has a completion column: that column is what the sequence-level student trained on and this trainer does not use it, which is precisely the comparison the two labs set up.

=Expected result

Expect this to be several times slower per example than the previous lab’s supervised run, because every step generates. That is the trade the first lesson’s timeline diagram describes, and watching it is the point of running the lab.

5. Read the loss, which is not a cross entropy

The number falling in your terminal is a divergence between two distributions, and it reads differently from the cross entropy of a supervised run.

• It has a floor that is not zero for any interesting pair. The student cannot represent the teacher’s distribution exactly; it has fewer parameters. A divergence that flattens well above zero is the expected shape, not a stalled run.
• It is not comparable with a supervised loss. Do not put the sequence-level run’s final loss and this run’s final loss in the same column. They measure different things.
• It is comparable with itself across runs of the same pair at the same beta and temperature. Change either and the numbers move for reasons that have nothing to do with the student.

TRL documents the metrics the trainer logs, and two are worth watching beside the loss. completions/mean_length tells you how long the student’s own generations are; a collapse towards the minimum usually means the student has learned to stop early rather than to answer well. entropy is “average entropy of token predictions across generated completions”; a falling divergence with a collapsing entropy is a student narrowing rather than learning, which is the characteristic risk of the reverse-KL direction.

✕Common mistake

Reading a flat divergence as a broken run and raising the learning rate. The floor is real and it is set by the capacity gap. Raise the rate and you get a student that has narrowed onto one of the teacher’s modes faster, which shows as a lower loss and a less useful model. Judge this run on the evaluation in task 6, not on the loss.

6. Compare against the sequence-level student

Export the logit student the same way as the previous lab, serve it under a third alias, and run the comparison.

RunnableAll tracks

merge, convert and quantise the logit student

python ~/llm-course/merge-adapter.py \
    --adapter runs/logit-qwen3-1.7b \
    --merged-dir models/logit-qwen3-1.7b-merged
python ~/llama.cpp/convert_hf_to_gguf.py models/logit-qwen3-1.7b-merged \
    --outfile models/logit-qwen3-1.7b.gguf \
    --outtype bf16
~/llama.cpp/build/bin/llama-quantize \
    models/logit-qwen3-1.7b.gguf \
    models/logit-qwen3-1.7b-Q4_K_M.gguf \
    Q4_K_M

RunnableAll tracks

compare-students.py

"""Put the sequence-level student and the logit student side by side, with what each cost.

Purpose: the logit lab's conclusion. Runs Part 10's harness over the same tasks for
    both students, asks the judge to compare them with the order swapped so position
    bias shows as a flip rate rather than hiding inside the answer, and reads the lab
    notebook for what each route cost in tokens, seconds and watt-hours. The output
    is a decision, not a score: which of the two routes to use next time, and on what
    evidence.
Platform: all (pure Python over HTTP; both students are reached through an
    OpenAI-compatible API, so they may be served by any engine on any track)
Minimum memory: 12 GB on the machine serving the students; this script needs very little
Assumes: Python 3.10 or newer; Part 10's run-eval.py and judge.py in --harness-dir;
    both students reachable at --base-url under the names given; distillog.py next to
    this file; the lab notebook holding the generate, filter, train and train-logit
    records written by the other scripts in this part.

Usage: python3 compare-students.py --harness-dir ~/eval --tasks my-tasks.json \\
           --base-url http://127.0.0.1:4000/v1 \\
           --sequence-student local/student-seq --logit-student local/student-logit \\
           --judge-model local/chat --quant Q4_K_M --engine llama.cpp \\
           --engine-version v0.4.0 --out-dir compare-out --labbook labbook.md

       Add --skip-eval to reuse results files already in --out-dir, which is what you
       want when only the judging step needs rerunning.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
from pathlib import Path
from typing import Any

import distillog

# Which run-log stage paid for which route. The sequence-level student paid for
# generation and filtering as well as its own training; the logit student paid
# only for its training run, because it never generated a dataset.
ROUTE_STAGES = {
    "sequence": ("generate", "filter", "train"),
    "logit": ("train-logit",),
}


def run(script: Path, arguments: list[str]) -> None:
    command = [sys.executable, str(script), *arguments]
    print("+ " + " ".join(command))
    subprocess.run(command, check=True)


def checks(results_path: Path) -> tuple[int, int]:
    data = json.loads(results_path.read_text(encoding="utf-8"))
    return sum(1 for r in data["results"] if r["checks"]["passed"]), len(data["results"])


def route_cost(labbook: str, route: str) -> dict[str, Any]:
    """Add up the cost blocks of the stages that belong to one route.

    A stage that recorded null for a field contributes nothing to that field and is
    counted under `incomplete`, so a total is never quietly made up of half the run.
    """
    totals = {"prompt_tokens": 0, "completion_tokens": 0, "seconds": 0.0, "watt_hours": 0.0}
    present = {k: 0 for k in totals}
    stages_found = []
    for record in distillog.read_stages(labbook):
        if record.get("stage") not in ROUTE_STAGES[route]:
            continue
        stages_found.append(record["stage"])
        cost = record.get("cost") or {}
        for key in totals:
            value = cost.get(key)
            if isinstance(value, (int, float)):
                totals[key] += value
                present[key] += 1
    return {
        "stages": stages_found,
        "totals": {k: (round(v, 2) if isinstance(v, float) else v) for k, v in totals.items()},
        "incomplete": [k for k in totals if present[k] == 0],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--harness-dir", required=True)
    parser.add_argument("--tasks", required=True)
    parser.add_argument("--base-url", default="http://127.0.0.1:8080/v1")
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--sequence-student", required=True, help="the student trained on teacher text")
    parser.add_argument("--logit-student", required=True, help="the student trained against teacher logits")
    parser.add_argument("--judge-model", default=None)
    parser.add_argument("--quant", default="unknown")
    parser.add_argument("--engine", default="llama.cpp")
    parser.add_argument("--engine-version", default="unknown")
    parser.add_argument("--out-dir", default="compare-out")
    parser.add_argument("--skip-eval", action="store_true", help="reuse results files already in --out-dir")
    parser.add_argument("--labbook", default="labbook.md")
    args = parser.parse_args()

    harness = Path(args.harness_dir)
    run_eval, judge = harness / "run-eval.py", harness / "judge.py"
    if not run_eval.is_file():
        raise SystemExit(f"{run_eval} not found; point --harness-dir at your Part 10 evaluation directory")
    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    files = {"sequence": out / "results-sequence.json", "logit": out / "results-logit.json"}
    models = {"sequence": args.sequence_student, "logit": args.logit_student}

    for route, path in files.items():
        if args.skip_eval and path.is_file():
            print(f"reusing {path}")
            continue
        run(run_eval, [
            "--base-url", args.base_url,
            *(["--api-key", args.api_key] if args.api_key else []),
            "--model", models[route],
            "--quant", args.quant,
            "--engine", args.engine,
            "--engine-version", args.engine_version,
            "--tasks", args.tasks,
            "--out", str(path),
            "--notes", f"Part 15 student comparison, {route} route",
        ])

    if args.judge_model:
        if not judge.is_file():
            raise SystemExit(f"{judge} not found, but --judge-model was given")
        run(judge, [
            "compare",
            "--base-url", args.base_url,
            *(["--api-key", args.api_key] if args.api_key else []),
            "--judge-model", args.judge_model,
            "--results-a", str(files["sequence"]),
            "--results-b", str(files["logit"]),
            "--labbook", args.labbook,
        ])
        print("\nRead the flip rate the comparison printed before reading the winner. A high "
              "flip rate means the judge is answering the order of the answers rather than "
              "their quality, and the head-to-head number means little.")

    summary: dict[str, Any] = {}
    print("\n" + "=" * 72)
    print(f"{'route':<12}{'served as':<26}{'checks':>10}")
    for route, path in files.items():
        passed, total = checks(path)
        cost = route_cost(args.labbook, route)
        summary[route] = {"model": models[route], "checks_passed": passed,
                          "checks_total": total, "cost": cost}
        print(f"{route:<12}{models[route]:<26}{passed:>4}/{total:<5}")

    print("\nWhat each route cost, from the run log:")
    for route in files:
        cost = summary[route]["cost"]
        totals = cost["totals"]
        stages = ", ".join(cost["stages"]) or "no stages found in the notebook"
        print(f"  {route:<10} stages: {stages}")
        print(f"             completion tokens {totals['completion_tokens']}, "
              f"{totals['seconds'] / 60:.1f} minutes, "
              f"{totals['watt_hours']} Wh")
        if cost["incomplete"]:
            print(f"             not recorded anywhere: {', '.join(cost['incomplete'])}")

    seq, log = summary["sequence"], summary["logit"]
    print()
    if seq["checks_passed"] == log["checks_passed"]:
        print("The two students passed the same number of deterministic checks. Choose on "
              "cost, on the judge comparison, and on which route you can run again next "
              "month, not on a difference that is not there.")
    else:
        better = "sequence" if seq["checks_passed"] > log["checks_passed"] else "logit"
        print(f"The {better} route passed more deterministic checks on this task set. That is "
              "one task set on one machine: report it with the settings above and treat it "
              "as evidence rather than as a general result.")

    report_path = out / "comparison.json"
    report_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(f"\nwritten to {report_path}")

    rec = distillog.record(
        labbook=args.labbook,
        lab="part-15/compare-students",
        stage="compare",
        teacher=None,
        student={"sequence": models["sequence"], "logit": models["logit"]},
        dataset={"tasks": args.tasks, "tasks_sha256": distillog.file_sha256(args.tasks)},
        hyperparameters={"engine": args.engine, "engine_version": args.engine_version,
                         "quant": args.quant, "judge_model": args.judge_model},
        seed=None,
        scores=summary,
        config_path=__file__,
        notes=None,
    )
    print(f"recorded {rec['stage']} as {rec['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()

Download compare-students.py199 lines

RunnableAll tracks

two students, same tasks, with what each cost

python compare-students.py \
    --harness-dir ~/eval \
    --tasks my-tasks.json \
    --base-url http://127.0.0.1:4000/v1 \
    --sequence-student local/student-distilled \
    --logit-student local/student-logit \
    --judge-model local/chat \
    --quant Q4_K_M \
    --engine llama.cpp --engine-version v0.4.0 \
    --out-dir compare-out \
    --labbook labbook.md

Output — what you should see

route       served as                     checks
sequence    local/student-distilled        ../..
logit       local/student-logit            ../..

What each route cost, from the run log:
  sequence   stages: generate, filter, train
             completion tokens ..., ... minutes, ... Wh
  logit      stages: train-logit
             completion tokens 0, ... minutes, ... Wh

The cost lines are the point. The sequence route paid for generation, filtering and training; the logit route paid only for its own training, and generated nothing it kept. Whether that is cheaper depends entirely on your machine, and now you have the number for yours.

The script also runs Part 10’s judge in compare mode, which asks twice with the answers swapped. Read the flip rate before the winner: a high flip rate means the judge is answering the order of the answers rather than their quality.

7. Optional: the off-policy end, with GKD

The classic offline recipe, where the student is scored against a fixed dataset rather than its own generations, is GKDTrainer with a mixing fraction of zero. TRL documents lmbda=0.0 as reducing “to supervised JSD where the student is trained with the token-level probabilities of the teacher”.

RunnableAll tracks

the off-policy end of the same dial

python train-logit-distil.py \
    --teacher Qwen/Qwen3-8B \
    --student Qwen/Qwen3-1.7B \
    --prompts seeds/prompts.jsonl \
    --trainer gkd --lmbda 0.0 --beta 0.0 \
    --output-dir runs/gkd-offpolicy \
    --labbook labbook.md

--beta 0.0 selects forward KL, the mean-seeking direction, which with lmbda 0.0 is as close as these trainers get to the 2015 recipe. Run it if you have time; the interesting comparison is not which scores better but how much faster the off-policy run is, since it does not generate.

⌚Version note

GKDTrainer is imported from trl.experimental.gkd, a different module from trl, and TRL marks it experimental as of 2026-09-09 against TRL 1.12.0 · verified 2026-09-08. Its own defaults differ from the stable trainer’s: temperature 0.9, lmbda 0.5, beta 0.5, and max_new_tokens 128, which is short enough to truncate many answers. The script passes your values explicitly for exactly that reason.

Check vocabulary alignment and memory before the trainer starts

Compare the teacher and student tokenisers, including token-to-ID mappings and special tokens. Equal vocabulary size alone does not establish that a logit at one index refers to the same event. If the method requires shared tokenisation and the check fails, stop and choose a compatible pair or the sequence-level lab.

Budget for both models plus student training state and activation/logit buffers. Test a tiny batch with the intended sequence length and dtype before the full run. Confirm that teacher parameters are frozen and that gradients reach the student. Record whether examples are teacher-generated, student-generated or supplied from the dataset; this determines which distillation experiment you are performing.

Compare with the sequence-level student using the same held-out task file and deployment precision. Do not compare raw losses from different objectives as a quality ranking. Keep the teacher access method, temperature, divergence configuration, data mixture and model revisions. If the full path does not fit, report the reduced experiment and its changed settings rather than silently offloading or shortening one candidate. The useful result is a controlled task comparison plus measured resource cost, not merely a lower divergence value during training.

Validation

You are done when all of the following are true:

• the tokeniser check reported compatible for your pair, and you have seen what an incompatible pair looks like;
• the memory budget you wrote down before the run is in the notebook next to the peak the script reported;
• the trainer printed a trainable-parameter count matching the arithmetic for your student: about 17.4 million for Qwen3-1.7B at rank 16, about 10.1 million for Qwen3-0.6B;
• the divergence fell from the first logged step and flattened above zero;
• completions/mean_length did not collapse towards the minimum over the run;
• a quantised GGUF of the logit student exists and an engine has generated from it;
• compare-out/comparison.json exists with both routes’ checks and both routes’ costs;
• the run log has a train-logit line and a compare line, and you can state which route you would use again on this machine and why.

Expected outcome

Two students trained two ways, one comparison, and a cost figure for each route on your own hardware. The table is what to measure; the validation pass will fill it in.

Pending validationWhat to record from this lab, per track

Track	Pair	Peak memory, GB	Training wall clock	Final divergence	Checks passed vs sequence student
S: DGX Spark, 128 GB	Qwen3-8B to Qwen3-1.7B	to be measured	to be measured	to be measured	to be measured
X: Ryzen AI Max+ 395	Qwen3-8B to Qwen3-1.7B	to be measured	to be measured	to be measured	to be measured
M: Apple silicon, MPS float32	Qwen3-4B to Qwen3-0.6B	to be measured	to be measured	to be measured	to be measured
N: NVIDIA, 32 GB and above	Qwen3-8B to Qwen3-1.7B	to be measured	to be measured	to be measured	to be measured
N: NVIDIA, 12 to 16 GB	Qwen3-4B to Qwen3-0.6B	to be measured	to be measured	to be measured	to be measured

the four platform tracks, one machine each · TRL DistillationTrainer with a PEFT LoRA adapter on the student transformers 5.16.1, trl 1.12.0, peft 0.20.0 · teacher and student as listed per row, both models BF16 during training, except Track M in float32; Q4_K_M after export · 256 tokens of context · 2026-09-09

Not yet run on hardware on any track. The divergence column is comparable only between rows with the same pair, the same beta and the same temperature; it is not comparable with the cross-entropy loss of the sequence-level lab.

Troubleshooting

The script refuses to run, saying the tokenisers do not match. That is the check working. Pick a same-family pair, or use sequence-level distillation, which needs only text. Override it only to reproduce the fault in this part’s challenge.

Out of memory before the first step. The two models did not fit. In order: drop to the reduced pair, lower --max-completion-length, and confirm nothing else is holding accelerator memory. The budget diagram tells you which bar you are fighting; the teacher’s is the one you cannot shrink.

Out of memory during the first generation. The weights fit and the generation cache did not. Lower --max-completion-length; it bounds both the cache and the number of positions the divergence is computed over.

Extremely slow steps. Generation dominates. On Track S or N, TRL documents vLLM-powered generation for this trainer, which is the intended fix. Otherwise lower --max-completion-length, and accept that on-policy training is slower per example than supervised training by design.

The divergence goes down and entropy collapses. The student is narrowing onto one mode rather than matching the distribution, which is the characteristic risk of reverse KL. Try --beta 0.5, the symmetric divergence, or --beta 0.0 for forward KL, and compare on the evaluation rather than on the loss.

ImportError on trl.experimental.gkd. The experimental modules move between releases. Check the trainer page for the version you have installed rather than assuming the import path here is still right; this is what “experimental” means in practice.

The logit student is worse than the sequence-level student. A legitimate result, and worth reporting as one. Check three things before concluding: that completions/mean_length did not collapse, that the comparison used the same task file and settings for both, and that you served the model you think you served.

Cleanup

Keep both adapters, both quantised exports, compare-out/ and the notebook. The project at the end of this part can run either route, and Part 16’s quantisation lab wants a fine-tuned model to measure.

RunnableAll tracks

reclaim the disk

rm -rf runs/logit-qwen3-1.7b/checkpoint-*
rm -rf models/logit-qwen3-1.7b-merged models/logit-qwen3-1.7b.gguf

The base models stay in the Hugging Face cache. The teacher is the largest download here and is reused by the project.

What you learned

• The tokeniser check is not a formality. It is the difference between logit distillation and a silent, plausible-looking failure, and it costs thirty seconds.
• Two resident models is a different kind of budget. The teacher’s weights are the largest term and nothing shrinks them, which is why this page has the only real memory floor in the part.
• The stable trainer is on-policy. It takes prompts, not answers, and generates the completions it trains on, which is why it costs what it costs.
• A divergence is not a cross entropy. It has a non-zero floor set by the capacity gap, it is comparable only with itself at the same settings, and the metrics beside it say more about the run than the loss does.
• The comparison has two columns. Quality and cost. On one machine the logit route wins; on another the sequence-level route wins by not needing the memory at all, and now you know which yours is.

Record in the notebook: the tokeniser report; the predicted and actual peak memory; the trainable parameter count; beta, temperature and max_completion_length; the first and final divergence; whether completions/mean_length or entropy collapsed; the two students’ scores on the same task file; and the cost of each route from the run log.

Check your understanding

Question 1. Why does this lab check the tokenisers before it checks anything else?

• Because loading a tokeniser is fast
• Because a vocabulary mismatch does not raise an error: the loss is well defined, it falls, and the student is trained towards the wrong tokens, which is the one failure mode the run itself will not reveal
• Because TRL requires it
• Because the tokeniser decides the learning rate

Show the answer and why

Answer: Because a vocabulary mismatch does not raise an error: the loss is well defined, it falls, and the student is trained towards the wrong tokens, which is the one failure mode the run itself will not reveal

TRL's documentation warns that a teacher with a different vocabulary trains against the wrong tokens, silently when its vocabulary is no larger than the student's. Everything else in the run has a visible symptom; this one does not.

Question 2. TRL's DistillationTrainer takes a prompt-only dataset. What does that tell you about how it works?

• It fills the completions from the teacher
• It is on-policy: the student generates the completions it is trained on, and the teacher is asked what it would have said at each position
• It only supports classification
• It requires vLLM

Show the answer and why

Answer: It is on-policy: the student generates the completions it is trained on, and the teacher is asked what it would have said at each position

That is why the teacher answers from the previous lab are not used here, and why every step costs a generation. It also means the completions the student trains on drift as the student improves, which is the point of on-policy distillation.

Question 3. Your divergence falls quickly and then flattens well above zero. What should you conclude?

• The run is stalled and the learning rate is too low
• The student lacks the capacity to represent the teacher's distribution exactly, so a non-zero floor is the expected shape; judge the run on the evaluation, not on the loss
• The tokenisers do not match
• The teacher is quantised

Show the answer and why

Answer: The student lacks the capacity to represent the teacher's distribution exactly, so a non-zero floor is the expected shape; judge the run on the evaluation, not on the loss

A smaller model cannot match a larger one exactly, so the divergence has a floor set by the capacity gap. Raising the learning rate to push it lower usually produces a student that has narrowed onto one mode, which reads as a better loss and a worse model.

Question 4. Which of these are real reasons the sequence-level route may beat the logit route on your machine? Select all that apply.

• It needs only one model resident at a time
• Its generation cost is paid once and the dataset is reusable
• It works when teacher and student have different tokenisers
• Its loss is always lower

Show the answer and why

Answer: It needs only one model resident at a time, Its generation cost is paid once and the dataset is reusable, It works when teacher and student have different tokenisers

The two losses measure different things and cannot be compared at all, which is why the last option is wrong. The first three are the practical advantages, and on a 24 GB machine the first one alone often settles the decision.

Question 5. What does beta = 0.0 do in TRL's DistillationConfig, and when might you reach for it?

• It disables the teacher
• It selects forward KL, the mean-seeking direction, which is worth trying when the student's entropy is collapsing under the reverse-KL default
• It sets the learning rate to zero
• It adds a KL penalty against a reference model, as in GRPO

Show the answer and why

Answer: It selects forward KL, the mean-seeking direction, which is worth trying when the student's entropy is collapsing under the reverse-KL default

This beta selects the divergence itself and is not GRPO's reference-model penalty; the documentation says so explicitly. Reverse KL, the default, is mode-seeking, and a collapsing entropy alongside a falling divergence is its characteristic failure.

Check my answersClear