Train your own draft model.

A draft model is the small, fast helper that guesses the next few words so your big model only has to check them, and this lesson is about training that helper to guess the way your big model would, not just to write well.

On your Spark that is the difference between a speed-up you can barely see and one you can put on the scoreboard. Every guess the big model throws away is wasted work, and every guess it keeps is a word you got almost for free. The helper you used last time is a decent writer with its own habits, so a lot of its guesses get thrown away. Teach it your big model’s habits and more of them survive, with no change at all to what your server finally writes. The same move, a small model trained to copy a big one, is how labs make cheap fast models out of expensive slow ones.

This picks up exactly where Speculative decoding left off, with your guess-and-check decoder, its speed formula and your best result so far. It leans on Memory math for models to decide what to save to disk, and it is a small, complete training run, a cousin of Train a small LM. Next, Fine-tuning and LoRA shows how to do this kind of retraining while touching only a sliver of the model, and Inside vLLM and SGLang is where the fancier helpers you meet at the end get deployed for real.

Today

• Say precisely what makes a helper model good: not “smart”, but close to the big model in how it spreads its bets over the next word. Show that the fraction of guesses kept is exactly the overlap of the two models’ bets, which is one minus a standard measure of how far apart they are.
• Explain how a small “student” model learns from a big “teacher”: it trains on the teacher’s full list of odds for every possible next word instead of on the one word that was picked, why that carries far more signal per example, what the knob that flattens those odds does, and the two directions you can measure “how different” in, which push the student differently.
• Build a training set from your own 8B: have it answer prompts like your real traffic, then save only its 64 favourite next words (with their scores) per position, about 388 bytes instead of 250 KB.
• Retrain Llama 3.2 1B two ways (once on the 8B’s text alone, once on its saved odds), build a second helper from the 8B’s own first 4 layers, and measure how many guesses get kept before and after for each.
• Plug the winner into your speculative decoder, record the end-to-end speed-up on SCOREBOARD.md, and, as a stretch, train a few tiny add-ons on the big model itself that do the guessing with no second model at all.

You need Speculative decoding first: you will reuse your spec_greedy, spec_sample and sweep.py, and the speed formula built from how often guesses are kept, how many you make per round and how much each one costs. Why ~17 tokens per second? and Memory math for models help. If you haven’t done Train a small LM yet, this is probably your first time training a language model, and a gentle first taste of the full loop. Plan on an evening for the idea and a weekend for the build, most of it waiting for the Spark.

The Problem

Your K sweep from the speculative decoding lesson ended somewhere like this: Llama 3.2 1B drafting for Llama 3.1 8B, greedy, α̂ around 0.7, c around 0.3, best K of 2 or 3, about 1.4x. A real gain. And the calculator told you which levers are left. Lowering c means engineering on the draft’s step (CUDA graphs, fused kernels). Raising α means a better draft.

So you print the rejections. Pick twenty rounds where d1 was thrown out and look at what the 1B proposed against what the 8B wanted. You might see something like this:

• The 8B starts answers with “Here’s a” and the 1B guesses “Sure”.
• The 8B likes a markdown header before a list; the 1B goes straight to “1.”
• In code, the 8B writes math.pi * r ** 2 and the 1B writes 3.14 * r * r.

None of the 1B’s guesses are wrong. They are fine English and working Python. They are just not what this 8B would say. The 1B is a perfectly decent writer with its own habits, and you asked it to impersonate a different writer whose habits it has never studied.

That reframes the job. You don’t need a smarter draft. You need a better impersonator. And you own the thing it should impersonate: the 8B is sitting on your Spark, happy to show its work on as much text as you like, including its full probability vector at every step, not just the token it picked. Training a small model to copy a big model’s probabilities has a name, and it is older than LLMs.

The Idea

What a good draft is, exactly

Recall the rule from the last lesson. At each step the draft proposes x from its distribution q; the target keeps it with probability min(1, p(x)/q(x)). The chance that a proposal is kept, averaged over what the draft might propose, is

P(keep)  =  Σ_x q(x) · min(1, p(x)/q(x))  =  Σ_x min(p(x), q(x))

That’s the area of the shared part of the two bar charts, the solid blue in Figure 003 of the last lesson. There is a standard name for how far apart two distributions are in exactly this sense. The total variation distancetotal variation distanceHalf the sum of absolute differences between two probability lists: ½ Σ |p(x) − q(x)|. It’s the largest possible difference in the probability the two give to any single event. 0 means identical, 1 means no overlap at all., TV for short, is ½ Σ |p − q|. A two-line identity connects them: for any two numbers, |p − q| = p + q − 2·min(p, q), so summing over the vocabulary,

TV(p, q)  =  ½ (1 + 1 − 2 Σ min(p, q))  =  1 − Σ min(p, q)

so   α  =  1 − TV(p, q)

Leviathan et al. state this in their paper too. (I checked the identity numerically on random 50-token distributions; the two sides agree to 16 digits, as they must.)

Read that as a design spec. The acceptance rate at one step is a pure function of how far apart the two probability vectors are, in one specific sense. Not of how good the draft is at anything. Average it over the steps your server actually sees, and you have the α in the speedup formula. To raise α, make q close to p, on the text your server writes. That’s the whole job.

This also explains why some tokens are easy and some hard. After “The capital of France is”, both models put 0.8+ on ” Paris”: huge overlap. After “Here is a poem about”, both are spread out over hundreds of tokens, and the overlap of two spread-out distributions with different preferences is small. The draft is “wrong” there in the sense that it has different taste.

Things to try:

• The start. An off-the-shelf small model, α = 0.63. Over-proposing NONE and SELF costs as much as under-proposing MATH: both are probability in the wrong place.
• Distill halfway (0.50). α = 0.82, and the KL from p to q drops from 0.32 to 0.08 nats. At K = 3, c = 0.3, the speedup goes from about 1.2x to about 1.6x. Same model size, same c, one change: it imitates better.
• K and c. With the distilled draft, push K up. The best K grows, because long runs now survive. Then raise c and watch a slow draft eat the gain.

Learning from the teacher’s whole answer

The obvious way to make the 1B write like the 8B: have the 8B write a lot of text, and fine-tune the 1B on it with the ordinary language-modelling loss. That works, and it has a name: sequence-level distillationsequence-level distillationKim and Rush, 2016. Train the student on full outputs generated by the teacher, using the normal next-token loss. The teacher’s text replaces human-written text.. The 1B learns the 8B’s habits (“Here’s a”, the markdown header) because the training text is full of them.

But look at what each training example tells the student. The ordinary loss is cross-entropycross-entropyThe standard next-token loss: minus the log of the probability the model gave to the token that actually came next. Low when the model put high probability on it. against a one-hotone-hotA probability list that is 1 for one entry and 0 for all others. A “label” that says only which token came next, and nothing about the alternatives. label. The 8B sampled ” math” at some step, so the label says “math: 1, everything else: 0”. But the 8B knew more than that. It had a full vector: ” math” 0.46, ” r” 0.14, ” 3” 0.10, and so on. Sampling one token from it threw nearly all of it away.

knowledge distillationknowledge distillationTraining a small “student” model to match a big “teacher” model’s output probabilities, not just its chosen outputs. The teacher’s full probability vector is the training target. keeps it. Instead of the one-hot label, the target at each position is the teacher’s whole vector, the soft labelssoft labelsThe teacher’s full probability vector used as a training target, instead of a single correct answer. Also called soft targets. “Soft” because the probability is spread over many entries.. The loss measures how far the student’s vector is from the teacher’s. The natural choice is the KL divergenceKL divergenceA measure of how different one probability distribution is from another, in nats (natural log) or bits. Zero only if they are identical; never negative; not symmetric, so KL(p‖q) and KL(q‖p) are different numbers.:

KL(p ‖ q)  =  Σ_x p(x) · log( p(x) / q(x) )
           =  Σ_x p(x) · (−log q(x))   −   Σ_x p(x) · (−log p(x))
              └── cross-entropy against p ──┘   └── teacher's entropy: constant ──┘

The second line is worth staring at. KL against the teacher is cross-entropy where the label is the whole vector p instead of a one-hot, minus a constant the student can’t change. So “distillation loss” is just “the loss you already know, with a soft label”. The gradient with respect to the student’s logits is the familiar q − p, the same shape as ordinary training’s q − onehot. One line of code changes.

Why it matters so much: a one-hot label is one sample from p. The student has to see the same kind of context many times before the counts reveal p’s shape. A soft label hands it p directly.

Temperature, and the T² factor

Hinton’s paper had one more knob. A trained model’s distribution is often very peaky: 0.97 on one token and tiny amounts spread over the rest. The tiny amounts carry the dark knowledge, but they contribute almost nothing to the loss, so the student barely learns them. Their fix was to divide both models’ logits by a distillation temperaturedistillation temperatureA number T > 1 that divides the logits before softmax during distillation, flattening both the teacher’s and student’s distributions so the small probabilities matter more in the loss. T before the softmax. T = 2 or 4 flattens both distributions, which makes the small entries visible to the loss. One side effect: dividing logits by T shrinks the gradients by about 1/T², so they multiply the loss by T² to keep its scale.

For a draft model, be careful with this knob. Hinton wanted the student to learn the teacher’s knowledge, and evaluated it on its top answer. You want the student to match the teacher’s distribution at the temperature you serve, because that’s what α is computed on. A student trained to match flattened distributions has been optimizing for something slightly different. So start at T = 1 (the real distribution), and if you try T = 2, measure α at your serving temperature, not at 2. One of the Break It experiments does exactly this.

Which way round the KL goes

KL isn’t symmetric, and the two directions push the student differently when it is too small to copy the teacher perfectly. Which, for a 1B imitating an 8B, it always is.

• Forward KL, KL(p ‖ q), averages log(p/q) weighted by the teacher’s p. It is huge wherever the teacher has real probability and the student has almost none. So the student is pushed to cover everything the teacher might say, even if that means spreading thin. This is called mode-coveringmode-coveringBehavior of forward KL: the student spreads its probability to cover every place the teacher has mass, even if it must put mass in between where the teacher has none.. It’s what Hinton used, and what “soft-label cross-entropy” computes.
• Reverse KL, KL(q ‖ p), is weighted by the student’s own q. It is huge wherever the student puts probability on something the teacher considers unlikely, and it doesn’t care what the student ignores. So the student hides inside the teacher’s biggest peak and gives up the rest: mode-seekingmode-seekingBehavior of reverse KL: the student concentrates on one high-probability region of the teacher and ignores the others, rather than putting mass anywhere the teacher wouldn’t.. MiniLLM (Gu et al., 2024) argued for this direction for distilling LLMs.

Which matches acceptance? Neither, exactly. α = 1 − TV, and TV is symmetric: it charges the same for a unit of probability in the wrong place whichever direction it’s wrong in. Both KLs are bounds on it (by Pinsker's inequalityPinsker's inequalityA classic result: TV(p, q) ≤ √(KL/2), for either direction of KL. Driving either KL to zero forces TV to zero, but two students with the same KL can have quite different TV.), so driving either to zero drives α to 1. But at the student’s limit, when neither can reach zero, they disagree about which mistakes to make, and that changes α.

The figure is the cleanest version of that. The teacher has two bumps. The student can only make one bump (a mean and a width). Pick an objective; the figure finds the best student under it and reports α.

What you should see (and what I got in numpy at the same settings):

• Second bump 0.35. Forward KL straddles both bumps: α = 0.557, and its top token sits in the valley where the teacher has almost nothing. Reverse KL takes the big bump: α = 0.650. Maximizing α directly: 0.656, close to reverse KL.
• Second bump 0.50. Now forward KL (0.558) beats reverse KL (0.500), and maximizing α directly agrees with forward KL. Giving up half the mass is too expensive.
• The greedy column. If you serve greedily, acceptance is just “does the draft’s top token match the target’s top token”. Forward KL’s straddling student can fail that badly even when its α at temperature 1 is fine.

Real students aren’t one-bump toys, and the effect is milder. DistillSpec’s finding was exactly this kind of answer: the best divergence depends on the task and on how you decode (greedy or sampling, which temperature). There’s no universal winner. That’s a measurement, not a debate, so in the Build you train with forward KL (the standard, and easiest with stored logits) and you can try others cheaply afterward. You can even train on 1 − Σ min(p, q) directly, the true objective; DistillSpec includes TV among the losses it compares.

Whose text to learn on

Distillation needs contexts: prefixes at which you ask the teacher “what would you say next?” Which prefixes?

The draft is used at very specific moments: after the real text so far (which the target wrote), for d1; and after the real text plus the draft’s own earlier guesses, for d2 through dK. So you want the draft to be good on the target’s text, on prompts like your traffic, and ideally on its own half-finished guesses too.

Three choices, in order of how closely they match that:

1. Generic text (web pages, books). What the 1B was pretrained on. Wrong prompts, wrong writer.
2. The target’s generations on prompts like your traffic. Right prompts, right writer. This is what you’ll build: off-policyoff-policyTraining on text produced by some other process than the model being trained, here by the teacher. Cheap to produce once and reuse, but the student never sees its own mistakes. data, made once, reused for every experiment.
3. The draft’s own generations, scored by the target. The student writes, the teacher reports its probability vector at every position of the student’s text, and the student is trained toward those vectors. This is on-policy distillationon-policy distillationTraining the student on text it generated itself, with the teacher’s probabilities at each position as the target. It teaches the student what the teacher would do after the student’s own mistakes. Agarwal et al. 2024 (GKD); used by DistillSpec.. It fixes a problem with choice 2 called exposure biasexposure biasA model trained only on someone else’s text has never seen prefixes containing its own typical mistakes, so it handles them badly when it meets them at run time.: a student trained only on the teacher’s text has never seen a prefix containing its own slips, but at d2 through dK of every round, it is reading exactly that.

DistillSpec (Zhou et al. 2023) found that drafts trained on student-generated data did as well as or better than teacher-generated data, and student data is much cheaper to make, since the small model is the one writing. The catch for you is plumbing: on-policy data has to be regenerated as the student changes, so you need the teacher running during training. On the Spark both fit easily (16 GB plus a few GB), so it’s a stretch task, after the off-policy version works.

The dataset, and why you store only the top 64

Here is the pipeline you’ll build.

The memory math is the reason for the top-64 box, and it’s a good example of the habit from Memory math for models: multiply it out before you write the code.

The other number per position is the teacher’s logsumexplogsumexplog Σ exp(logit) over the whole vocabulary: the normalizer of the softmax. With it, any single logit can be turned into an exact probability: p(x) = exp(logit(x) − logsumexp).. With it, your 64 stored logits become 64 exact probabilities, and one minus their sum is the exact tail masstail massThe total probability the teacher gives to all tokens outside the stored top 64. Unknown token by token, but its total is known if you stored the logsumexp. the teacher spreads over the other 128,192 tokens. Without it, all you could do is renormalize the 64 to sum to 1, which quietly claims the teacher puts zero probability everywhere else.

Is 64 enough? A useful fact: if the student matched the renormalized top-k of the teacher perfectly, its α would be exactly the teacher’s top-k mass (the student would give the top-k a bit more than p on each, and the shared part is all of p’s top-k). So the average top-64 mass on your data is roughly a ceiling on what training from stored logits can teach, and you’ll measure it. For peaky chat text it’s likely high, with a tail of open-ended positions where it isn’t. Keeping the tail mass as its own term in the loss (one “everything else” bucket) recovers a little of the rest.

Why store at all, instead of running the 8B live during training? You could: the Spark holds both. But a training step for a 1B costs about 6 × 1.2 B ≈ 7 GFLOP per token (forward and backward), while one 8B forward costs about 2 × 8 B = 16 GFLOP per token. Running the teacher live roughly triples the cost of every epoch of every experiment. Scoring once and storing is the cheap, reusable choice, and live teachers are for on-policy training, where the text keeps changing.

What shape should a draft be

Two drafts you’ll build, and a few you’ll meet.

A: a separate small model is what you have: Llama 3.2 1B, 16 layers, 2048 wide, with its own embedding table shared with its output layer (tied embeddingstied embeddingsUsing the same matrix for the input embedding lookup and the output layer that produces logits. Saves parameters; Llama 3.2 1B does it, Llama 3.1 8B does not.). Its only input is tokens. Distillation will make it a better impersonator. Hard rule: it must use the same tokenizer as the target. Token id 1234 has to mean the same string to both, or every comparison of p and q is comparing different things. Llama 3.2 and Llama 3.1 share one; check anyway, the way your common.py already does.

B: a truncated copy of the target. Take the 8B, keep its embedding table, its first N layers, its final normalization and its output layer, and throw away layers N+1 through 32. Now it’s a draft. Its arithmetic, for N = 4: each 8B layer is about 218 M parameters (attention projections plus the big MLP), so 4 layers are 0.87 B. The output layer is 4096 × 128,256 ≈ 525 M parameters, 1.05 GB in bf16, read every step. The embedding table is a lookup, so only one row is read per token. Total read per draft step: about 2.8 GB, a bit more than the 1B’s 2.5 GB.

So why bother? Two reasons. First, it is shallow and wideshallow and wideFew layers, each with a large hidden size. At batch 1, a draft step’s time is partly set by how many layers run one after another, since each layer launches its own kernels; a shallow model has fewer sequential steps for the same bytes.: 4 layers instead of 16, so a quarter of the per-layer kernel launches. Your c measurement in the last lesson probably showed the 1B was far slower than its bytes predicted, for exactly this reason. Second, it starts inside the target: same embeddings, same output layer, and its first 4 layers are literally the target’s. In the same process you can point the draft’s embedding and head at the target’s tensors, so they cost zero extra memory. You’ll freeze those and train only the 4 layers.

C: Medusa heads (Cai et al., 2024) drop the separate draft entirely. During verification the target already computes its final hidden state at every position, the vector right before the output layer. Medusa bolts a few tiny Medusa headsMedusa headsSmall extra modules on the target’s final hidden state. Head k takes the hidden state at position t and predicts the token at t + k + 1, skipping over the tokens in between. The target itself stays frozen (Medusa-1) or is fine-tuned jointly (Medusa-2). onto that vector. In the paper each head is one residual block: it adds a small learned correction to the vector, h + SiLU(W·h), where SiLU is a smooth ReLU-like curve and W is a 4096 × 4096 matrix (17 M parameters), followed by its own copy of an output layer. Head 1 guesses the token two ahead, head 2 three ahead, and so on, all from the same h. Train them on the target’s own text while the target stays frozen. Since they ride on the verify pass you were running anyway, there is no separate draft loop: c is close to zero.

The trade: head 2 has to guess the token three positions ahead without knowing the two tokens in between. That’s much harder than a draft model that sees its own previous guesses, so per-head accuracy falls quickly with distance. Medusa compensates by having each head propose its top few candidates and verifying a small tree of combinations in one pass (tree attentiontree attentionVerifying several candidate continuations in one target pass by laying them out as a tree and using an attention mask so each branch only sees its own ancestors.). The paper reports about 2.2x for Medusa-1 and 2.3 to 3.6x for Medusa-2. One honest caveat: for sampling, Medusa’s default “typical acceptance” rule is a relaxation that is not exactly the target’s distribution. With greedy decoding, or with the standard rejection rule you already wrote, it stays exact.

D: EAGLE (Li et al., 2024) is among the strongest practical drafting methods today. A draft that only sees tokens has to rebuild, from scratch, everything the target already worked out. So EAGLE’s draft is a single transformer layer that reads the target’s top-layer features (plus the embedding of the next token), predicts the next feature vector, and turns it into a token using the target’s own output layer. Then it loops on its own predicted features to draft further. It sees what the target was “thinking”, not just what it said. EAGLE-2 made the draft tree’s shape depend on the draft’s confidence, and EAGLE-3 changed the training so the draft practices on its own predictions, the same “learn from your own slips” idea as on-policy distillation. vLLM and SGLang ship EAGLE support; you’ll meet it in Inside vLLM and SGLang.

You’ll build A and B, and C as a stretch. D is the natural next step if you want to go further, and the reason it works is something you’ll feel in t9: the target’s hidden state is a far richer input than a few tokens.

Measuring a draft without running speculation

Running the full speculative decoder for every candidate is slow and noisy. There’s a shortcut you get for free from the identity at the top of this lesson.

Take held-out text the target generated. Run both models over it with the true prefix at every position (this is teacher forcingteacher forcingFeeding a model the true previous tokens at each step, instead of its own outputs, so every position can be scored in one parallel pass., the same mini-prefill trick as verification). At every position you now have both full vectors p and q. Compute Σ min(p, q) there: that is exactly the probability the first draft of a round would be kept at that position. Average it. That’s an offline αoffline αThe acceptance rate estimated directly from both models’ probability vectors on held-out target text, without running speculative decoding. Exact for the first draft position; an approximation for later ones., with no sampling noise from the accept/reject coin flips at all.

For greedy serving, the matching number is top-1 agreementtop-1 agreementThe fraction of positions where the draft’s most likely token equals the target’s most likely token. The acceptance rate for greedy speculative decoding.: how often argmax q == argmax p. Compute that on held-out text the target generated greedily, since that’s the text your greedy server will see.

Offline α is what you’ll use to compare training runs. The real decoder then confirms it, and adds the two things offline α can’t see: positions 2 through K, where the prefix contains the draft’s own guesses, and c.

Build It

Add to your inference/ folder:

inference/
  draft/
    prompts.py         # t1: build train / held-out prompt sets
    generate.py        # t2: 8B writes the answers (vLLM, batched)
    score_topk.py      # t3: 8B re-reads them, store top-64 logits
    offline_alpha.py   # t4: Σ min(p,q) and top-1 agreement, held-out
    train_draft.py     # t5, t6: CE / KL fine-tune, truncated target
    onpolicy.py        # t7 (stretch)
    medusa.py          # t9 (stretch)
  spec/                # from the last lesson: sweep.py gains per-position stats (t8)
  checkpoints/

The models: Llama 3.1 8B Instruct (the target and teacher, the model your scoreboard tracks) and Llama 3.2 1B Instruct (the student). Everything runs in the same NVIDIA PyTorch container as before. Training a 1B in bf16 with AdamW needs roughly 16 bytes per parameter for weights, gradients and optimizer state (about 20 GB), plus activations; the training memory lesson does that math properly. Comfortable on 128 GB, even with the 8B loaded alongside for scoring.

# draft/prompts.py
import json, random
from datasets import load_dataset

random.seed(0)
ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft")
firsts = [ex["messages"][0]["content"] for ex in ds if 20 < len(ex["messages"][0]["content"]) < 2000]

# never train on what you benchmark with
bench = {json.loads(l)["prompt"] for l in open("../prompts/short.jsonl")}   # adjust field name
firsts = [p for p in dict.fromkeys(firsts) if p not in bench]               # dedupe, drop bench prompts
random.shuffle(firsts)

splits = {"train": firsts[:10_000], "heldout": firsts[10_000:10_500]}
for name, ps in splits.items():
    with open(f"prompts_{name}.jsonl", "w") as f:
        for p in ps: f.write(json.dumps({"prompt": p}) + "\n")

# draft/generate.py  -- check the vLLM docs for your version; the shape is stable, names drift.
import json, sys
from vllm import LLM, SamplingParams

split, temp = sys.argv[1], float(sys.argv[2])       # e.g. train 1.0, heldout 1.0, heldout 0.0
prompts = [json.loads(l)["prompt"] for l in open(f"prompts_{split}.jsonl")]
llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct", dtype="bfloat16", max_model_len=4096)
params = SamplingParams(temperature=temp, top_p=1.0, top_k=-1, max_tokens=256, seed=0)
outs = llm.chat([[{"role": "user", "content": p}] for p in prompts], params)

with open(f"gen_{split}_T{temp}.jsonl", "w") as f:
    for o in outs:
        f.write(json.dumps({"prompt_ids": o.prompt_token_ids,
                            "resp_ids": list(o.outputs[0].token_ids)}) + "\n")

@torch.inference_mode()
def score_one(target, prompt_ids, resp_ids, k=64):
    """Returns (topk_ids int32 [R, k], topk_logits fp16 [R, k], lse fp32 [R]) where R = len(resp_ids).
    Row j describes the 8B's distribution for response token j, given everything before it."""
    ...

@torch.inference_mode()
def offline_alpha(draft, target, seqs, T=1.0):
    """seqs: list of (prompt_ids, resp_ids). For every response position, compute
       a_t   = Σ_x min(p_t(x), q_t(x))     with p, q = softmax(logits / T), full vocab, fp32
       top_t = argmax p_t == argmax q_t
    Return mean a_t, mean top_t, mean KL(p_t ‖ q_t), and the per-position a_t array."""
    ...

def ce_loss(student_logits, resp_ids):
    """(a) Sequence-level distillation: ordinary next-token cross-entropy on the 8B's text.
    student_logits: [R, V] rows predicting each response token (same alignment as t3)."""
    return F.cross_entropy(student_logits.float(), resp_ids)

def kl_loss(student_logits, t_ids, t_logits, t_lse, T=1.0):
    """(b) Token-level distillation: forward KL(p ‖ q) over the stored top-64 plus one tail bucket.
    t_ids [R, k] int64, t_logits [R, k], t_lse [R]."""
    logq = F.log_softmax(student_logits.float() / T, dim=-1)        # full vocab
    logq_k = logq.gather(-1, t_ids)                                 # student at the teacher's top-k
    p_k = torch.exp(t_logits.float() / T - t_lse[:, None])          # exact at T = 1 (see note)
    # tail: p_tail = 1 − Σ p_k, q_tail = 1 − Σ exp(logq_k). Add p_tail · log(p_tail / q_tail).
    # loss = Σ_k p_k · (log p_k − logq_k) + tail term, averaged over rows, times T².
    ...

import copy
def truncated_draft(target, n_layers=4):
    """A copy of the target with only its first n decoder layers. Embedding, final norm and
    output layer are the target's own tensors (shared, frozen). Only the layers train."""
    d = copy.deepcopy(target)          # simplest; or build from config and load weights
    d.model.layers = d.model.layers[:n_layers]
    d.config.num_hidden_layers = n_layers
    # share and freeze: point d.model.embed_tokens.weight and d.lm_head.weight at the target's
    # parameters, set requires_grad=False on them and on d.model.norm, train only d.model.layers
    ...

def onpolicy_step(student, target, prompt_batch, opt, T=1.0, max_new=128):
    """1) student generates (no grad), 2) target scores (no grad, full logits or top-64),
    3) KL(p ‖ q) on the student's own sequences, one optimizer step."""
    ...

python bench.py --target spark-spec-distilled --base-url http://localhost:8000/v1 \
  --model meta-llama/Llama-3.1-8B-Instruct --precision bf16 \
  --prompts prompts/short.jsonl --concurrency 1,4,16 --n 200 --max-tokens 256 --ignore-eos \
  --change "spec decoding: distilled 1B draft (KL, top-64), K=<yours>, greedy"

class MedusaHead(nn.Module):
    """h + SiLU(W h), then the 8B's frozen output layer. Head k predicts token t + 1 + k from h_t."""
    def __init__(self, d=4096):
        super().__init__()
        self.w = nn.Linear(d, d)
        nn.init.zeros_(self.w.weight); nn.init.zeros_(self.w.bias)   # starts as the identity
    def forward(self, h):
        return h + F.silu(self.w(h))

# loss for head k on a sequence: cross_entropy(lm_head(head_k(H[:-(k+1)])), ids[k+1:])

Measure

The dataset (t2, t3).

Offline α (t4-t7). Held-out, teacher forced, full logits.

In the decoder (t8, t9). Greedy, 20 short prompts, 256 tokens each, each draft at its own best K.

What to expect. These are rough expectations to check, not results.

• Dataset size. Within a few percent of 388 bytes × tokens, or something is padded. Top-64 mass: probably high on average for chat answers, with a long low tail on open-ended positions. If the mean is well under 0.9, try k = 128.
• CE vs KL. Both should raise α over the off-the-shelf 1B, because both teach the 8B’s habits. Expect KL to beat CE at the same number of steps, by more on α at T = 1 than on top-1 agreement (CE on sampled text is already a decent signal for “what’s the top token”). How much? DistillSpec’s 10 to 45% extra speedup is the published range for their setups; a gain of several points of α is a reasonable thing to hope for. Measure.
• The truncated draft. Near-useless before training, as predicted. After training it should improve a lot. Whether it passes the 1B depends on how much 4 layers can learn from 2M tokens; it starts with the target’s own representations, but it has far fewer parameters doing the work than the 1B’s 16 layers. Its c may be noticeably lower than the 1B’s in eager PyTorch. Whichever wins, the comparison is the lesson.
• Offline vs online. Offline top-1 agreement should match position-1 acceptance in the greedy decoder to within a couple of points. Later positions: flat or rising is common, for the selection reason in the last Predict.
• Speedup. Use last lesson’s calculator, or Figure 001 here. At c = 0.3, going from α = 0.7 to 0.8 takes the best speedup from about 1.37x (K = 2) to about 1.55x (K = 3). At c = 0.15, the same α change goes from about 1.75x to 2.1x. A distilled draft helps more the cheaper your draft step is, so this lesson and a lower c multiply.
• Medusa. Head 1 alone may rival the distilled 1B’s position-1 acceptance or fall short of it; heads 2 and 3 will be clearly worse, since they guess blind. But c is near zero, so even modest acceptance can win. Without tree verification you’re leaving some of Medusa’s gain on the table; that’s fine for a first build.
• Scoreboard. Tok/s at concurrency 1 should rise by roughly the ratio of the new speedup to the old one. Throughput at 16 is still bounded by the lock. Same story as last lesson: faster per request, not wider.

Break It

Off by one. In kl_loss, shift the teacher rows by one position relative to the student’s (use rows s through len − 1 instead of s − 1 through len − 2). Train for a few hundred steps. The loss still goes down, which is the scary part: the student is learning to predict the token after the next one, which is possible, just hard. Offline α collapses. Lesson: a falling loss proves the student learned something. Only α tells you it learned the right thing. Keep the alignment self-check from t3.

Hinton’s temperature, wrong job. Distill with T = 4 (renormalizing over the top-64, times T²), then measure offline α at T = 1 and top-1 agreement. Compare with the T = 1 run. The student has learned to match flattened distributions, which usually means it spreads too much probability at serving temperature, and α drops, especially at T = 1. Lesson: the distillation temperature is a choice about what to copy. For a draft, copy what you serve.

Top-1 only. Set k = 1 in the KL loss and drop the tail term. That’s cross-entropy against the teacher’s argmax, a one-hot label again (just a different one from the sampled token). Watch α at T = 1 fall back toward the CE run while top-1 agreement holds up or even improves. Lesson: the soft part of soft labels is exactly what makes a draft good at sampling.

Leak the benchmark. Add your prompts/short.jsonl prompts (with the 8B’s answers) to the training set and retrain. Your scoreboard speedup jumps. Then run the sweep on a fresh set of 20 chat prompts the draft has never seen, and watch it come back down. Lesson: drafts memorize. A speedup measured on training prompts is a speedup on your test, not on your traffic. (A real server sees repeats, which is why caching and n-gram lookup exist, but that’s a different claim, and you should measure it as one.)

Explain It Back

Answer each out loud before opening it.

Go Deeper

• “Distilling the Knowledge in a Neural Network”, Geoffrey Hinton, Oriol Vinyals, Jeff Dean, 2015 (arXiv 1503.02531). Short and readable. Soft targets, temperature, the T² factor, and the “dark knowledge” argument.
• “Model Compression”, Cristian Buciluă, Rich Caruana, Alexandru Niculescu-Mizil, KDD 2006. The original: label lots of data with a big model’s outputs and train a small one to copy them.
• “Sequence-Level Knowledge Distillation”, Yoon Kim and Alexander Rush, EMNLP 2016. Why training on the teacher’s generated text works, and how it relates to token-level KL. Your (a) vs (b) comparison, for translation.
• “DistillSpec: Improving Speculative Decoding via Knowledge Distillation”, Yongchao Zhou et al., ICLR 2024 (arXiv 2310.08461). This lesson’s direct source: distilling drafts for acceptance, on-policy data, and the comparison of FKL, RKL, JSD and TV across tasks and decoding strategies.
• “On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes” (GKD), Rishabh Agarwal et al., ICLR 2024 (arXiv 2306.13649). The general method behind t7, and a clean discussion of exposure bias in distillation. TRL ships a GKDTrainer based on it.
• “MiniLLM: Knowledge Distillation of Large Language Models”, Yuxian Gu, Li Dong, Furu Wei, Minlie Huang, ICLR 2024. The case for reverse KL when distilling generative models, with good intuition about mode-seeking.
• “Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads”, Tianle Cai et al., ICML 2024 (arXiv 2401.10774). The heads, tree attention, typical acceptance, and Medusa-1 vs Medusa-2. The code repo (FasterDecoding/Medusa) is small enough to read.
• “EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty”, Yuhui Li, Fangyun Wei, Chao Zhang, Hongyang Zhang, ICML 2024 (arXiv 2401.15077), then EAGLE-2 and EAGLE-3. Why drafting on features beats drafting on tokens. The natural next build after this lesson.

Connections

Built on: Speculative decoding (the acceptance rule, α, the K sweep, your decoder) and Memory math for models (the 510 GB vs 0.78 GB decision). The bandwidth argument for sharing Medusa’s output layer is Why ~17 tokens per second? again.

This was a small, complete training run: data generation, a scoring pass, a training loop, held-out evaluation, and a downstream metric that isn’t the loss. Train a small LM, end to end does the same loop from raw text at larger scale, Training memory explains the 16 bytes per parameter you budgeted for, and Fine-tuning and LoRA shows how to do this kind of fine-tune with a tiny fraction of the trainable parameters. The habit of judging a model by a held-out metric that matches the real goal, not by its training loss, is the whole subject of why evals are separate.

On the inference side, Inside vLLM and SGLang is where you’d deploy a real EAGLE or Medusa draft, and Quantization stacks with all of it: a 4-bit draft lowers c, and a lower c makes every point of α you trained for worth more.