Skip to content

frankandleaf/draft-maker

Repository files navigation

Draft-Adapter

Turn standard Hugging Face LLMs into vLLM-compatible draft models for speculative decoding — with a single command.

Given any GQA decoder model (Llama, Qwen3, Mistral, Gemma2), Draft-Adapter automatically produces a compact draft model via width compression, depth pruning, and on-policy distillation. The output is a standard HF-format model that plugs directly into vLLM's speculative_model parameter.

Default factors (es=0.5, ls=0.75) produce a draft ~10% the size of the original.

Quick Start

# Install
pip install -e .

# Generate a draft model from Qwen3-1.7B
draft-adapter --model Qwen/Qwen3-1.7B --output ./draft_qwen

# With distillation for better acceptance rate
draft-adapter --model Qwen/Qwen3-1.7B --output ./draft_qwen --distill

# Use the draft model with vLLM
python -c "
from vllm import LLM
llm = LLM(model='Qwen/Qwen3-1.7B', speculative_model='./draft_qwen')
"

Pipeline

Original Model (HF)
    │
    ├── [1] inspect     — Architecture detection, target dimension computation
    ├── [2] calibrate   — Swift-SVD incremental covariance aggregation
    ├── [3] compress    — SliceGPT global PCA projection + width slicing
    ├── [4] prune       — ShortGPT Block Influence (BI) layer ranking
    ├── [5] distill     — DistillSpec on-policy top-K KL (optional)
    └── [6] export      — HF-format draft model (config.json + safetensors)

Parameters

Flag Range Description
--es (0, 1] Embed dimension multiplier
--ls (0, 1] Layer count multiplier
--distill flag Enable on-policy distillation
--kl-mode reverse/forward/tvd KL divergence type for distillation

Technical Details

Width Compression (SliceGPT)

We compute a global PCA projection matrix Q from calibration data, then apply it to all weight matrices. Unlike naive dimension slicing, the PCA rotation ensures we delete only the least important principal components.

For weights interacting with the residual stream:

  • Both-side projection (q_proj, o_proj): W' = Q_top.T @ W @ Q_top
  • Input-side (k_proj, v_proj, gate_proj, up_proj): W' = W @ Q_top
  • Output-side (down_proj): W' = Q_top.T @ W
  • Embeddings (embed_tokens, lm_head): W' = W @ Q_top

Depth Pruning (ShortGPT)

Block Influence (BI) measures how much each layer transforms hidden states: BI_i = 1 - mean(cos_sim(X_i, X_{i+1})). Lower BI = more redundant. First and last layers are always protected.

Distillation (DistillSpec)

On-policy distillation: the student (draft) generates tokens, the teacher scores them. Top-K sparse KL divergence prevents the student from wasting capacity on near-zero teacher logits.

Supported Models

Family Model Type Status
Llama 2/3/4 llama Supported
Qwen2 / Qwen2.5 qwen2 Supported
Qwen3 qwen3 Supported
Mistral mistral Supported
Gemma 2 gemma2 Supported
StableLM stablelm Supported
MoE variants Not supported
MLA (DeepSeek) Not supported
Mamba (SSM) Not supported

References

About

一个能设置参数+一键启动的软件包,对标准Transformer模型(Llama/Qwen/Mistral等)进行处理,产出一个显存占用极小、接受率达标、可直接接入vLLM的Draft Model。

Resources

License

Stars

6 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages