Osprey

Code, checkpoints, and data for Osprey: Target-agnostic Pre-training Makes Stronger Drafters in Speculative Decoding, accepted at EMNLP 2026.

• Paper: https://arxiv.org/abs/2609.09338
• Checkpoints and data: https://huggingface.co/BobbieBieee/Osprey-Speculative-Decoding

Speculative drafters are usually trained from scratch against one target on one data distribution, and their acceptance rate collapses when the workload shifts. Osprey instead bootstraps the drafter from a pretrained small language model: prune it to a shallow backbone, restore its language modeling with target-agnostic next-token pretraining, then adapt that one backbone to each target through vocabulary alignment, zero-initialized QKV expansion, and on-policy EAGLE-3 distillation.

Stage	What it does	Where
1. Prune	Keep the embedding, LM head, and first N blocks of a small LM	not included; the released backbone is already pruned
2. Pre-train	Next-token prediction on FineWeb; one reusable backbone	not included; the 55k-step backbone is on the Hub
3. Convert	Align vocabulary, expand QKV with zero-initialized target taps	scripts/osprey/convert_checkpoint.py
4. Train	On-policy EAGLE-3 distillation against the frozen target	scripts/osprey/run_adaptation.py

This repository implements Stages 3 and 4. The walkthrough below reproduces the Qwen3-8B cross-domain matrix (Figure 2). The two larger targets follow the same steps from the same backbone:

• Llama-3.3-70B-Instruct — Table 1
• MiniMax-M2.5 — Tables 2–3

Built on SpecForge; everything the paper does not use has been removed.

Installation

Python 3.11 and NVIDIA GPUs. torch, transformers, and sglang are pinned in pyproject.toml.

conda create -y -n osprey python=3.11 && conda activate osprey
pip install --upgrade pip && pip install .
MAX_JOBS=8 pip install -v ".[fa]" --no-build-isolation   # optional flash-attn; training uses flex_attention either way
bash sglang_patches/apply_patch.sh                        # required to serve or evaluate an Osprey drafter
PYTHONPATH=$PWD python -m unittest discover -s tests -p "test_*.py"

Stock SGLang serves only single-layer EAGLE-3 drafts and overwrites the draft's embedding with the target's at load time; the patch adds multi-layer drafts and keeps the drafter's own aligned embedding. Training does not need it. If a trained drafter evaluates at an accepted length near 1.0, SGLang is unpatched.

Every command below runs from the repository root with one output root set:

export OSPREY_OUTPUT=/path/with/room     # checkpoints, results, caches

Qwen3-8B

Five single-domain drafters — chat, code, commonsense, finance, math — each evaluated on all five domains.

1. Data

Each domain has 65k training conversations whose assistant turns were regenerated by Qwen3-8B, plus 512 held-out prompts. Records are {"id": str, "conversations": [{"role", "content"}, ...]}.

Option A — download the released splits:

hf download BobbieBieee/Osprey-Speculative-Decoding \
  --include "data/*_qwen3_8B_4096.jsonl" --local-dir osprey
export OSPREY_DATA=$PWD/osprey/data

Option B — regenerate them. Start from prompt-side JSONL in the same schema (data/convert_dataset.py converts a HuggingFace or ShareGPT-style dataset) and let the target write the responses; the script runs a pool of SGLang servers:

export OSPREY_DATA=/path/to/prompts            # holds <domain>_prompts.jsonl
for domain in chat code commonsense finance math; do
  INPUT_FILE="$OSPREY_DATA/${domain}_prompts.jsonl" \
  OUTPUT_FILE="$OSPREY_DATA/${domain}_train_65k_qwen3_8B_4096.jsonl" \
  MODEL=Qwen/Qwen3-8B MAX_TOKENS=4096 \
    bash scripts/osprey/regenerate_data.sh
done

2. Convert

Download the Stage-2 backbone (Qwen3-4B pruned to 2 layers, 55k FineWeb steps) and convert it for the target:

hf download BobbieBieee/Osprey-Speculative-Decoding \
  --include "pretrained/qwen3-4b-2layer-fineweb-55k/*" --local-dir osprey

python scripts/osprey/convert_checkpoint.py \
  --architecture qwen3 --target qwen3-8b \
  --source-checkpoint osprey/pretrained/qwen3-4b-2layer-fineweb-55k \
  --output-dir "$OSPREY_OUTPUT/qwen3-2l-init"

This widens every Q/K/V projection to 2 * hidden_size with the pretrained weights on the hidden-state half and zeros on the new target-feature columns, sets fc to the identity on the deepest tapped target layer, and remaps the embedding and LM head into Qwen3-8B's vocabulary. At initialization the target features contribute nothing, so the drafter computes exactly what the backbone computed. The alignment report should read direct: 151669, single: 0, averaged: 0, unmapped: 0; on the released backbone this reproduces the paper's Stage-3 output bit for bit.

3. Train

export OSPREY_DRAFT_INIT="$OSPREY_OUTPUT/qwen3-2l-init"
bash examples/qwen3-8b/train.sh                # all five domains, one after another

Defaults are the paper's: 2 GPUs with the target served in-process by SGLang at tp=1, lr 1e-4, 3 epochs, sequence length 4096, TTT length 5, qwen3-thinking template. 65k × 3 / 2 GPUs = 97,500 steps per domain. Checkpoints land in $OSPREY_OUTPUT/qwen3-8b/osprey/<domain>/epoch_2_step_97500.

On an 8-GPU node train four domains at once — one launch per domain, its own GPU pair and tokenizer cache; torchrun --standalone picks a free port per job:

gpu=0
for domain in chat code commonsense finance; do
  CUDA_VISIBLE_DEVICES=$gpu,$((gpu + 1)) OSPREY_DOMAINS="$domain" \
  OSPREY_CACHE_DIR="$OSPREY_OUTPUT/cache/$domain" \
    bash examples/qwen3-8b/train.sh 2 > "$OSPREY_OUTPUT/train-$domain.log" 2>&1 &
  gpu=$((gpu + 2))
done
wait

Disk: a run at the default SAVE_INTERVAL=5000 writes 20 checkpoints of ~6 GB. Training also writes tokenizer caches to $OSPREY_CACHE_DIR, HF dataset caches to $HF_DATASETS_CACHE, and DataLoader sockets to $TMPDIR; on a full filesystem the job dies with no error in its log, and clearing $TMPDIR during a run kills it. Point all of them at a volume with room.

4. Evaluate

Either evaluate your own checkpoints or download the released ones:

hf download BobbieBieee/Osprey-Speculative-Decoding \
  --include "qwen3-8b/*" --local-dir osprey

One run scores one drafter on all five domains — one row of the matrix:

bash examples/qwen3-8b/eval.sh osprey/qwen3-8b/osprey-math      # or $OSPREY_OUTPUT/qwen3-8b/osprey/math/epoch_2_step_97500

Rows are independent, so fill the grid in parallel — one GPU and one port each:

gpu=0
for domain in chat code commonsense finance math; do
  CUDA_VISIBLE_DEVICES=$gpu PORT=$((30000 + gpu)) RESULT_NAME="osprey-$domain" \
    bash examples/qwen3-8b/eval.sh osprey/qwen3-8b/osprey-$domain > "$OSPREY_OUTPUT/eval-$domain.log" 2>&1 &
  gpu=$((gpu + 1))
done
wait

Serving configuration is the paper's: (batch, steps, topk, draft_tokens) = (1, 5, 1, 6), max_tokens 4096, greedy, tp=1, all 512 prompts per domain. Accepted length is averaged per request and is deterministic under greedy decoding; a subsample reads high (the first 32 prompts overstate commonsense by 0.44), so quote only full-split numbers. Results go to $OSPREY_OUTPUT/results/<name>/ as JSON plus a flat CSV.

The filled grid reproduces Figure 2 — Osprey averages 3.686 accepted tokens against 3.176 for the from-scratch EAGLE-3 baseline, ahead in all 25 cells and by more out of domain (+0.52) than in domain (+0.49):

To serve a drafter rather than benchmark it, hand it to SGLang directly; the server applies the target's own chat template:

python -m sglang.launch_server --model-path Qwen/Qwen3-8B \
  --speculative-algorithm EAGLE3 --speculative-draft-model-path osprey/qwen3-8b/osprey-chat \
  --speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6 \
  --tp-size 1 --dtype bfloat16 --mem-fraction-static 0.8 --port 30000

Repository layout

scripts/osprey/         convert_checkpoint.py (Stage 3), run_adaptation.py (Stage 4 presets),
                        regenerate_data.sh (on-policy data), align_tokenizer.py
scripts/                train_eagle3.py (training loop) and the weight-surgery converters
specforge/              TTT loss, draft models, SGLang target backend, data pipeline
benchmarks/             bench_eagle3.py plus one adapter per benchmark
configs/                draft configs for the from-scratch EAGLE-3 arms of examples/qwen3-8b/train.sh
data/                   convert_dataset.py: any chat dataset -> the JSONL schema above
examples/               qwen3-8b/, llama33-70b/, minimax-m25/ — train.sh, eval.sh, README
sglang_patches/         the SGLang patch required to serve an Osprey drafter
tests/                  unit tests; test_osprey_release.py guards the invariants

Citation

@inproceedings{bie2026osprey,
  title     = {Osprey: Target-agnostic Pre-training Makes Stronger Drafters in
               Speculative Decoding},
  author    = {Bie, Fengxiang and Jian, Yuqing and Yu, Yifan and Zhou, Zhongzhu and
               Shao, Zelei and Athiwaratkun, Ben and Song, Shuaiwen Leon and
               Xu, Chenfeng and Wu, Xiaoxia and Zhang, Tianyi},
  booktitle = {Proceedings of the 2026 Conference on Empirical Methods in Natural
               Language Processing (EMNLP)},
  year      = {2026},
}

Acknowledgements

Osprey builds directly on SpecForge and SGLang. We thank the SpecForge, SGLang, and EAGLE contributors for the training and serving infrastructure this project depends on.