# deepseek-ai/deepspec architecture

> Architecture diagram of the GitHub repository deepseek-ai/deepspec, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/deepseek-ai/deepspec
- Repository: https://github.com/deepseek-ai/deepspec

## Overview

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.

Main entry points: [`deepspec/modeling/dspark`](https://github.com/deepseek-ai/deepspec/tree/main/deepspec/modeling/dspark) (DSpark), [`deepspec/modeling/eagle3`](https://github.com/deepseek-ai/deepspec/tree/main/deepspec/modeling/eagle3) (Eagle3), [`deepspec/trainer/base_trainer.py`](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/base_trainer.py) (Base trainer), [`deepspec/eval/dspark/evaluator.py`](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/dspark/evaluator.py) (DSpark eval)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_grp_pipeline["Pipeline"]
  node_train_py["train.py<br/>entrypoint<br/>[train.py]"]
  node_eval_py["eval.py<br/>entrypoint<br/>[eval.py]"]
  node_train_sh["train.sh<br/>launcher<br/>[train.sh]"]
  node_eval_sh["eval.sh<br/>launcher<br/>[eval.sh]"]
end

subgraph group_grp_data["Data Prep"]
  node_prepare_data["Prep data<br/>pipeline<br/>[prepare_data.sh]"]
  node_download_split["Split prompts<br/>prep step"]
  node_generate_train_data["Gen targets<br/>prep step"]
  node_launch_sglang[("Target server<br/>external service")]
  node_prepare_cache[("Target cache<br/>data product")]
  node_jsonl_data["JSONL data<br/>dataset io<br/>[jsonl_dataset.py]"]
  node_parser["Parser<br/>[parser.py]"]
  node_cache_dataset[("Cache dataset")]
end

subgraph group_grp_models["Models"]
  node_dspark_model["DSpark<br/>draft model"]
  node_eagle3_model["Eagle3<br/>draft model"]
  node_model_common["Common core<br/>model core<br/>[common.py]"]
  node_losses["Losses<br/>training objective<br/>[loss.py]"]
  node_dspark_heads{{"Heads<br/>[markov_head.py]"}}
end

subgraph group_grp_train["Training"]
  node_prefetcher["Prefetcher<br/>throughput<br/>[cuda_prefetcher.py]"]
  node_trainer_base["Base trainer<br/>trainer core<br/>[base_trainer.py]"]
  node_dspark_trainer["DSpark trainer<br/>[dspark_trainer.py]"]
  node_eagle3_trainer["Eagle3 trainer<br/>[eagle3_trainer.py]"]
  node_ckpt_manager[("Checkpoints<br/>artifact mgmt<br/>[ckpt_manager.py]")]
  node_runtime_utils["Runtime utils<br/>support lib"]
end

subgraph group_grp_eval["Evaluation"]
  node_eval_base["Base evaluator<br/>evaluator core<br/>[base_evaluator.py]"]
  node_dspark_eval["DSpark eval<br/>evaluator<br/>[evaluator.py]"]
  node_eagle3_eval["Eagle3 eval<br/>evaluator<br/>[evaluator.py]"]
end

subgraph group_grp_assets["Assets"]
  node_benchmarks[("Benchmarks<br/>eval data")]
  node_configs["Configs<br/>experiment matrix"]
end

node_train_sh -->|"launches"| node_train_py
node_eval_sh -->|"launches"| node_eval_py
node_prepare_data -->|"splits"| node_download_split
node_prepare_data -->|"drives"| node_generate_train_data
node_generate_train_data -->|"queries"| node_launch_sglang
node_generate_train_data -->|"packs"| node_prepare_cache
node_download_split -->|"reads"| node_jsonl_data
node_jsonl_data -->|"parsed by"| node_parser
node_prepare_cache -->|"feeds"| node_cache_dataset
node_cache_dataset -->|"batched by"| node_prefetcher
node_train_py -->|"uses"| node_trainer_base
node_train_py -->|"loads"| node_configs
node_trainer_base -->|"specializes"| node_dspark_trainer
node_trainer_base -->|"specializes"| node_eagle3_trainer
node_dspark_trainer -->|"optimizes"| node_dspark_model
node_eagle3_trainer -->|"optimizes"| node_eagle3_model
node_dspark_model -->|"shares core"| node_model_common
node_eagle3_model -->|"shares core"| node_model_common
node_dspark_model -->|"trains with"| node_losses
node_eagle3_model -->|"trains with"| node_losses
node_dspark_model -->|"includes"| node_dspark_heads
node_dspark_model -->|"checkpointed by"| node_ckpt_manager
node_eagle3_model -->|"checkpointed by"| node_ckpt_manager
node_trainer_base -->|"depends on"| node_runtime_utils
node_eval_py -->|"uses"| node_eval_base
node_eval_py -->|"loads"| node_configs
node_eval_base -->|"specializes"| node_dspark_eval
node_eval_base -->|"specializes"| node_eagle3_eval
node_dspark_eval -->|"runs on"| node_benchmarks
node_eagle3_eval -->|"runs on"| node_benchmarks
node_dspark_eval -->|"compares to"| node_launch_sglang
node_eagle3_eval -->|"compares to"| node_launch_sglang
node_dspark_eval -->|"loads checkpoint"| node_ckpt_manager
node_eagle3_eval -->|"loads checkpoint"| node_ckpt_manager

click node_train_py "https://github.com/deepseek-ai/deepspec/blob/main/train.py"
click node_eval_py "https://github.com/deepseek-ai/deepspec/blob/main/eval.py"
click node_train_sh "https://github.com/deepseek-ai/deepspec/blob/main/scripts/train/train.sh"
click node_eval_sh "https://github.com/deepseek-ai/deepspec/blob/main/scripts/eval/eval.sh"
click node_prepare_data "https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/prepare_data.sh"
click node_download_split "https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/download_and_split.py"
click node_generate_train_data "https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/generate_train_data.py"
click node_launch_sglang "https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/launch_sglang_server.sh"
click node_prepare_cache "https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/prepare_target_cache.py"
click node_jsonl_data "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/jsonl_dataset.py"
click node_parser "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/parser.py"
click node_cache_dataset "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/target_cache_dataset.py"
click node_prefetcher "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/cuda_prefetcher.py"
click node_dspark_model "https://github.com/deepseek-ai/deepspec/tree/main/deepspec/modeling/dspark"
click node_eagle3_model "https://github.com/deepseek-ai/deepspec/tree/main/deepspec/modeling/eagle3"
click node_model_common "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/modeling/dspark/common.py"
click node_losses "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/modeling/dspark/loss.py"
click node_dspark_heads "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/modeling/dspark/markov_head.py"
click node_trainer_base "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/base_trainer.py"
click node_dspark_trainer "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/dspark_trainer.py"
click node_eagle3_trainer "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/eagle3_trainer.py"
click node_ckpt_manager "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/ckpt_manager.py"
click node_runtime_utils "https://github.com/deepseek-ai/deepspec/tree/main/deepspec/utils"
click node_eval_base "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/base_evaluator.py"
click node_dspark_eval "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/dspark/evaluator.py"
click node_eagle3_eval "https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/eagle3/evaluator.py"
click node_benchmarks "https://github.com/deepseek-ai/deepspec/tree/main/eval_datasets"
click node_configs "https://github.com/deepseek-ai/deepspec/tree/main/config"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train_py,node_eval_py,node_train_sh,node_eval_sh toneBlue
class node_prepare_data,node_download_split,node_generate_train_data,node_launch_sglang,node_prepare_cache,node_jsonl_data,node_parser,node_cache_dataset toneAmber
class node_dspark_model,node_eagle3_model,node_model_common,node_losses,node_dspark_heads toneMint
class node_prefetcher,node_trainer_base,node_dspark_trainer,node_eagle3_trainer,node_ckpt_manager,node_runtime_utils toneRose
class node_eval_base,node_dspark_eval,node_eagle3_eval toneIndigo
class node_benchmarks,node_configs toneTeal
```

## Components

### Pipeline

Top-level execution stages and entrypoints

- [train.py](https://github.com/deepseek-ai/deepspec/blob/main/train.py): entrypoint · Main training entrypoint · `train.py`
- [eval.py](https://github.com/deepseek-ai/deepspec/blob/main/eval.py): entrypoint · Main evaluation entrypoint · `eval.py`
- [train.sh](https://github.com/deepseek-ai/deepspec/blob/main/scripts/train/train.sh): launcher · User-facing training launcher · `scripts/train/train.sh`
- [eval.sh](https://github.com/deepseek-ai/deepspec/blob/main/scripts/eval/eval.sh): launcher · User-facing evaluation launcher · `scripts/eval/eval.sh`

### Data Prep

Prompt processing and target-cache generation

- [Prep data](https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/prepare_data.sh): pipeline · Orchestrates prompt split, target generation, and cache prep · `scripts/data/prepare_data.sh`
- [Split prompts](https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/download_and_split.py): prep step · Downloads and partitions raw prompts · `scripts/data/download_and_split.py`
- [Gen targets](https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/generate_train_data.py): prep step · Runs target-model inference to regenerate answers · `scripts/data/generate_train_data.py`
- [Target server](https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/launch_sglang_server.sh): external service · External inference server serving the target model · `scripts/data/launch_sglang_server.sh`
- [Target cache](https://github.com/deepseek-ai/deepspec/blob/main/scripts/data/prepare_target_cache.py): data product · Packed target outputs used by training · `scripts/data/prepare_target_cache.py`
- [JSONL data](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/jsonl_dataset.py): dataset io · Raw prompt and benchmark JSONL handling · `deepspec/data/jsonl_dataset.py`
- [Parser](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/parser.py): parser · Dataset and record parsing utilities · `deepspec/data/parser.py`
- [Cache dataset](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/target_cache_dataset.py): dataset · Training dataset backed by target cache · `deepspec/data/target_cache_dataset.py`

### Models

Draft-model implementations and variants

- [DSpark](https://github.com/deepseek-ai/deepspec/tree/main/deepspec/modeling/dspark): draft model · DSpark draft-model family · `deepspec/modeling/dspark`
- [Eagle3](https://github.com/deepseek-ai/deepspec/tree/main/deepspec/modeling/eagle3): draft model · Eagle3 draft-model family · `deepspec/modeling/eagle3`
- [Common core](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/modeling/dspark/common.py): model core · Shared draft-model logic and algorithm glue · `deepspec/modeling/dspark/common.py`
- [Losses](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/modeling/dspark/loss.py): training objective · Algorithm-specific loss functions · `deepspec/modeling/dspark/loss.py`
- [Heads](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/modeling/dspark/markov_head.py): head · DSpark confidence and Markov-style heads · `deepspec/modeling/dspark/markov_head.py`

### Training

Optimization, checkpointing, and runtime support

- [Prefetcher](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/data/cuda_prefetcher.py): throughput · CUDA prefetching for training batches · `deepspec/data/cuda_prefetcher.py`
- [Base trainer](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/base_trainer.py): trainer core · Shared optimization and execution flow · `deepspec/trainer/base_trainer.py`
- [DSpark trainer](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/dspark_trainer.py): trainer · DSpark-specific training orchestration · `deepspec/trainer/dspark_trainer.py`
- [Eagle3 trainer](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/eagle3_trainer.py): trainer · Eagle3-specific training orchestration · `deepspec/trainer/eagle3_trainer.py`
- [Checkpoints](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/trainer/ckpt_manager.py): artifact mgmt · Writes and manages experiment checkpoints · `deepspec/trainer/ckpt_manager.py`
- [Runtime utils](https://github.com/deepseek-ai/deepspec/tree/main/deepspec/utils): support lib · Distributed, optim, logging, sampling, and IO helpers · `deepspec/utils`

### Evaluation

Benchmark evaluation and speculative-decoding metrics

- [Base evaluator](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/base_evaluator.py): evaluator core · Shared benchmark evaluation flow · `deepspec/eval/base_evaluator.py`
- [DSpark eval](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/dspark/evaluator.py): evaluator · DSpark acceptance and draft-behavior evaluation · `deepspec/eval/dspark/evaluator.py`
- [Eagle3 eval](https://github.com/deepseek-ai/deepspec/blob/main/deepspec/eval/eagle3/evaluator.py): evaluator · Eagle3 benchmark evaluation · `deepspec/eval/eagle3/evaluator.py`

### Assets

Configs and benchmark datasets

- [Benchmarks](https://github.com/deepseek-ai/deepspec/tree/main/eval_datasets): eval data · Local JSONL benchmark suite · `eval_datasets`
- [Configs](https://github.com/deepseek-ai/deepspec/tree/main/config): experiment matrix · Algorithm, backbone, and size configuration set · `config`

## Connections

- train.sh → train.py: launches
- eval.sh → eval.py: launches
- Prep data → Split prompts: splits
- Prep data → Gen targets: drives
- Gen targets → Target server: queries
- Gen targets → Target cache: packs
- Split prompts → JSONL data: reads
- JSONL data → Parser: parsed by
- Target cache → Cache dataset: feeds
- Cache dataset → Prefetcher: batched by
- train.py → Base trainer: uses
- train.py → Configs: loads
- Base trainer → DSpark trainer: specializes
- Base trainer → Eagle3 trainer: specializes
- DSpark trainer → DSpark: optimizes
- Eagle3 trainer → Eagle3: optimizes
- DSpark → Common core: shares core
- Eagle3 → Common core: shares core
- DSpark → Losses: trains with
- Eagle3 → Losses: trains with
- DSpark → Heads: includes
- DSpark → Checkpoints: checkpointed by
- Eagle3 → Checkpoints: checkpointed by
- Base trainer → Runtime utils: depends on
- eval.py → Base evaluator: uses
- eval.py → Configs: loads
- Base evaluator → DSpark eval: specializes
- Base evaluator → Eagle3 eval: specializes
- DSpark eval → Benchmarks: runs on
- Eagle3 eval → Benchmarks: runs on
- DSpark eval → Target server: compares to
- Eagle3 eval → Target server: compares to
- DSpark eval → Checkpoints: loads checkpoint
- Eagle3 eval → Checkpoints: loads checkpoint