DeepSpec is a full-stack codebase for training and evaluating draft models for speculative decoding. It contains data preparation utilities, draft model implementations, training code, and evaluation scripts.
Install the Python dependencies:
python -m pip install -r requirements.txtData preparation additionally requires an inference engine to serve the target model when regenerating answers; see scripts/data/README.md for details.
Run the stages in order — each stage's output feeds the next:
- Data Preparation — download prompts, regenerate target answers, and build the target cache.
- Training — train a draft model against the cached target outputs.
- Evaluation — measure speculative-decoding acceptance on benchmark tasks.
See scripts/data/README.md for the step-by-step data pipeline:
- download and split training data,
- regenerate answers,
- prepare the target cache (storage warning: this can be very large — roughly 38 TB for the default
Qwen/Qwen3-4Bsetting).
bash scripts/train/train.shtrain.sh launches train.py, which spawns one worker per visible GPU. Select the algorithm and target model by pointing config_path at one of the configs under config/ (e.g. config/dspark/dspark_qwen3_4b.py); see the script header for the full list of configs, how to override config_path / target_cache_dir, and how to use --opts to override individual config fields. Checkpoints are written to ~/checkpoints/<project_name>/<exp_name>/step_*.
Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce CUDA_VISIBLE_DEVICES.
train.sh above trains the DSpark draft with offline SFT against the cached target outputs. DSpark-OPD is an optional on-policy fine-tuning stage that runs on top of an SFT-trained draft: the draft rolls out its own block predictions, the target model scores them token-by-token (KL), and the draft is updated with a policy-gradient + confidence loss. This mainly reduces exposure bias in the markov/confidence heads. It is built as a verl 0.7.0 recipe under third_party/verl/recipe/dspark_opd/.
Prerequisites
- A target cache (from Data Preparation), same as SFT.
- An SFT-trained DSpark draft checkpoint to start from (the output of
train.sh). - The verl environment (separate from the SFT env); see docs/opd/env-setup.md.
Point the config at your paths in third_party/verl/recipe/dspark_opd/config/dspark_trainer.yaml:
actor_rollout_ref.model.path— the SFT draft checkpoint to fine-tune,data.dspark.target_cache_path/override_config.dspark_target_cache_path— the target cache dir,override_config.dspark_teacher_path— the target model (e.g.Qwen/Qwen3-4B).
Train
# single GPU, 1 epoch (default)
bash third_party/verl/recipe/dspark_opd/run.sh
# 8-GPU, 1 epoch
NGPUS=8 BATCH=64 EXP=run1 CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
bash third_party/verl/recipe/dspark_opd/run.shTraining length follows standard verl semantics — epochs are the primary control (EPOCHS, default 1 = one full pass over the dataset; total_training_steps is derived as len(dataloader) × EPOCHS). To cap at a fixed step count instead (e.g. a short smoke run), set STEPS — it overrides epochs and may stop mid-epoch:
NGPUS=2 BATCH=16 STEPS=3 EXP=smoke CUDA_VISIBLE_DEVICES=0,1 \
bash third_party/verl/recipe/dspark_opd/run.shrun.sh knobs: NGPUS, BATCH (data.train_batch_size, must be divisible by GPU count), EPOCHS, STEPS (optional cap), SAVE_FREQ, EXP. Dataset size is data.dspark.n_samples in the config (-1 = the full cache, ~1.34M samples — so 1 epoch is long; lower it for a shorter pass). Checkpoints are written to third_party/verl/checkpoints/dspark_opd/<EXP>/global_step_<N>/.
The default data path is DSPARK_HIDDEN_MODE=recompute: the driver dispatches only tokens and each worker recomputes target_hidden_states from the co-resident teacher (fastest, and matches the hidden states used at inference). Set DSPARK_HIDDEN_MODE=cache (worker reads cached hidden) or dispatch (driver ships cached hidden) to switch data paths. See docs/opd/worker-side-cache-read-design.md.
Convert + evaluate
The verl checkpoint is FSDP-sharded; convert it to the HF format eval.sh expects, then evaluate as usual:
# 1) convert verl checkpoint -> HF-style draft checkpoint
PYTHONPATH=$(pwd) python scripts/opd/convert_ckpt.py \
--verl-ckpt third_party/verl/checkpoints/dspark_opd/run1/global_step_200 \
--out /path/to/dspark_opd_qwen3_4b/step_200
# 2) evaluate (see the Evaluation section) with draft_name_or_path=/path/to/dspark_opd_qwen3_4b/step_200
bash scripts/eval/eval.shDesign docs: docs/DSpark-OPD.md (master plan) and docs/opd/ (fused-step, data-flow optimizations, tensor contract, env setup).
bash scripts/eval/eval.sheval.sh runs eval.py against a trained draft checkpoint over the speculative-decoding benchmarks in eval_datasets/ (gsm8k, math500, aime25, humaneval, mbpp, livecodebench, mt-bench, alpaca, arena-hard-v2). Set:
target_name_or_path— the target model the draft was trained against (e.g.Qwen/Qwen3-4B),draft_name_or_path— the draft checkpoint, e.g.~/checkpoints/deepspec/dspark_block7_qwen3_4b/step_latest, or one of the Hugging Face repo IDs listed in Released Checkpoints.
The checkpoints below are the ones used for Table 1 in the paper. Each checkpoint was trained on open-perfectblend data generated by its corresponding target model in non-thinking mode, and is the direct output of the corresponding training configuration under config/.
| Algorithm | Qwen/Qwen3-4B |
Qwen/Qwen3-8B |
Qwen/Qwen3-14B |
google/gemma-4-12B-it |
|---|---|---|---|---|
| Eagle3 | deepseek-ai/eagle3_qwen3_4b_ttt7 | deepseek-ai/eagle3_qwen3_8b_ttt7 | deepseek-ai/eagle3_qwen3_14b_ttt7 | deepseek-ai/eagle3_gemma4_12b_ttt7 |
| DFlash | deepseek-ai/dflash_qwen3_4b_block7 | deepseek-ai/dflash_qwen3_8b_block7 | deepseek-ai/dflash_qwen3_14b_block7 | deepseek-ai/dflash_gemma4_12b_block7 |
| DSpark | deepseek-ai/dspark_qwen3_4b_block7 | deepseek-ai/dspark_qwen3_8b_block7 | deepseek-ai/dspark_qwen3_14b_block7 | deepseek-ai/dspark_gemma4_12b_block7 |
Important
If you cite these results in a new paper, align your setup with the training settings in this repository; otherwise, the comparison is not meaningful. For domain-specific use, fine-tune the draft model again for better results, especially if the target model is expected to run in thinking mode.
Currently, DeepSpec includes three draft models: DSpark, DFlash and Eagle3.
DeepSpec is released under the MIT License. It includes code adapted from third-party projects under their own licenses; see NOTICE for the full attribution.
DeepSpec builds on the ideas and code of several excellent open-source projects:
- SpecForge (Apache-2.0) — the overall training framework and Eagle3 implementation; portions of the Eagle3 modeling, loss, optimizer, attention, and evaluation code are adapted from it. Adapted files carry an in-file attribution comment, and the full notice is recorded in NOTICE.
- DFlash (MIT) — the DFlash draft-model design and training recipe.
- Qwen3 and Gemma — the target model families supported in this repo.
We thank the authors and maintainers of these projects. Contributions of new algorithms are welcome.