deepseek-ai/deepspec

1. What this repo is

DeepSpec is a Python-based training/evaluation stack for speculative decoding draft models. The repository is organized around a clear pipeline: prepare target-model data, train a draft model against that cache, then evaluate the draft’s acceptance behavior on benchmarks.

2. Top-level execution entrypoints

The main runtime entrypoints are train.py and eval.py. The shell scripts under scripts/train/ and scripts/eval/ are the intended user-facing launchers, while scripts/data/ contains the preprocessing pipeline. This makes the repo split into three operational stages rather than a single monolithic app.

3. Data preparation boundary

Data preparation is intentionally externalized from training. Prompts are downloaded/split, then a separate inference engine serves the target model to regenerate answers, and finally those outputs are packed into a target cache. The scripts/data/launch_sglang_server.sh file and the README notes make it clear that serving the target model is a prerequisite here; this is not done inside the trainer.

4. Training data and cache layer

The deepspec/data/ package is the data boundary between raw JSONL prompts and GPU-ready training batches. It includes:

- JSONL dataset loading,

- parsing utilities,

- a target-cache dataset abstraction,

- CUDA prefetching for throughput.

This layer is what lets training consume precomputed target outputs instead of repeatedly querying the target model.

5. Model families and algorithm split

The codebase supports three draft-model algorithms: DSpark, DFlash, and Eagle3. Each has its own modeling subtree under deepspec/modeling/ and its own config set under config/. The modeling layout is further split by backbone family:

- qwen3/

- gemma4/

That means architecture is parameterized by both algorithm and target backbone.

6. Model internals

The modeling packages contain the actual draft-model implementations:

- algorithm-common code in common.py,

- model definitions in modeling.py,

- training losses in loss.py,

- algorithm-specific heads such as markov_head.py for DSpark and its confidence-related components.

This suggests the repo is not just wrappers around Hugging Face models; it implements custom heads and losses for speculative draft behavior.

7. Trainer subsystem

Training orchestration lives in deepspec/trainer/. There is a shared base_trainer.py plus specialized trainers for dspark and eagle3, along with checkpoint management. The trainer layer is responsible for distributed execution, optimization, logging, and writing checkpoints to the configured experiment directory.

8. Distributed/runtime assumptions

The utilities under deepspec/utils/ show the runtime assumptions: distributed helpers, optimizer setup, sampling, metrics, I/O, logging, and a special suspend/resume helper. The README states the default configs assume a single node with 8 GPUs, so the training architecture is designed around multi-GPU local distributed training rather than a generic cluster scheduler.

9. Configuration architecture

Config files under config/ encode the full experiment matrix: algorithm × target backbone × model size. For example, there are separate configs for dspark_qwen3_4b, dflash_gemma4_12b, and eagle3_qwen3_14b. These configs are the main control plane for selecting model family, draft algorithm, cache location, and training hyperparameters.

10. Evaluation subsystem

Evaluation lives under deepspec/eval/ with a base_evaluator.py and algorithm-specific evaluators for DSpark and Eagle3. The evaluation flow runs the trained draft checkpoint against benchmark datasets in eval_datasets/ and measures speculative-decoding behavior such as acceptance rates and task-specific outcomes.

11. Benchmark and dataset boundary

eval_datasets/ contains JSONL benchmark inputs for tasks like GSM8K, MATH500, HumanEval, MBPP, LiveCodeBench, MT-Bench, Alpaca, and Arena-Hard. There are also helper scripts to convert and download some of these datasets. This repository therefore treats evaluation data as a curated local asset rather than fetching benchmarks dynamically at runtime.

12. Release and checkpoint lifecycle

Checkpoints are written under ~/checkpoints/<project_name>/<exp_name>/step_*, and the README points to released Hugging Face checkpoints for the paper results. So the lifecycle is: train locally, save checkpoints, optionally publish them externally, and then evaluate either local or released artifacts.

13. External dependencies that matter architecturally

The system depends on PyTorch-style GPU training, an external inference server for target-answer regeneration, and model families from Qwen3 and Gemma4. These dependencies shape the architecture: DeepSpec is a draft-model framework built around existing target LLMs, not a standalone foundation-model training stack.

14. Overall data flow

The critical flow is:

raw prompts → target-model regeneration → target cache → draft-model training → checkpoint → benchmark evaluation against the target model.

Any diagram of this repo should show the target model as an external service in the data-prep and evaluation phases, and the draft model as the artifact being trained and measured inside DeepSpec.