Speculative Decoding

Experiments in speculative decoding for LLM inference acceleration on Mixtral-8x7B-Instruct as the target model. Two strategies are implemented and benchmarked across four datasets.

Strategies

Strategy	Draft mechanism	Extra cost per step	Status
Standard spec decoding	Separate Mistral-7B draft model	1 draft pass × k + 1 target pass	✅ Complete
Standard spec decoding + tuned draft	Mistral-7B fine-tuned via QLoRA distillation	1 draft pass × k + 1 target pass	✅ Complete
Medusa	N lightweight MLP heads on target hidden state	2 target passes	🚧 Work in progress

Both strategies implement rejection sampling (Leviathan et al., 2023): the output distribution is provably identical to sampling from the target model alone.

---

Results

Baseline: std/tokens_per_second ≈ 7.21 tok/s (Mixtral-8x7B autoregressive, 4×A40)

Standard speculative decoding — Mistral-7B draft (Untuned vs QLoRA fine-tuned)

Dataset	k	Acceptance rate (Untuned → Tuned)	Effective k (Untuned → Tuned)	Speedup (Untuned → Tuned)	Spec tok/s (Untuned → Tuned)
HumanEval	2	0.394 → 0.406	1.78 → 1.81	1.45× → 1.48×	10.49 → 10.63
HumanEval	4	0.406 → 0.409	2.59 → 2.62	1.41× → 1.41×	10.14 → 10.19
HumanEval	6	0.400 → 0.399	3.32 → 3.36	1.39× → 1.40×	10.04 → 10.08
HumanEval	8	0.400 → 0.403	4.08 → 4.17	1.42× → 1.42×	10.20 → 10.25
GSM8K	2	0.370 → 0.387	1.74 → 1.75	1.43× → 1.44×	10.33 → 10.40
GSM8K	4	0.377 → 0.375	2.49 → 2.49	1.35× → 1.35×	9.75 → 9.71
GSM8K	6	0.348 → 0.352	3.05 → 3.04	1.27× → 1.26×	9.20 → 9.08
GSM8K	8	0.317 → 0.317	3.49 → 3.52	1.20× → 1.20×	8.65 → 8.65
Wikitext	2	0.379 → 0.387	1.76 → 1.77	1.44× → 1.45×	10.39 → 10.45
Wikitext	4	0.363 → 0.364	2.44 → 2.46	1.31× → 1.32×	9.46 → 9.49
Wikitext	6	0.334 → 0.321	2.99 → 2.92	1.23× → 1.20×	8.90 → 8.66
Wikitext	8	0.301 → 0.281	3.39 → 3.25	1.15× → 1.10×	8.30 → 7.92
MBPP	2	0.363 → 0.379	1.72 → 1.75	1.42× → 1.45×	10.24 → 10.45
MBPP	4	0.373 → 0.380	2.48 → 2.51	1.35× → 1.36×	9.71 → 9.84
MBPP	6	0.354 → 0.351	3.08 → 3.08	1.29× → 1.29×	9.34 → 9.29
MBPP	8	0.332 → 0.338	3.59 → 3.67	1.24× → 1.26×	8.97 → 9.12

Key observations

• k=2 is the sweet spot across all datasets: highest speedup with low overhead cost. Beyond k=4 the acceptance rate drops faster than the effective-k grows.
• QLoRA fine-tuning gives marginal gains (~+0.01 acceptance rate, ~+0.03× speedup). The draft model was already reasonably aligned with the target; distillation has limited headroom here.
• Code tasks (HumanEval, MBPP) accept better than open-domain text (Wikitext): Mixtral's code distribution is well-approximated by Mistral-7B, which shares the same pretraining data mix.
• Acceptance rate is stable across k for HumanEval/MBPP, but degrades with larger k for Wikitext/GSM8K — likely because longer draft chains drift further from the target distribution on structured outputs.

---

Project Structure

speculative_decoding/
│
├── core/
│   ├── sampler.py          # SpeculativeDecoder, StandardDecoder, GenerationMetrics
│   ├── data.py             # Dataset loaders (HumanEval, GSM8K, Wikitext, MBPP)
│   └── benchmark.py        # Sweep runner: iterates (dataset, k), writes JSON, logs W&B
│
├── draft_tuning/
│   └── draft_tuning.py     # QLoRA fine-tuning of draft model via KL + CE distillation
│
├── medusa/                 # 🚧 Work in progress
│   ├── medusa_head.py      # MedusaHeads: N independent 2-layer MLP heads + tied lm_head
│   ├── medusa_sampler.py   # MedusaDecoder: propose (heads) → verify (target) → reject
│   └── medusa_tuning.py    # Training loop for Medusa heads on frozen target
│
├── scripts/
│   ├── run.py              # CLI: standard spec decoding benchmark (+ optional LoRA adapter)
│   ├── run_medusa.py       # CLI: Medusa spec decoding benchmark
│   └── log_medusa_wandb.py # Utility: batch-log result JSONs to W&B
│
└── slurm/
    ├── run_speculative.slurm        # Baseline spec decoding (no tuning)
    ├── run_speculative_tuned.slurm  # Spec decoding with QLoRA draft adapter
    ├── run_draft_tuning.slurm       # Draft model distillation job
    └── run_medusa.slurm             # Medusa training + benchmark (single job)

---

Model Pair

Role	Model	Params	Vocab
Draft	mistralai/Mistral-7B-v0.1	7B	32,000
Target	mistralai/Mixtral-8x7B-Instruct-v0.1	47B	32,000

⚠️ Vocab must match. Mistral-7B-v0.3 uses vocab 32,768 — do not pair it with Mixtral-8x7B-v0.1.

---

Hardware Layout (4 × A40, 48 GiB each)

GPU 0  →  Mistral-7B draft model  (~14 GiB bf16)
GPU 1  ┐
GPU 2  ├→ Mixtral-8x7B target     (~94 GiB bf16, sharded across 3 GPUs)
GPU 3  ┘

For Medusa (no separate draft), all 4 GPUs are available to the target model.

---

Datasets

All benchmarks run on held-out test splits, disjoint from any training data used for draft tuning or Medusa head training.

Dataset	Task	Split used
humaneval	Python code generation	openai/openai_humaneval — test
gsm8k	Grade-school math	openai/gsm8k — test
wikitext	Open-domain text continuation	Salesforce/wikitext wikitext-103 — test
mbpp	Python programming problems	google-research-datasets/mbpp sanitized — test

---

Usage

1. Standard speculative decoding

python scripts/run.py \
    --draft_model  mistralai/Mistral-7B-v0.1 \
    --target_model mistralai/Mixtral-8x7B-Instruct-v0.1 \
    --datasets humaneval gsm8k wikitext mbpp \
    --k_values 2 4 6 8 \
    --n_samples 100 \
    --dtype bfloat16 \
    --results_dir results

2. Standard spec decoding with a tuned draft (QLoRA adapter)

# Step A — fine-tune the draft model
python draft_tuning/draft_tuning.py \
    --max_steps 500 \
    --output_dir ./draft_adapter

# Step B — benchmark with the adapter merged in
python scripts/run.py \
    --draft_adapter ./draft_adapter \
    --datasets humaneval gsm8k wikitext mbpp \
    --k_values 2 4 6 8

3. Medusa 🚧

# Step A — train N heads on the frozen target model
python medusa/medusa_tuning.py \
    --n_heads 8 \
    --max_steps 1000 \
    --n_wikitext_samples 5000 \
    --learning_rate 5e-4 \
    --output_dir ./medusa_heads

# Step B — benchmark
python scripts/run_medusa.py \
    --medusa_heads ./medusa_heads \
    --n_heads 8 \
    --k_values 2 4 6 8 \
    --datasets humaneval gsm8k wikitext mbpp

Quick smoke test

python scripts/run.py \
    --datasets mbpp --n_samples 5 --k_values 4 --max_new_tokens 50 --no_wandb

python draft_tuning/draft_tuning.py \
    --max_steps 20 --no_wandb --output_dir /tmp/draft_adapter

python medusa/medusa_tuning.py \
    --max_steps 10 --no_wandb --output_dir /tmp/medusa_heads

---

SLURM Jobs

Submit from mistral_experiments/ (the project root on the cluster):

sbatch speculative_decoding/slurm/run_speculative.slurm
sbatch speculative_decoding/slurm/run_speculative_tuned.slurm
sbatch speculative_decoding/slurm/run_draft_tuning.slurm
sbatch speculative_decoding/slurm/run_medusa.slurm

All jobs write logs to mistral_experiments/logs/.

---

Output Format

Each (dataset, k) combination writes one JSON file:

{
  "spec/acceptance_rate":    0.72,
  "spec/tokens_per_second":  18.4,
  "spec/effective_k":        2.9,
  "std/tokens_per_second":   7.2,
  "speedup":                 2.56,
  "n_prompts":               100
}

Metric	Description
acceptance_rate	Fraction of draft tokens accepted by rejection sampling
effective_k	Average tokens generated per target forward pass
speedup	spec_tps / std_tps — wall-clock speedup vs autoregressive baseline

---

Metrics & Theory

Speedup upper bound for k drafts with acceptance rate α:

E[tokens per step] = (1 - α^(k+1)) / (1 - α)
Cost per step      = 1 draft pass × k  +  1 target pass   [standard]
                   = 2 target passes                        [Medusa]

In practice, Medusa underperforms standard spec decoding at equal k because heads predict tokens independently from the same hidden state h_t, without chain conditioning between them. Acceptance rate is the primary driver of speedup — increasing k beyond 1/(1-α) yields diminishing returns.

---

References

• Leviathan et al. (2023). Fast Inference from Transformers via Speculative Decoding. arXiv:2211.17192
• Cai et al. (2024). Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv:2401.10774
• Chen et al. (2023). Accelerating Large Language Model Decoding with Speculative Sampling. arXiv:2302.01318