# MTP draft-head self-distillation (FastMTP recipe)
Raise the single-stream speculative-decode throughput of ANY checkpoint that ships an
`mtp.*` draft head by fine-tuning **only that head** (~0.9B params on Qwen3.6-35B-A3B) on the
serving model's own greedy continuations. The verify model is untouched, so decode output
stays **bitwise plain-greedy** — this changes acceptance (speed), never content.
Why it works: the stock head is trained generically; self-distillation trains it on the exact
distribution it is scored on in the engine — "predict what THIS trunk will emit greedily".
Reference: FastMTP (arXiv 2509.18362); measured baseline here: per-depth acceptance
0.80/0.67/0.62/0.46, E[accepted]=1.825 ⇒ tokens/round = 2.825.
**Fair-comparison rule**: a retrained head changes the ARTIFACT, not the engine. Benchmarks
against other engines (e.g. vLLM `--speculative-config '{"method":"mtp",...}'`) must either
use the stock checkpoint on both sides, or the patched checkpoint on both sides. Never mix.
## Pipeline (4 steps, ~1h/cycle on the 4×V100 box)
### 1. Data generation (the engine dumps its own training pairs)
The shard worker taps every last-stage batch-step column when `OSFKB_DUMP_HIDDEN=<path>` is
set: record = `pos u32 | btrow u32 | token u32 | h[hidden] f32` where `h` is the pre-final-norm
residual — exactly the tensor the draft chain seeds from. The file is **append-mode**: run as
many coordinator sessions as you like; the loader segments streams on position resets.
```bash
# workers (dump on the LAST stage only; plain serving — no OSFKB_MTP_SPEC needed)
OSFKB_WGPU_ADAPTER=0 lfm2-shard-worker --model $M --layers 0:20 --listen 127.0.0.1:9300 &
OSFKB_WGPU_ADAPTER=1 OSFKB_DUMP_HIDDEN=/data/mtp-data/dump.bin \
lfm2-shard-worker --model $M --layers 20:40 --listen 127.0.0.1:9301 &
# one request per line of a VARIED prompts file (domains × languages × formats), repeated
# with different files until the dump holds ≥50k positions (v1 failed at 9.2k: memorized,
# held-out d4 collapsed 0.5→0.14).
lfm2-shard-run --model $M --split 0 --workers 127.0.0.1:9300,127.0.0.1:9301 \
--prompts-file prompts_a.txt --max-tokens 500 --mb-k 16
```
Positions ≈ Σ(prompt_len + max_tokens) per line per session (prefill columns are real-text
pairs and count too). 3 sessions × 40 prompts × ~515 ≈ 60k.
### 2. Training (single V100, fp32, ~20-40 min with early stop)
```bash
docker run --rm --gpus device=2 --entrypoint python3 \
-v /data:/data -v $REPO/training/train_mtp.py:/train.py \
qwen36-v100-vllm:v121 /train.py \
--model /data/osfql-models/qwen36-35b-a3b \
--dump /data/mtp-data/dump.bin \
--out /data/mtp-data/mtp_trained.safetensors \
--depths 4 --steps 600 --bs 16 --lr 5e-6
```
Guard rails built in: 10% held-out windows; every 25 steps it prints held-out per-depth
accuracy and `E = Σ_d Π a_i` (the engine acceptance proxy); keeps the BEST-E tensors; stops
after 6 evals without improvement. **Sanity gates**: the step −1 baseline depth-acc must sit
near the engine's measured per-depth acceptance (d1 ≈ 0.7–0.85 — proves the chain replication
is exact); training loss collapsing toward 0 while val E stalls = memorization, feed more data.
V100 note: fp32 only (fp16 NaNs on V100), tensors saved bf16.
### 3. Surgery (reversible checkpoint patch)
```bash
python3 patch_mtp.py --model /data/osfql-models/qwen36-35b-a3b \
--trained /data/mtp-data/mtp_trained.safetensors
# writes model-mtp-trained.safetensors + patches the index (backup: *.pre-mtp-train);
# docker writes root-owned files — chown BEFORE relaunching workers.
# revert = restore the index backup.
```
### 4. Measure (the only judge is the engine)
```bash
# restart workers (loader reads the patched index automatically), then:
bash accept_bench.sh trained 4 # 10 fixed prompts × 256 tok: per-depth + E + tok/s
```
Ship if held-out E gain confirms on the bench (target: E 1.825 → ≥2.1; >108.8 tok/s needs
E ≥ 2.3 at a 30 ms round). If it regresses: restore the index backup, iterate on data volume
(50→200k), lr (5e-6→2e-6), or freeze fc/norms.
## Porting to another model
Engine side is generic: the dump tap, `--prompts-file`, and `MtpEngine` (any arch's head runs
as a one-layer model through the generic plan builder). Model-specific is `MtpHead` in
`train_mtp.py` — it must replicate the checkpoint's draft recurrence EXACTLY:
1. Tensor names: adjust the `mtp.*` map (`fc`, pre-fc norms, `norm`, the layer, shared
embed/lm_head).
2. Architecture of the draft layer: attention flavor (here: doubled q_proj with sigmoid output
gate, partial RoPE, GQA), MoE vs dense MLP, norm convention (`(1+w)` here — the engine
folds +1 at load, the trainer works on RAW weights).
3. The fc merge order (which half multiplies the embedding — check the loader).
4. Validate before any long run: `--steps 50` — step −1 held-out depth-acc MUST match the
engine's measured per-depth acceptance for the stock head. If it doesn't, the replication
is wrong; fix that first (this gate caught the v1 fc-half swap).
## Files
- `train_mtp.py` — data loader (segmented streams), exact chain replication, weighted CE
(β=0.6 geometric), held-out early-stop, best-tensor save.
- `patch_mtp.py` — safetensors surgery (extra shard + patched index, reversible).
- Engine tap: `src/shard.rs` (`OSFKB_DUMP_HIDDEN`); data-gen driver: `src/bin/shard_run.rs`
(`--prompts-file`).