The Problem

Training speculative‑decoding draft models (MTP, DFlash, DFly, etc.) requires tightly coupled data pipelines, custom loss functions, and a disaggregated inference‑training architecture. Existing open‑source frameworks either support a single draft family or force heavyweight external orchestration, making rapid experimentation and production deployment cumbersome.

What This Does

AngelSpec delivers a torch‑native training stack that unifies six draft architectures under a single configurable pipeline. Core code lives in angelspec/:

  • Model definitions – angelspec/models/draft/*.py (e.g., dflash.py, mtp.py) expose a common forward that the trainer calls.
  • Training orchestration – angelspec/training/*.py (e.g., dflash_trainer.py, mtp_trainer.py) implement the on‑policy TTT rollout, sequence packing, and loss composition.
  • Controller layer – angelspec/controller/setup.py is the package entry point; it parses CLI configs, builds the model, and dispatches to the appropriate trainer.

Configuration files (configs/*.yaml, angelspec/config/*.py) drive every switch: model family, learning‑rate schedule, and packing policy are changed without code edits.

How It Is Wired

Execution starts at the setup entry point angelspec/controller/setup.py. The script:

  1. Loads a config (e.g., configs/qwen3-8b-dfly.yaml).
  2. Calls angelspec.controller.training_controller.validate_packing_candidates (40 functions, 2 classes) to sanity‑check the packing layout.
  3. Instantiates the selected trainer (angelspec/training/dflash_trainer.DFlashTrainer or .../mtp_trainer.MTPTrainer). Each trainer’s __init__ (≈20 functions) invokes: _init_target_lm_head → from_pretrained (9 call sites) to materialize the draft model. _split_hidden_states and _forward to run the forward pass; these call utility functions like exists (20 call sites) and padding (9 call sites).

The forward path reaches the hub module angelspec/utils/logging (imported by 47 modules, instability 0.04). All trainers and the inference engine funnel logging through this file, so changes here have the widest blast radius.

Model‑specific code (angelspec/models/draft/llama3_eagle.py) defines low‑level helpers (_compile_and_cache_with_disk, _patch_cutlass_compilation) that are invoked by the trainer’s weight‑initialisation step. The trainer then writes checkpoints via angelspec/training/checkpoint.py (file I/O) and optionally launches an online evaluation loop (angelspec/controller/online_eval.py) which spawns a subprocess (subprocess.Popen) to run a serving engine.

A circular import exists among angelspec/controller/__init__.py, online_eval.py, and loop.py. This adds maintenance friction: any modification to shared symbols forces careful import ordering or deferred imports.

The call graph shows a handful of high‑traffic functions (exists, backward, _make_dflash_model) called from many locations, indicating they are critical for correctness and performance.

How To Use It

# Clone and install (editable, with vLLM backend)
git clone https://github.com/moses-y/AngelSpec
cd AngelSpec
pip install -e ".[vllm]"          # pulls torch, mooncake-transfer-engine, etc.
pip install mooncake-transfer-engine

Configuration – edit a YAML under configs/ (e.g., configs/qwen3-8b-dfly.yaml) or override on the command line:

./examples/qwen3-8b-dfly/run.sh training.learning_rate=5e-5 training.num_train_steps=500

Run – the example script ultimately executes python -m angelspec.controller.setup --config configs/qwen3-8b-dfly.yaml. The controller builds the model, launches the trainer, and (optionally) starts angelspec/inference/engine/vllm_engine.py for concurrent inference.

Real‑World Use

A cloud‑native service can launch two GPU worker groups: one runs angelspec/controller/setup.py in training mode, the other runs angelspec/inference/engine/vllm_engine.py as a Mooncake‑backed inference server. As checkpoints are written, online_eval.py streams them to the inference server, which reports acceptance length metrics used to adapt the loss weighting in the next training step.

Code Health & Issues

  • High – lockfile missing – pyproject.toml declares dependencies but no poetry.lock/requirements.txt lockfile.
  • High – CI never runs tests – .github/workflows/lint.yml lacks a test step despite 8 test files.
  • Medium – broad exception handling – 14 places (e.g., angelspec/utils/logging.py) swallow all exceptions.
  • Medium – hub module – angelspec/utils/logging.py is imported by 47 modules; instability is low but any change propagates widely.
  • Medium – deep nesting – up to 6‑level indentation in angelspec/models/draft/dflash.py and others, reducing readability.
  • Medium – import cycle – angelspec/controller/__init__.py, online_eval.py, loop.py form a circular dependency.
  • Medium – missing least‑privilege GITHUB_TOKEN permissions – workflow lint.yml declares no permissions.
  • Low – no job timeout – workflow jobs run with default six‑hour limit.

No license issues, secrets, or Dockerfile are present.

The Bottom Line

AngelSpec offers a comprehensive, torch‑native framework for training a range of speculative‑decoding drafts, with clear config‑driven switching and built‑in online evaluation. However, the codebase suffers from maintenance hotspots (a logging hub, circular imports) and lacks reproducible builds and CI test enforcement. Teams ready to invest in stabilising the CI pipeline and refactoring the high‑impact utilities will find a solid foundation for research‑grade speculative decoding.