Draft-OPD: On-Policy Distillation for Speculative Draft Models

Paper | Models | Project Page | Quick Start | Evaluation | Acknowledgements | Citation

News

• [July 2026] The training data for Draft-OPD has been released and is available at /data.
• [May 2026] Draft-OPD is available on arXiv: 2605.29343.
• [May 2026] Released Draft-OPD model checkpoints are available in the Hugging Face collection.

Introduction

Draft-OPD trains speculative draft models with on-policy target feedback. Instead of only learning from fixed target-generated trajectories, the drafter is supervised on draft-induced states exposed during speculative verification, including the positions where draft proposals are rejected.

This repository contains the public training and evaluation code for Draft-OPD. The main training stack is built on verl and sglang-dflash, while the evaluation utilities live under diffusion/.

Repository Layout

Path	Purpose
verl/	Training code and the public OPD DFlash training entrypoint.
sglang-dflash/	DFlash / SGLang runtime code used by training and evaluation.
diffusion/	Draft-OPD evaluation utilities, with the main benchmark workflow in diffusion/dflash/.

Training Entry Point

This repository provides a single public training entrypoint:

verl/examples/on_policy_distillation_trainer/run_qwen_gsm8k_forward-ins.sh

The script launches DFlash on-policy distillation through verl. It wraps run_qwen_gsm8k.sh, so run it from the repository root after installing the verl training environment.

Install

From the repository root, run:

bash install.sh

This installs the editable sglang-dflash and verl packages and their dependencies. No other manual setup is required.

Quick Start

Set your local model and data paths, then run:

cd /path/to/opd

MAIN_MODEL_PATH=/path/to/main/model \
DRAFT_MODEL_PATH=/path/to/draft/model \
TRAIN_JSONL=/path/to/train.jsonl \
bash verl/examples/on_policy_distillation_trainer/run_qwen_gsm8k_forward-ins.sh \
  "data.val_files=['/path/to/aime24.jsonl','/path/to/gsm8k.jsonl','/path/to/math500.jsonl','/path/to/mbpp.jsonl']"

Required paths:

• MAIN_MODEL_PATH: target/main model. For example, use Qwen/Qwen3-4B.
• DRAFT_MODEL_PATH: initialized draft model used for speculative decoding and training. For example, download z-lab/Qwen3-4B-DFlash-b16 from Hugging Face. Alternatively, you can train your own draft model with SpecForge.
• TRAIN_JSONL: training data in JSONL format.
• data.val_files: validation JSONL files, passed as a Hydra override.

Common optional overrides:

LR=3e-4 \
train_epochs=8 \
STUDENT_WORLD_SIZE=7 \
TEACHER_WORLD_SIZE=1 \
TRAIN_PROMPT_BSZ=21 \
PPO_MICRO_BATCH_SIZE_PER_GPU=1 \
MAX_PROMPT=512 \
MAX_RESPONSE_LENGTH=4096 \
ENABLE_THINKING=False \
bash verl/examples/on_policy_distillation_trainer/run_qwen_gsm8k_forward-ins.sh

Useful DFlash-specific options:

• DFLASH_LM_HEAD_CHUNK_SIZE: LM-head chunk size, default 512.
• RANDOM_RESPONSE_ANCHOR_ENABLED: enable random response-anchor ablation, default False.
• RANDOM_RESPONSE_ANCHOR_SEED: random-anchor seed, default 42.
• TEACHER_GPU_MEMORY_UTILIZATION: teacher inference memory fraction, default 0.1.

Checkpoints are saved under:

verl/checkpoints/verl-dflash-opd/

Use verl/scripts/fsdp_to_dflash.sh from the repository root to extract the draft model from saved actor weights.

Evaluation

Draft-OPD evaluation utilities live under diffusion/, with the main benchmark workflow in diffusion/dflash/.

The diffusion/dflash/ folder is adapted from an early version of DFlash. You can also directly use the DFlash repository to evaluate DFlash draft models.

See diffusion/dflash/README.md for the DFlash evaluation entrypoints and links to the English / Chinese usage guides.

Acknowledgements

We thank DFlash and EAGLE3 for their inspiring work on speculative decoding and draft-model training. We also thank SpecForge, SGLang, and verl for the open-source infrastructure that this repository builds on.

Citation

If you find our work useful, please consider citing our paper:

@misc{lei2026draftopdonpolicydistillationspeculative,
      title={Draft-OPD: On-Policy Distillation for Speculative Draft Models}, 
      author={Haodi Lei and Yafy Li and Haoran Zhang and Shunkai Zhang and Qianjia Cheng and Xiaoye Qu and Ganqu Cui and Bowen Zhou and Ning Ding and Yun Luo and Yu Cheng},
      year={2026},
      eprint={2605.29343},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2605.29343}, 
}