蚂蚁 inclusionAI 开源 ASpec:面向投机解码 Draft 模型的高吞吐训练框架
inclusionAI/ASpec
蚂蚁 inclusionAI 在 GitHub 开源 ASpec,一个面向投机解码 Draft 模型的高吞吐训练框架,基于 TorchSpec 代码库构建。
中文概览 · Documentation · Roadmap · Quickstart
Faster draft training. More room to experiment.
ASpec is a high-throughput training framework for speculative decoding draft models. It optimizes Draft forward, backward and optimizer execution, overlaps training with Target feature generation, and shares compatible Target features across independent experiments.
ASpec is developed from the TorchSpec codebase, building on its training infrastructure with further execution and workflow optimizations.
Use one workflow to explore algorithms and model structures, train at long context, compare independent experiments, and take a checkpoint through evaluation and export.
Why ASpec?
- Speed up Draft training itself. Optimized attention, fused kernels and selective head/loss computation reduce the work and memory traffic in each forward and backward pass.
- Keep the training pipeline moving. Overlap Target feature generation with Draft training, move features through a GPU tensor data path, and schedule training by useful work.
- One Target, multiple concurrent Draft experiments. Compare algorithms, structures or hyperparameters while generating compatible Target features once. Each Draft keeps its own weights, optimizer, progress and checkpoints, without a shared optimizer-step barrier. See shared Target training for the single-node workflow.
- Make long-context training practical. Combine sequence packing, context parallelism and optimized attention execution to reduce padding waste and distribute long sequences across GPUs.
- Explore without rebuilding the workflow. Switch between algorithm and model recipes with consistent configuration, reporting, recovery, evaluation and export.
Choose your Draft
Choose the algorithm and model structure for your workload. Draft and Target models can have different architectures:
| Dimension | Implemented choices |
|---|---|
| Draft algorithm | DSpark, DFlash, EAGLE3, EAGLE4, native MTP; experimental DFlash2 |
| FFN and routing | Dense, Qwen MoE, Ling MoE |
| Attention structure | MHA/GQA, materialized MLA, sliding-window attention (SWA), CSA, QSA, DSA |
| Execution | Sequence packing, context parallelism, shared Target features, static multi-node training |
See algorithm details and structure stages for the supported combinations, training phases and export formats. DSpark/DFlash with Qwen and Ling are the main optimization and recipe priorities; other implemented paths have explicit support boundaries, and DFlash2 remains experimental.
Scale from a first multi-GPU run to static multi-node training, with NVLink, TCP and RDMA data paths. Check supported paths and limitations, then browse the recipes for starting points.
Start here
- Run source setup; containers are also available.
- Choose Qwen3-8B or Ling tiny for the same two-GPU Plan, train, report, resume and export workflow.
- Choose a training recipe and prepare training data.
Recipes cover Qwen3-8B, Qwen3.5-9B Native MTP, Qwen3.8-27B long context, and Ling tiny/flash 3.0/flash 3.1.
From the repository root:
bash setup.sh source .venv/bin/activate CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output outputs/ling-tiny-001 # Or choose Qwen3-8B: CUDA_VISIBLE_DEVICES=0,1 aspec quickstart qwen3-8b --output outputs/qwen3-8b-001
Both profiles use bundled raw conversations, their own tokenizer/template and an eight-step
train/resume/native-export check. Add --target /absolute/path/to/model to use a local checkpoint.
For training with Target-generated data and held-out evaluation, continue with the
Qwen3-8B DSpark recipe.
Performance
ASpec improves both Draft execution and the complete online training pipeline. The DSpark comparison measures Qwen3-8B and Qwen3.8-27B on Hopper GPUs at short and long context. A representative Qwen3-8B / 4K result:
| Measurement | ASpec input tokens/s | / TorchSpec | / SpecForge | / AngelSpec |
|---|---|---|---|---|
| Draft-only, 1 GPU | 15,256.44 | 1.96× | 2.07× | 2.48× |
| Online steady window, 8 GPUs | 77,721.37 | 2.28× | 2.50× | 2.68× |
The comparison uses one original document per Draft rank and microstep (MBS1), aligned per-document anchor budgets and logical batch sizes, normal interconnects and sufficient buffers. Draft-only excludes Target execution, transport and feature loading; the online steady window includes them after warmup update eight. These are training input rates, not inference speedups or time to a given model quality. The full matrix reports software stacks, GPU splits and complete-loop rates separately.
A separate fixed-sequence Qwen3-8B / CP1 study compares training quality: BSHD validation CE stays close to TorchSpec; measured mean accepted length (tokens per verification round, including the Target bonus token) is 3.665717 for TorchSpec, 3.691665 for ASpec BSHD and 3.746300 for ASpec THD. With deterministic inference enabled, all three match Target-only greedy output token IDs on 256 questions across three repeats. These findings cover that training sequence and evaluation configuration.
For your own training run, see DSpark performance settings for cuDNN attention and direct-prefill options. Resource splits depend on the workload.
See benchmark results and methodology for measured workloads and reproducible comparisons. To assess the drafts you train, use the evaluation and export guide for held-out metrics, acceptance and downstream validation.
Learn more
- Configuration · Training and recovery
- Evaluation and export · Troubleshooting
- Contributing · Security
- Architecture · Project roadmap and status
- Supported paths · Contributor validation
Citation
If you find ASpec useful, please consider citing it:
@misc{aspec2026, author = {Zhou, Aihui and Long, Lingkun and Yu, Chaofan and Yang, Zhe and Peng, Cunyin and Zhang, Zhiqiang and Zhou, Jun}, title = {{ASpec}: A High-Throughput Training Framework for Speculative Decoding}, year = {2026}, url = {https://github.com/inclusionAI/ASpec} }
Citation metadata is available in CITATION.cff.
License
ASpec is licensed under Apache 2.0. Third-party material retains its own licenses and notices; see the source headers and license texts. Model and dataset licenses are separate.
Acknowledgements
We thank TorchSpec, AngelSpec, and SpecForge for their contributions to the speculative decoding training ecosystem.
来源:蚂蚁 inclusionAI:GitHub 新仓库 · github.com