DeepSpec 🚀

DeepSpec is a full-stack codebase for training and evaluating speculative decoding algorithms (DSpark-style and others). It provides dataset utilities, modeling components, training scripts, and evaluation tooling used to develop and benchmark speculative decoding and related heads/algorithms.

Table of Contents

• Key features
• Stack
• Repository layout
• Quickstart — run training and evaluation
• Configuration and examples
• How it fits together
• Datasets and formats
• Extending / Adding models and evaluators
• Maintainers
• PR checklist & contributing notes
• Cite us
• Contributing
• License

Key features

• Modular dataset loading and caching (JSONL support, prefetchers).
• Modeling building blocks for speculative decoding heads (confidence, Markov, loss functions).
• Training script and evaluation pipeline for end-to-end experiments.
• Config-driven experiments (per-model/per-dataset config files under config/).

Stack

• Language(s): Python (primary)
• Runtime: CPython (GPU support via PyTorch expected)
• Notable libraries: PyTorch (modeling/training), Hugging Face tokenizers / transformers (likely used in configs), numpy / tqdm for data and utilities.

Repository layout

README.md                    # (this file)
LICENSE
NOTICE
DSpark_paper.pdf             # paper describing DSpark/approach
requirements.txt             # Python deps
train.py                     # training entrypoint
eval.py                      # evaluation entrypoint
config/                      # experiment configs by model / dataset
  dspark/
  dflash/
  eagle3/
deepspec/                    # core python package
  __init__.py
  data/                      # dataset loaders, parser, prefetchers
    jsonl_dataset.py
    parser.py
    target_cache_dataset.py
    cuda_prefetcher.py
  eval/                      # evaluators and evaluation utilities
    base_evaluator.py
    dspark/
      evaluator.py
      draft_ops.py
      confidence_head.py
  modeling/                  # modeling components, heads, and loss
    dspark/
      common.py
      loss.py
      markov_head.py
      ... (gemma4, qwen3 subpackages)
  trainer/                   # training orchestration (hooks, loops)
  utils/                     # misc helpers
scripts/                     # utility scripts
eval_datasets/               # dataset-specific evaluation assets

Quickstart

Prerequisites

• Python 3.8+ (virtualenv recommended)
• pip
• (Optional) CUDA-enabled GPU and matching PyTorch

Install dependencies:

git clone https://github.com/chang5-ctrl/DeepSpec.git
cd DeepSpec
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Run training (example):

# Example: run a training job using a config under config/dspark
python train.py --config config/dspark/dspark_qwen3_14b.py \
  --data_dir /path/to/data \
  --output_dir /path/to/output \
  --num_epochs 1

Run evaluation (example):

python eval.py --config config/dspark/dspark_qwen3_14b.py \
  --checkpoint /path/to/output/checkpoint.pt \
  --data_dir /path/to/eval_data \
  --output /path/to/eval_results.json

Notes:

• The repository provides config/ files for concrete model/dataset combinations (see config/dspark/*).
• Environment variables you may want to set:
  • DATA_DIR or --data_dir for dataset root
  • OUTPUT_DIR or --output_dir for checkpoints & logs
  • HF_HOME if you use Hugging Face cache/custom paths

Configuration and examples

• Config files under config/ encode dataset, optimizer, model head, and training/eval hyperparameters.
• Example configs:
  • config/dspark/dspark_qwen3_14b.py
  • config/dspark/dspark_gemma4_12b.py
  • Use these as templates for new experiments and replace dataset and output paths.

How it fits together

• Data flow: deepspec/data/* parses JSONL-style datasets into PyTorch datasets (see jsonl_dataset.py, target_cache_dataset.py) and provides prefetchers for GPU feeding.
• Modeling: deepspec/modeling/dspark/* contains implementations of loss functions and speculative decoding heads (confidence estimators, Markov heads).
• Training: train.py orchestrates the training loop and uses deepspec/trainer utilities (hooks, checkpointing).
• Evaluation: eval.py uses deepspec/eval/base_evaluator.py and the deepspec/eval/dspark evaluators to run end-to-end metrics described in the accompanying DSpark paper.

Datasets and formats

• The codebase expects dataset JSONL files (one example per line) and includes parsing utilities in deepspec/data/parser.py.
• jsonl_dataset.py shows the required fields / expected tokenization step — examine it to prepare custom datasets.

Extending the codebase

• Add a new model head:
  • Create a new module in deepspec/modeling/<family>/your_head.py
  • Implement forward(), any loss computation, and register to the config.
• Add a new evaluator:
  • Extend deepspec/eval/base_evaluator.BaseEvaluator
  • Add dataset-specific evaluation code under deepspec/eval/<family>/

Maintainers

• Primary: chang5-ctrl (GitHub user chang5-ctrl)
• For urgent issues, open an issue and tag the maintainers.

PR checklist & contributing notes

Before opening a PR, please:

• Run unit tests (if present) and confirm linting.
• Add or update docs if you change public behavior.
• Keep changes isolated and add a short description in the PR body.
• For experimental changes, open a draft PR for early feedback.

Cite us

If you use DeepSpec or the DSpark algorithms in your research or product, please cite:

• See DSpark_paper.pdf in the repository for the formal citation.

BibTeX example:

@inproceedings{DeepSpec2024,
  title = {DSpark: Speculative Decoding ...},
  author = {Authors},
  year = {2024}
}

Contributing

• Issues & PRs welcome — please open issues for bugs or feature requests.
• Follow the existing code style; tests and clear commit messages simplify reviews.
• For substantial changes, open a draft PR and mention maintainers for early feedback.

Troubleshooting & tips

• Out of memory during training: reduce batch size or sequence length, use gradient accumulation.
• Slow data loading: check cuda_prefetcher.py, and ensure dataset caching using target_cache_dataset.py is enabled.
• Reproducibility: commit config file and record requirements.txt / environment info.

License

DeepSpec is provided under the terms in the LICENSE file in this repository.