Skip to content

Latest commit

 

History

4 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

On-Policy Delta Distillation

Byeongho Heo, Jaehui Hwang, Sangdoo Yun, Dongyoon Han
NAVER AI Lab

arXiv

Abstract

On-policy distillation is an alternative post-training method that provides token-level supervision from a teacher model. In this work, we introduce the delta signal, defined as the difference between a reasoning-tuned teacher and its base model before reasoning tuning. This signal captures the changes induced by reasoning tuning and provides a more direct objective for transferring reasoning capabilities. Based on this idea, we propose On-Policy Delta Distillation (OPD²). Across mathematics, science, and code-reasoning benchmarks, OPD² consistently outperforms conventional on-policy distillation methods on multiple Qwen3 model sizes, in both thinking and non-thinking settings, as well as on Gemma-4-E4B.

Paper

Updates

  • Aug 7, 2026: Multilingual extension released (arXiv:2608.05802)
  • Aug 5, 2026: Code released
  • Jul 16, 2026: arXiv paper released

Trainers

Trainer Entry Per-token signal at the sampled token
OPD src/open_r1/opd.py teacher − student log-prob diff
ExOPD src/open_r1/exopd.py λ·(teacher − teacher_base) − (student − teacher_base)
OPD² src/open_r1/opd2.py delta signal (teacher − teacher_base) with expectation correction, gated by the (teacher − student) update direction

The signal is used directly as a dense per-token advantage in a GRPO PPO-clipped loss (built on trl's GRPOTrainer, adapted from open-r1).

Installation

The pinned training environment is defined by the Docker images (recommended):

docker build -f docker/Dockerfile_open-r1_vllm012 -t opd2:qwen3 .   # Qwen3 runs  (torch 2.9,  vLLM 0.12, trl 1.4)
docker build -f docker/Dockerfile_open-r1_gemma4  -t opd2:gemma4 .  # Gemma4 runs (torch 2.11, vLLM 0.23, trl 1.7)

Then, inside the container (or on your own torch + vLLM environment):

git clone https://github.com/bh-heo/opd2.git && cd opd2
pip install --no-deps -e .

Dataset

Training data comes from NVIDIA's open reasoning datasets: nvidia/OpenMathReasoning, nvidia/OpenCodeReasoning, and nvidia/OpenScienceReasoning-2.

For convenience we provide bhheo/nvidia_open_reasoning_balanced_100k — a pre-mixed subset of the three (33,333 conversations each, math : code : science = 1 : 1 : 1) in the sharegpt conversation format the trainers consume. It is referenced by dataset_name in every recipe and downloaded automatically from the Hugging Face Hub on first run — no manual step needed. Each row keeps a source field pointing back to the originating NVIDIA dataset.

Training

All experiments were run on NVIDIA H100 80GB nodes (8 GPUs per node): Qwen3-1.7B on 1 node (8×H100); Qwen3-4B / Qwen3-8B / Gemma4-E4B-it on 4 nodes (1 vLLM rollout server + 3 training nodes = 32×H100).

Recipes are self-contained: dataset, hyperparameters, teacher / teacher_base models and thinking-mode flags are all set per student model.

Student Recipes Teacher (teacher_base) Mode
Qwen3-1.7B recipes/Qwen3-1.7B/{opd,exopd,opd2}/ Qwen3-4B-Instruct-2507 (Qwen3-4B-Base) no-think
Qwen3-4B recipes/Qwen3-4B/{opd,exopd,opd2}/ Qwen3-30B-A3B-Instruct-2507 (Qwen3-30B-A3B-Base) no-think
Qwen3-8B recipes/Qwen3-8B/{opd,exopd,opd2}/ Qwen3-30B-A3B-Instruct-2507 (Qwen3-30B-A3B-Base) no-think
Gemma4-E4B-it recipes/Gemma4-E4B-it/{opd,exopd,opd2}/ gemma-4-31B-it (gemma-4-31B) thinking

For a thinking-mode Qwen3 run, override on the command line: --opd_teacher Qwen/Qwen3-4B-Thinking-2507 --opd_only_think_mode true --opd_no_think_mode false --opd_no_think_teacher false.

Single node (Qwen3-1.7B — teacher fits alongside the student, 8×80GB)

vLLM runs colocated with training; the teacher is sharded with ZeRO-3:

accelerate launch --config_file recipes/accelerate_configs/zero2.yaml \
    --num_processes 8 \
    src/open_r1/opd2.py \
    --config recipes/Qwen3-1.7B/opd2/config_open_nvidia_100k.yaml \
    --use_vllm true --vllm_mode colocate \
    --vllm_tensor_parallel_size 4 --vllm_gpu_memory_utilization 0.1 \
    --vllm_importance_sampling_correction false \
    --output_dir checkpoints/opd2_q17b

Swap opd2.py + the recipe path for opd.py / exopd.py to run the baselines.

Multi node (Qwen3-4B / 8B / Gemma4-E4B — 30B-class teachers)

We use 4 nodes: 1 vLLM rollout server + 3 training nodes (24 GPUs). The recipes' gradient_accumulation_steps assume this topology — adjust it to keep the effective batch (per_device_batch 2 × grad_accum × #GPUs = 256 for Qwen3-1.7B, 288 for the multi-node recipes) if you use a different GPU count.

# node 0 — rollout server
python -m trl.scripts.vllm_serve --model Qwen/Qwen3-8B \
    --tensor_parallel_size 8 --host 0.0.0.0 --port 8000 \
    --gpu_memory_utilization 0.9 --max_model_len 32768 --enforce_eager True

# nodes 1-3 — training (run on each node with RANK=0/1/2)
accelerate launch --config_file recipes/accelerate_configs/zero2.yaml \
    --num_machines 3 --num_processes 24 --machine_rank $RANK \
    --main_process_ip <node1 addr> \
    src/open_r1/opd2.py \
    --config recipes/Qwen3-8B/opd2/config_open_nvidia_100k.yaml \
    --use_vllm true --vllm_mode server \
    --vllm_server_host <node0 addr> --vllm_server_port 8000 \
    --vllm_importance_sampling_correction false \
    --output_dir checkpoints/opd2_q8b

Gemma4-E4B-it runs the same way with the opd2:gemma4 image and recipes/Gemma4-E4B-it/... (the recipe uses attn_implementation: sdpa; kernels-community/flash-attn also works, see the Dockerfile header).

Notes:

  • Any recipe value can be overridden on the command line (--learning_rate, --opd_teacher, ...).
  • OPD² accepts every opd_* flag; its own flags are --opd2_teacher_base (required) and --opd2_rewards_top_k (default 1024).
  • Multi-node generation payloads are broadcast over a gloo (CPU) process group (src/open_r1/trainers/_dist_utils.py) to avoid NCCL transport issues on some fabrics.

Evaluation

Benchmarks are run with evalchemy (vLLM backend, chat template applied, \boxed{} extraction). The paper averages use:

Domain Benchmarks
math AIME24, AIME25, AMC23, HMMT, MATH500, OlympiadBench, RGMath
code LiveCodeBenchv5, CodeForces, CodeContests, RGAlgorithmic
others GPQADiamond, SuperGPQA, SciBench

AIME/AMC/HMMT/GPQADiamond/LiveCodeBenchv5/CodeForces are evalchemy built-ins; the rest are the custom adapters in evaluation/chat_benchmarks/.

Default flow (no-think Qwen3 checkpoint):

TASK=AIME24                              # benchmark to run (see table above)
CKPT=$(pwd)/checkpoints/opd2_q17b        # output_dir of the training run

# 1) build the eval image
docker build -f docker/Dockerfile_eval2_vllm012 -t opd2:eval .

# 2) inside the container: copy the custom adapters into evalchemy (idempotent)
bash evaluation/setup_extended.sh

# 3) patch the checkpoint's chat template to the trained rollout mode
#    (first argument = the student's base HF repo, the chat-template source)
python evaluation/patch_no_think_tokenizer.py Qwen/Qwen3-1.7B "$CKPT"

# 4) run one benchmark (TP = #GPUs)
REPO=$(pwd)
export PYTHONINTMAXSTRDIGITS=0
export PYTHONPATH="$REPO/evaluation/eval_compat:/workspace/evalchemy:$(ls -d /workspace/evalchemy/eval/chat_benchmarks/*/ | tr '\n' ':')"
cd /workspace/evalchemy
python "$REPO/evaluation/extended_eval_setting.py"   # paper n_repeat settings
python -m eval.eval \
    --model vllm \
    --tasks "$TASK" \
    --apply_chat_template \
    --batch_size auto \
    --model_name "$CKPT,dtype=bfloat16,tensor_parallel_size=8,enforce_eager=True,gpu_memory_utilization=0.8" \
    --output_path "$REPO/results/$TASK"

Thinking-mode checkpoints (--opd_only_think_mode runs) — replace step 3 with:

python evaluation/patch_only_think_tokenizer.py Qwen/Qwen3-1.7B "$CKPT"

Gemma4 checkpoints — use the Gemma4 eval image and tokenizer patcher:

docker build -f docker/Dockerfile_eval2_gemma4 -t opd2:eval-gemma4 .   # step 1
python evaluation/patch_gemma_tokenizer.py google/gemma-4-E4B-it "$CKPT"  # step 3

Notes:

  • extended_eval_setting.py applies the per-benchmark repeat counts used in the paper (AIME24/25 ×16, MATH500 ×10, GPQADiamond ×10, most others ×3).
  • apps / taco (SkyThought stack) are not included in this release; the code-domain average here is over the four evalchemy-based benchmarks.
  • Checkpoints trained by this repo already carry the patched chat template (persisted at save time), so step 3 is only strictly needed when evaluating baseline HF models — it is a harmless no-op on trained checkpoints.

Citation

@article{heo2026opd2,
  title   = {On-Policy Delta Distillation},
  author  = {Heo, Byeongho and Hwang, Jaehui and Yun, Sangdoo and Han, Dongyoon},
  journal = {arXiv preprint arXiv:2607.15161},
  year    = {2026}
}

License

On-Policy Delta Distillation
Copyright (c) 2026-present NAVER Corp.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

About

Official implementation of On-Policy Delta Distillation (OPD2)

Topics

Resources

Stars

63 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages