Skip to content

Latest commit

 

History

17 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 

Repository files navigation

train-dspark-draft-models

Train and evaluate a DSpark draft model for Llama-3.2-1B-Instruct, using the speculators library and vLLM.

Speculative decoding speeds up inference by letting a small drafter propose several tokens at once, then having the full model verify the whole block in a single forward pass and keep the longest prefix it would have produced itself. Output is identical to running the verifier alone — the win is fewer verifier forward passes per token, not a different distribution. The metric that matters is acceptance length, the mean number of tokens kept per verification round.

DSpark builds on DFlash: instead of predicting the block autoregressively the way EAGLE-3 does, it predicts the entire block in one forward pass using anchored block diffusion, conditioned on hidden states read from selected verifier layers (--target-layer-ids 2 8 14 here). Pure block-parallel drafting leaves no dependency between tokens inside a block, so acceptance decays toward the block's end; DSpark restores that dependency with a Markov head — a low-rank logit bias conditioned on the previous token — and adds a confidence head that estimates per-position acceptance probability. The pos_0pos_7 decay in the Evaluation table below is exactly this within-block effect.

  • Drafter — 3 layers, ~0.3B params, proposes 8 tokens per cycle.
  • Verifier — the full Llama-3.2-1B-Instruct, checks each block in one forward pass, so output is identical to running it alone.
  • Data — 100,000 Open PerfectBlend prompts, regenerated on-policy by the verifier itself.
  • Result3.148 tokens accepted per verification round, 4.573 on HumanEval.
  • Scope — two notebooks: regenerate the data, train online against a live vLLM verifier, then serve and benchmark.

Contents

What's in here

.
├── notebooks/
│   ├── data-preparation.ipynb
│   ├── [A100] train-dspark-llama-3.2-1B-instruct-online.ipynb
│   ├── [2xT4 kaggle] train-dspark-llama-3.2-1B-instruct-online.ipynb
│   ├── llama-1b-magpie-kaggle/
│   └── qwen-3-0.6b/
├── LICENSE
└── README.md

Everything runs inside the notebooks. Nothing to install; speculators is cloned at run time.

Both training notebooks use the same Open PerfectBlend data and differ in hardware and configuration: the A100 notebook produced the published drafter, while the 2× T4 Kaggle notebook is a smaller free-tier variant (2 layers, block 4, 80,000 samples, --total-seq-len 4096).

notebooks/llama-1b-magpie-kaggle/ holds the earlier Magpie-trained drafter, and notebooks/qwen-3-0.6b/ an earlier drafter for Qwen/Qwen3-0.6B in both offline and online modes — see its own README.

Model

Architecture: 3 Qwen3 layers (hidden size 2048, intermediate size 8192, 32 attention heads over 8 KV heads, sliding-window attention with a 2048-token window), ~0.3B params, bfloat16. Block size 8, draft vocabulary reduced to 32,000 tokens, aux hidden-state layers 2/8/14, confidence head with Markov (rank 256).

Training: 100,000 Open PerfectBlend prompts regenerated by the verifier, 3 epochs, 96/4 train/validation split, AdamW at 3e-4 with a cosine schedule and 4% warmup, loss weights {"ce": 0.1, "tv": 0.9}. Prompts prepared at 1024 tokens and packed into an 8192-token training window, up to 1024 anchors per sample. Trained online on a single A100, with the verifier sharing the same card at --gpu-memory-utilization 0.25.

Usage

vllm serve rasyosef/Llama-3.2-1B-Instruct-DSpark \
  --port 8000 \
  --gpu-memory-utilization 0.8

Query the OpenAI-compatible endpoint at http://localhost:8000/v1. vLLM reads the verifier and the speculators_config block out of the drafter's own config, so neither the verifier nor --speculative-config is passed separately.

The drafter is not a standalone model — it only works paired with its verifier.

Pipeline

Notebook What it does
data-preparation.ipynb Regenerates Open PerfectBlend prompts with the verifier and pushes the JSONL to yosefw/magpie-llama-3.2-1b-instruct
[A100] train-dspark-llama-3.2-1B-instruct-online.ipynb Filters and tokenizes that data, trains the drafter online against a live vLLM verifier, pushes checkpoint_best, then serves and evaluates it

Online training keeps the verifier resident and fetches hidden states per batch instead of caching them to disk — near-zero disk, but the verifier has to stay up for the whole run. On the A100 both jobs share one card: vLLM takes 25% for the verifier and its KV cache, training takes the rest. The 2× T4 notebook instead gives the verifier GPU 0 and trains on GPU 1.

The data step is a one-off; its output is cached on the Hub, so retraining never re-runs generation.

Run it

Prerequisites: vllm>=0.22.0, an A100 (or two smaller GPUs), and a Hugging Face write token.

The A100 notebook reads the token via google.colab.userdata, the Kaggle ones via kaggle_secrets. Elsewhere those imports fail; replace those cells with a direct os.environ["HF_TOKEN"] = ....

Three things to keep in sync:

  • Sample count--limit on response regeneration and --max-samples on prepare_data.py, 100000 in the A100 run.
  • Prompt length — the prompt filter (≤768 tokens) and --seq-length 1024 on prepare_data.py, which leaves at least 256 response tokens per retained sample. This is deliberately not --total-seq-len 8192: that is the packing window several conversations share, not a per-conversation cap.
  • --target-layer-ids2 8 14 here. It must match in the vLLM launch and training cells, or the drafter trains against hidden states the server isn't exporting; nothing cross-checks it.

Budget for a long session: this configuration takes several hours on a single A100, and the 2× T4 variant roughly 10–11 hours on Kaggle's free accelerators.

Evaluation

Measured with evaluate.py throughput in the training notebook, against the drafter served in vLLM, across the nine RedHatAI/speculator_benchmarks subsets.

subset acceptance_length pos_0 pos_1 pos_2 pos_3 pos_4 pos_5 pos_6 pos_7
HumanEval 4.573 83.7% 68.5% 56.2% 46.0% 36.8% 28.3% 21.7% 16.1%
math_reasoning 4.436 83.6% 68.2% 55.0% 43.8% 34.5% 25.8% 19.2% 13.4%
tool_call 3.462 73.3% 55.5% 41.3% 30.1% 20.2% 13.3% 8.1% 4.3%
question 2.646 63.0% 38.9% 24.3% 15.8% 9.7% 6.4% 4.0% 2.5%
writing 2.636 61.7% 38.8% 24.4% 15.8% 9.9% 6.2% 4.1% 2.6%
rag 2.523 64.9% 39.4% 23.9% 13.3% 6.4% 2.8% 1.1% 0.5%
qa 2.258 55.4% 32.1% 18.2% 9.6% 5.2% 2.8% 1.6% 0.9%
summarization 2.182 58.5% 31.0% 16.3% 7.6% 3.2% 1.1% 0.4% 0.1%
translation 2.021 54.6% 28.6% 12.6% 4.3% 1.3% 0.5% 0.2% 0.0%

acceptance_length is the mean tokens committed per verification round, including the bonus token — floor 1.0, ceiling 9.0 at block size 8. pos_N is the percentage of blocks whose slot N survives verification, decaying across the block as intended. The two use different denominators, so pos_N does not sum to acceptance_length.

Weighted over 89,102 verification steps, acceptance length is 3.148 — an upper bound on single-stream speedup, since it does not charge for the drafter's own forward pass.

Acceptance is highest where the verifier's next token is most predictable — code, math, structured tool calls. HumanEval and math_reasoning are far ahead of everything else and hold their lead deep into the block: HumanEval's pos_4 (36.8%) is above summarization's pos_1 (31.0%), and both still accept better than one token in six at pos_7. The prose-like subsets cluster tightly at 2.0–2.6 and fall off sharply after pos_3, where the longer block buys little — translation is under 2% accepted from pos_4 onward. A block of 8 pays for itself on code, math, and tool calls, and mostly idles on prose traffic.

Credits and license

Built on

Trained on a single Colab A100; the earlier drafters on Kaggle's free 2× T4 accelerators.

License

  • This repository — Apache-2.0, see LICENSE, matching speculators.
  • The published drafter weights — subject to the Llama 3.2 Community License, inherited from the verifier they are trained against and can only run with. Meta's Acceptable Use Policy applies, as does the "Built with Llama" attribution requirement.
  • The regenerated dataset — Llama 3.2 outputs, so the same Llama 3.2 terms apply on top of the source dataset's own license.

About

Code to train and evaluate DSpark draft models for speculative decoding using the speculators library and vllm

Topics

Resources

Stars

2 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages