The Problem

Training speculative‑decoding drafts (e.g., EAGLE‑3) requires a tightly coupled data‑prep, hidden‑state extraction, and small‑model distillation pipeline. Existing research code is often monolithic, undocumented, and hard to integrate into production workflows, leading to duplicated effort and fragile experiments.

What This Does

tiny-speculators implements the full end‑to‑end workflow for a draft model:

  • Data preparation – tiny_speculators/scripts/prepare_data.py tokenizes Qwen‑3 chat data and builds loss masks.
  • Draft vocabulary construction – tiny_speculators/scripts/vocab_mapping.py selects frequent target tokens and writes the mapping; the core logic lives in tiny_speculators/eagle3/vocab.py.
  • Hidden‑state extraction – tiny_speculators/scripts/generate_hidden_states.py runs a frozen verifier with vLLM and stores early/middle/late layer states.
  • Draft training – tiny_speculators/scripts/train_eagle3.py fuses the three hidden‑state streams, runs the three‑step TTT rollout, and distills the verifier’s token distribution into the single‑layer draft (tiny_speculators/eagle3/model.py + attention.py).
  • Export – tiny_speculators/scripts/export_vllm.py converts the checkpoint to the vLLM Speculators format.

All code is pure Python (15 files) and declared in pyproject.toml. Tests (16 files) cover each stage.

How It Is Wired

Execution starts at the main function in tiny_speculators/modal_pipeline.py:96. From there:

  1. main parses CLI args (parse_args) and calls run.
  2. run (in modal_pipeline.py) invokes the pipeline (tiny_speculators/scripts/pipeline.py).
  3. pipeline.py orchestrates the stages: stage_dataset → prepare_data.py → prepare (file I/O). generate_hidden_states → generate_hidden_states.py → subprocess.run (external vLLM call). train → train_eagle3.py → train (calls shuffle_training, generate_hidden_states, save_checkpoint). export_vllm → export_vllm.py → export_vllm (writes checkpoint files).

The internal call graph shows 62 intra‑repo edges; the most‑used functions are validate_qwen3_8b_config (4 callers), train (4 callers), and run (3 callers). The hubs are:

  • tiny_speculators/config.py – imported by 8 modules, no outgoing imports (instability 0).
  • tiny_speculators/eagle3/vocab.py – imported by 4 modules, also a leaf.

No circular imports are present, so refactoring a hub (e.g., config.py) is low‑risk. The widest blast radius belongs to pipeline.py (8 functions, 3 external callers) because it touches the filesystem, runs a subprocess, and drives the entire workflow.

How To Use It

# Clone the repo
git clone https://github.com/moses-y/tiny-speculators
cd tiny-speculators

# Install dependencies (uv is present)
uv sync          # creates a virtualenv and installs from pyproject.toml

# Prepare data (example input path)
python -m tiny_speculators.scripts.prepare_data \
    --input data/raw_chat.json \
    --output data/prepared.pt

# Generate hidden states (requires a GPU with CUDA 13.0)
python -m tiny_speculators.scripts.generate_hidden_states \
    --verifier-model qwen3-8b \
    --prepared data/prepared.pt \
    --outdir hidden_states/

# Train the draft model
python -m tiny_speculators.scripts.train_eagle3 \
    --hidden-dir hidden_states/ \
    --config tiny_speculators/config.py \
    --output models/draft_eagle3.pt

# Export to vLLM format
python -m tiny_speculators.scripts.export_vllm \
    --checkpoint models/draft_eagle3.pt \
    --outdir vllm_speculators/

The repository supplies a README.md with a high‑level diagram but no explicit CLI flags; the commands above follow the function signatures observed in each script.

Real‑World Use

A SaaS provider could embed this pipeline in a nightly data‑processing job: ingest new chat logs, run prepare_data.py → generate_hidden_states.py → train_eagle3.py, then push the exported vLLM checkpoint to a model‑serving fleet. The scripts are modular enough to be called from an orchestration system (e.g., Airflow) by invoking the same module entry points.

Code Health & Issues

  • High – Missing CI – 15 source files, no .github/workflows or other CI config. Fix: add a GitHub Actions workflow that runs uv sync && pytest.
  • Medium – Dependabot – Only pyproject.toml present; no Dependabot/Renovate config. Fix: commit .github/dependabot.yml.
  • Medium – Resource safety – tiny_speculators/scripts/train_eagle3.py opens files without a context manager; replace open(...) with with open(...) as f:.
  • Medium – Cognitive load – tiny_speculators/scripts/pipeline.py has nesting depth 6; refactor into early‑return guard clauses or split into smaller functions.

No lockfile aside from uv.lock; the declared transformers range (>=4.56.1,<5.14.0) is one major version behind the current 5.15.0, which may cause incompatibilities.

The Bottom Line

tiny-speculators delivers a clear, reproducible implementation of the full speculative‑decoding training pipeline, with well‑scoped modules and comprehensive tests. However, the lack of CI, automatic dependency updates, and a few resource‑handling bugs make it a research‑grade codebase that needs modest engineering effort before production deployment. Suitable for teams comfortable with Python‑centric ML pipelines and willing to add basic DevOps scaffolding.