On-Policy Distillation Example

On-Policy Distillation Example#

This example shows how to run on-policy distillation (OPD) with vime. A small student (Qwen3-4B) learns to match a larger teacher (Qwen3-32B) by training only on the student’s own vLLM rollouts and applying a token-level KL penalty against the teacher’s log-probabilities.

Key Features#

  • OPD is orthogonal to advantage estimators: OPD adds a KL penalty on top of any advantage estimator (GRPO, PPO, REINFORCE++, etc.), not as a separate estimator.

  • Two teacher modes:

    • vllm: Teacher runs on an external vLLM server; teacher log-probs are fetched during rollout via --rm-url.

    • megatron: Teacher is loaded into Megatron via --opd-teacher-load; teacher log-probs are computed during the training forward pass.

  • Student rollout always uses vLLM (vime’s default rollout backend).

Files#

File

Description

run-qwen3-4b-32b-opd.sh

8×GPU colocate demo: Qwen3-4B student + external Qwen3-32B vLLM teacher on GSM8K

run-qwen3-8b-opd-megatron.sh

Megatron-loaded teacher (no external server); student rollout still uses vLLM

GPU Layout (run-qwen3-4b-32b-opd.sh)#

Single node, 8× GPU:

GPUs

Role

0–3

Student Megatron train (TP=2) + student vLLM rollout (TP=2), colocate

4–7

Teacher vLLM (Qwen3-32B, TP=4)

Key Arguments#

Argument

Description

--use-opd

Enable on-policy distillation.

--opd-type

vllm or megatron. Required when --use-opd is set.

--opd-kl-coef

OPD KL penalty coefficient (default: 1.0).

--opd-teacher-load

Teacher checkpoint path. Required for --opd-type=megatron; must not be set for --opd-type=vllm.

--rm-url

Teacher vLLM generate endpoint. Required for --opd-type=vllm.

--custom-rm-path

vime.rollout.on_policy_distillation.reward_func

--custom-reward-post-process-path

vime.rollout.on_policy_distillation.post_process_rewards

Components#

  • vime/rollout/on_policy_distillation.py implements the vLLM teacher path:

    • reward_func POSTs each rollout sample to the teacher vLLM server (--rm-url) and collects token-level log-probs.

    • post_process_rewards trims teacher log-probs to the response span and stores them on each Sample for the OPD KL term in training.

  • Megatron teacher mode computes teacher log-probs inside apply_opd_kl_to_advantages during the training forward pass.

OPD Data Flow#

Student vLLM rollout (GPU 0-3)
  → token sequence + student logprobs
Teacher vLLM (HTTP POST, GPU 4-7)          [vllm mode only]
  → teacher_log_probs per token
post_process_rewards
  → store teacher_log_probs; scalar_rewards=[0.0] (pure distillation)
apply_opd_kl_to_advantages (Megatron)
  → advantages -= opd_kl_coef * (student_logp - teacher_logp)
GRPO policy update

Running the Example#

Prerequisites#

# Models
hf download Qwen/Qwen3-32B --local-dir /root/models/Qwen3-32B
hf download Qwen/Qwen3-4B   --local-dir /root/models/Qwen3-4B

# Data
hf download --repo-type dataset openai/gsm8k --local-dir /root/datasets/gsm8k
# Or use a parquet copy with `messages` + `label` columns under /root/datasets/gsm8k/

Step 1: Convert student checkpoint#

Qwen3-4B uses tied embeddings (tie_word_embeddings=True) — do not pass --untie-embeddings-and-output-weights. Use TP=1 for conversion; training uses TP=2 at runtime. Pad vocab to 152064 for TP=2 training:

cd /root/vime
source scripts/models/qwen3-4B.sh

PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
  --hf-checkpoint /root/models/Qwen3-4B \
  --save /root/models/Qwen3-4B_torch_dist \
  --padded-vocab-size 152064

Step 2: Run OPD (vLLM teacher)#

cd /root/vime
bash examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh

The script will:

  1. Launch the Qwen3-32B teacher vLLM server on GPUs 4–7.

  2. Start Ray on GPUs 0–3 and submit the OPD training job.

  3. Tear down the teacher server and Ray when training finishes.

Step 3 (optional): Megatron teacher#

For same-architecture teacher/student pairs that fit in GPU memory together, use the Megatron teacher path — no external vLLM teacher server needed:

# Convert teacher (example uses Qwen3-8B; use a stronger checkpoint in practice)
cd /root/vime
source scripts/models/qwen3-8B.sh
PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
  ${MODEL_ARGS[@]} \
  --hf-checkpoint /root/models/Qwen3-8B \
  --save /root/models/Qwen3-8B_torch_dist

bash examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh

Edit --opd-teacher-load in the Megatron script to point at your teacher checkpoint.

Preliminary Results#

End-to-end run with run-qwen3-4b-32b-opd.sh (500 rollout steps, GRPO + --opd-kl-coef 1.0, GSM8K greedy eval, n=1319):

Model

GSM8K Accuracy

Qwen3-4B (pre-OPD)

78.8%

Qwen3-4B (post-OPD, 500 steps)

85.6% (+6.8 pp)

Qwen3-32B teacher

88.6%

Training health signal: rollout/opd_reverse_kl dropped from 0.216 → 0.110 (−49%) over 500 steps.

FAQ#

  1. Why two OPD modes?

    • vllm: Teacher on a separate vLLM server. Use when the teacher is larger or has a different architecture than the student.

    • megatron: Teacher loaded into Megatron. Use when teacher and student share architecture and fit in training GPU memory.

  2. Why is rollout/raw_reward always 0? Pure OPD distillation does not use an external reward model. The learning signal comes entirely from the OPD KL term applied to advantages.

  3. What if I set incompatible arguments? vime validates OPD args at startup:

    • --use-opd without --opd-type → error

    • --opd-type megatron without --opd-teacher-load → error

    • --opd-type vllm with --opd-teacher-load → error

  4. Qwen3-4B checkpoint conversion fails with vocab/TP errors? Re-convert with --padded-vocab-size 152064 and TP=1 (do not set --tensor-model-parallel-size 2 during conversion). Add --make-vocab-size-divisible-by 128 at training time.

References#

  1. https://thinkingmachines.ai/blog/on-policy-distillation/

  2. https://arxiv.org/abs/2306.13649

  3. https://arxiv.org/abs/2306.08543