Training a Draft Model: Medusa and EAGLE at Home

Nobody will publish a draft model for your fine-tune. That is the reason this lesson exists. A draft is trained to agree with one particular set of weights, and the moment you changed those weights in Part 13 you left the set of models other people build drafts for. By the end of this lesson you will be able to describe the two families of cheap draft, say what training data each needs and why it has to come from your own model, name the tools that exist and what state they were in on 2026-09-09, and budget the memory and the hours on your own track before you start.

Two places to attach a draft

The previous lessons treated the draft as a black box that proposes tokens. There are three ways to build that box, and they differ in what they reuse from the target.

A separate draft model reuses nothing but the tokeniser. It is a small model of the same family, run to completion for k steps and then verified. It needs no training at all if a suitable small sibling already exists, which is why it is the reduced path on every track in this part’s lab. Its weakness is cost: a whole second network’s weights, a second key-value cache, and k full forward passes per iteration.

A Medusa head reuses the target’s final hidden state. Medusa’s authors describe “adding extra decoding heads to predict multiple subsequent tokens in parallel”, with “a tree-based attention mechanism” that “constructs multiple candidate continuations and verifies them simultaneously in each decoding step”. One target forward pass produces the hidden state; the heads read it and each predicts a different future position. There is no autoregression in the draft at all, which makes it very cheap and also caps how well it can do: head number three is guessing three tokens ahead from information about position zero.

An EAGLE-style draft reuses the target’s internal features and adds autoregression over them. EAGLE’s authors argue that “autoregression at the feature (second-to-top-layer) level is more straightforward than at the token level”, and that the obstacle is uncertainty: “inherent uncertainty in feature level autoregression constrains its performance”. Their fix is to condition on “a token sequence advanced by one time step”, which “effectively resolves the uncertainty, enabling precise second-to-top-layer feature prediction”. The result is much smaller than a draft model and much better informed than a set of parallel heads.

Where each kind of draft attaches to the target

cheaper to run, more dependent on the target

1. Verification: one target forward pass over every proposalIdentical for all three families. This is where the output distribution is preserved.
2. Medusa headsRead the target's final hidden state and predict positions t+1, t+2, t+3 in parallel. Tree attention turns the several candidates per position into a set of continuations verified together.
3. EAGLE-style draftReads the second-to-top-layer features and autoregresses over them, conditioned on the token sequence one step ahead. Later versions replace the fixed tree with a context-aware one.
4. Target model layersUnchanged. None of the three families modifies the weights being served, which is why a draft can be added to a fine-tune without retraining it.
5. Separate draft modelA different network entirely, sharing only the tokeniser. No hidden states, no training, no coupling - and the largest cost per drafted token.

A separate draft model touches the target only through the tokeniser. A Medusa head reads the final hidden state and predicts several positions at once, in parallel and without autoregression. An EAGLE-style draft reads the second-to-top-layer features and runs its own small autoregressive loop over them. All three end at the same verification step, which is what keeps every one of them exact.

Two refinements matter for anyone choosing today. EAGLE-2’s authors observed that most methods “use a static draft tree, implicitly assuming that the acceptance rate of draft tokens depends only on their position”, and found that acceptance “is also context-dependent”; their dynamic tree reports speed-up ratios of 3.05 to 4.26 times and an improvement of twenty to forty per cent over EAGLE. EAGLE-3 changes the objective: it abandons feature prediction for direct token prediction and replaces top-layer features with “multi-layer feature fusion via a technique named training-time test”, which the authors say lets “the draft model to fully benefit from scaling up training data”. They report speed-ups of up to 6.5 times and a throughput improvement of about 1.38 times at a batch size of 64 in SGLang. Those figures are reported by the papers on their own hardware and models, not measured here.

The training data is your model’s own output

Whichever family you pick, the data is the same idea and it is not what people expect.

A draft is not being taught to write well. It is being taught to predict what one specific model will write next. So the training targets have to be that model’s own generations, not human text and not another model’s output. If you train a draft on a human-written corpus, it learns the corpus, and then it disagrees with your model exactly where your model is most distinctive, which is where you needed the agreement.

This is what Medusa’s repository calls self-distillation: its data_generation folder holds “code for self-distillation”, which the project offers so that Medusa can be added “to any fine-tuned model without the original training data”. You do not need the data your fine-tune was trained on. You need your fine-tune’s answers to prompts that look like your traffic.

An EAGLE-style draft needs a second thing on top: the target’s own hidden states. It regresses on features, so its training targets are the vectors the target produced, and no serving API exposes those. They come from a local forward pass with output_hidden_states set, which is why the lab’s data preparation has two stages and why the second one is the expensive one.

Preparing the data for a draft

1. 1Collect prompts that look like your trafficFrom the gateway's logs, from the evaluation set you built in Part 10, or from the fine-tuning set you built in Part 13. The distribution of prompts decides where the draft will agree, so a draft trained on chat prompts will not help your code completion.
2. 2Generate answers with the target model itselfThrough the served endpoint, so no weights are loaded twice. This is the self-distillation step, and it is all a Medusa head needs.
3. 3Optionally, run the conversations back through the target locallyWith output_hidden_states enabled, saving the second-to-top layer per token. Only an EAGLE-style draft needs this, and only this stage needs the weights on the training machine.
4. 4Train the head against those targetsThe backbone stays frozen in the cheap recipes, so the gradients touch only the new parameters and the memory bill is much smaller than a fine-tune.
5. 5Serve and measureAcceptance rate and tokens per second against the no-draft baseline, plus the byte-for-byte check from the first lesson. A draft that trained without error and accepts nothing is a real and common outcome.

▤Memory budget

The hidden-state files are the surprise in this pipeline. One file holds a vector per token for every example, at the model’s hidden width, and a few hundred conversations of a thousand tokens each will fill several gigabytes before you have trained anything. Check the size of the directory after the first fifty examples rather than after the last one, and cap the sequence length deliberately.

The tools, and what state they are in

All three projects below move quickly, and none of them is pinned in this course’s tool reference, so the course names them with their repositories and the date they were read rather than with a version. Read them again before you rely on this table.

Project	Licence	What it trains	Entry point	Data it reads
Medusa	Apache-2.0	Medusa heads on a frozen or co-trained backbone	medusa/train/train_legacy.py under torchrun	ShareGPT-shaped conversations; create_data.py converts, data_generation/ self-distils
EAGLE	Apache-2.0	EAGLE and EAGLE-3 drafts	eagle/traineagle3/main.py under deepspeed	Conversations plus the target’s features
SpecForge	MIT	EAGLE3 and several newer methods, aimed at SGLang	specforge train --config <yaml>	One of raw conversations, pre-tokenised prompts, or precomputed hidden states

Three notes that are easy to miss and expensive to discover late.

The EAGLE repository points elsewhere for EAGLE-3. Its README, read on 2026-09-09, says the authors “strongly recommend using SpecForge for out-of-the-box training of EAGLE-3 with SGLang” rather than the repository’s own implementation. It also sets a useful expectation about scale, describing its work as “trainable (within 1-2 days) and testable on 8x RTX 3090 GPUs. So even the GPU poor can afford it” — which is a lower bar than a frontier lab and a considerably higher one than a single desktop.

SpecForge takes everything from one YAML file. Its training guide describes a single entry point, specforge train --config <file>, with typed sections for model, data, training, runtime and others, and dotted overrides on the command line rather than flags. Its data sources are mutually exclusive: data.train_data_path for raw or preformatted conversations, data.prompts_path for pre-tokenised rows, or data.hidden_states_path for precomputed features. That last option is what lets you separate the expensive feature pass from the training loop, and run them on different machines or different days.

Check that your serving engine can load what you trained. This is the trap. vLLM’s speculative decoding page, read on 2026-09-09, lists its method values as draft_model, ngram, suffix, mtp, eagle3 and dflash; Medusa is not among them on that page. SGLang’s speculative decoding page lists its algorithms as UNO, DFLASH, EAGLE, EAGLE3, STANDALONE, NGRAM and NEXTN, with STANDALONE being the separate-draft-model case. And llama.cpp’s server takes a draft model and nothing else. So a Medusa head is the cheapest thing on this page to train and the hardest to serve on the engines this course pins.

!Important

Decide how you will serve the draft before you decide what to train. If you serve with llama.cpp, the only thing you can use is a small draft model of the same family, and no training is needed. If you serve with vLLM or SGLang, EAGLE-3 through SpecForge is the path their own documentation supports. Train a Medusa head because it is the clearest way to understand what a draft head is and the cheapest run that fits on a small machine, which is exactly why this part’s lab uses it on Track N, and not because it is what you will deploy.

What it costs on your machine

The cheap recipes freeze the backbone, so the memory bill is unlike a fine-tune: the backbone carries no gradients and no optimiser state, and only the new parameters do. What makes a draft head larger than intuition suggests is its output projection, which is the model’s hidden width by its vocabulary size, per head. A model with a large vocabulary has expensive heads.

A draft head on a frozen 8B backbone at bfloat16, on a 24 GB machine (estimate)

Frozen backbone (bf16)

16.4 GB

Head parameters, gradients, optimiser state

4 GB

Activations and workspace, short sequences

2 GB

Free

1.6 GB

Total

24 GB

An estimate, not a measurement. The backbone figure is the bfloat16 size recorded in the course model reference; the head reserve covers the new parameters with their gradients and Adam states, and is dominated by an output projection of hidden width by vocabulary size per head, so a model with a larger vocabulary or more heads needs more. Measure yours with one short run before committing to a long one.

The same head with the backbone loaded at four bits, on a 16 GB machine (estimate)

Frozen backbone (4-bit)

5 GB

Head parameters, gradients, optimiser state

4 GB

Activations and workspace, short sequences

2 GB

Free

5 GB

Total

16 GB

The QLoRA trick from Part 13 applies here for the same reason: a frozen backbone only has to produce the right hidden states, so it can be quantised. The caveat is that the head then learns against four-bit features. If you will serve the target at bfloat16, train against bfloat16 features or expect the acceptance rate you measured to move.

Time is harder to estimate than memory, because it depends on how many conversations you generate and how long they are, and generating them is itself a serving job. The table below is the shape of the budget per track; the values are what the validation pass will fill in.

Pending validationTraining a draft for an 8B target: budget per track

Track	Data generation	Hidden states (EAGLE only)	Head training	Serve and measure
S — DGX Spark, 128 GB	pending	pending	pending	pending
X — Ryzen AI Max+, 64–128 GB	pending	not attempted: no documented path	pending	pending
M — Apple silicon, 24 GB and up	pending	not attempted: no documented path	not attempted	pending
N — NVIDIA desktop, 12–16 GB	pending	pending, at reduced sequence length	pending	pending

the four course tracks, as installed · transformers plus the training repository the commit recorded in the lab notebook · Qwen3-8B or your Part 13 fine-tune, bf16 backbone, or 4-bit where memory requires it · 1,024 tokens of context · to be filled in by the validation pass

Estimates and measurements both belong here once they exist. Until then the useful content is the shape: data generation is a serving job you can leave running, the hidden-state pass is the memory-hungry one, and head training on a frozen backbone is the cheapest of the three.

Per track, the honest position on 2026-09-09 is this. Track S has the memory and the CUDA stack for every recipe here, and is the comfortable path. Track N at 12 to 16 GB can train a Medusa-style head against a quantised 8B backbone, which is the reduced but real version of the exercise. Track X can generate the data and serve a draft model, but neither Medusa’s nor SpecForge’s documentation names a ROCm path for training, so the course does not claim one. Track M has no documented training path in any of the three repositories, and its useful contribution is the serving side: a small draft model in llama.cpp, or the prompt cache from the previous lesson.

✱Hardware note

Track M readers should not read this as a dead end. mlx-lm’s generate takes --draft-model and --num-draft-tokens, confirmed in its argument parser on 2026-09-09, and llama.cpp’s server takes a draft GGUF on every platform including Metal. A Mac can run speculative decoding well; what it cannot currently do, according to these projects’ own documentation, is train the head.

Optimise agreement with the deployed target

A draft model’s task is to propose tokens the target’s verification process will accept. Generic benchmark quality is an indirect signal. Use prompts representative of the target service and preserve the target checkpoint, template and sampling configuration used to generate training data.

Split by prompt source before generating traces. Keep an untouched evaluation set with both common workloads and a shifted domain. Compare an existing draft with the trained draft using the same target and speculative method. Record acceptance by prompt category and position, then measure end-to-end latency and memory.

Check the deployment interface before training. A separate autoregressive draft, a hidden-state predictor and extra prediction heads require different serving support and training signals. Training one does not produce an artefact another method can load. Run a miniature export-and-load test first. The training is useful when the saved serving cost across expected requests justifies preparation and maintenance, while preserving the target’s required output behaviour. A better draft loss alone does not establish that economic or operational result.

Recap

A draft can attach to a target in three places, and they trade cost against how much of the target they reuse: a separate small model shares only the tokeniser, Medusa heads read the final hidden state and predict several positions in parallel without autoregression, and an EAGLE-style draft autoregresses over the target’s second-to-top-layer features. Whichever you choose, the training data is your own model’s generations, because the draft’s job is agreement with these weights and not fluency in general; an EAGLE-style draft additionally needs the target’s hidden states, which only a local forward pass can produce. Medusa, EAGLE and SpecForge are the three projects, all permissively licensed, and none pinned in this course. The decision that matters most is not which is best in a paper but which one your serving engine can load, because a head you cannot serve is an afternoon spent on understanding rather than on speed. And the memory bill is small by fine-tuning standards, because the backbone is frozen: on a 16 GB machine with the backbone quantised, a draft head for an 8B target is within reach.

Check your understanding

Question 1. Why must a draft be trained on the target model's own generations rather than on a high-quality human corpus?

• Because human text is copyrighted
• Because the draft's job is to predict what this particular model will write next, so it must be trained to agree with this model rather than to write well
• Because human text has a different tokeniser
• Because the draft would otherwise be too large

Show the answer and why

Answer: Because the draft's job is to predict what this particular model will write next, so it must be trained to agree with this model rather than to write well

Acceptance is agreement with the target, position by position. A draft trained on human text disagrees with your fine-tune exactly where the fine-tune is distinctive, which is where you most needed the agreement. Medusa's repository ships self-distillation code precisely so that a head can be added to a fine-tune whose original training data is unavailable.

Question 2. You trained a Medusa head for your fine-tune and plan to serve it with llama.cpp. What will you find?

• It will work: llama.cpp loads Medusa heads with --spec-draft-model
• llama.cpp's server takes a separate draft model and nothing else, so the head cannot be served there; check the engine's supported methods before choosing what to train
• It will work after converting the head to GGUF
• Medusa heads are engine-independent by design

Show the answer and why

Answer: llama.cpp's server takes a separate draft model and nothing else, so the head cannot be served there; check the engine's supported methods before choosing what to train

This is the trap the lesson is built around. vLLM's method list as read on 2026-09-09 names draft_model, ngram, suffix, mtp, eagle3 and dflash; SGLang names UNO, DFLASH, EAGLE, EAGLE3, STANDALONE, NGRAM and NEXTN; llama.cpp takes a draft model. Decide how you will serve before you decide what to train.

Question 3. Which of these stages needs the target model's weights loaded on the training machine? Select all that apply.

• Generating the answers that make up the training set
• Extracting the second-to-top-layer hidden states for an EAGLE-style draft
• Training a Medusa head against a frozen backbone
• Measuring acceptance rate against a no-draft baseline

Show the answer and why

Answer: Extracting the second-to-top-layer hidden states for an EAGLE-style draft, Training a Medusa head against a frozen backbone

Generation can go through the served endpoint, and measurement talks to servers over HTTP, so neither needs weights on the training machine. The hidden-state pass does, because no serving API exposes internal features. Training a head against a frozen backbone also does, because the head reads the backbone's activations, even though the backbone carries no gradients.

Question 4. True or false: freezing the backbone means a draft head can be trained in far less memory than a LoRA fine-tune of the same model.

• True
• False

Show the answer and why

Answer: False

False, or at least not reliably. Both freeze the base, so both avoid optimiser state for the backbone. But a LoRA adapter is a handful of low-rank matrices, whereas a draft head carries an output projection of hidden width by vocabulary size, per head. For a model with a large vocabulary the head can be the larger of the two. Budget it, and measure with one short run before starting a long one.

Check my answersClear