From ed25bc4c3b2b1c4a77062ff32a58d44806fc0a14 Mon Sep 17 00:00:00 2001
From: kaiyuan <kyxiezju@163.com>
Date: Wed, 8 Jul 2026 22:03:38 +0800
Subject: [PATCH 1/4] Add on-policy distillation example for Qwen3-4B +
 Qwen3-32B.

Signed-off-by: kaiyuan <kyxiezju@163.com>
---
 examples/README.md                            |   2 +-
 examples/on_policy_distillation/README.md     | 176 +++++++++++++++
 .../run-qwen3-4b-32b-opd.sh                   | 211 ++++++++++++++++++
 .../run-qwen3-8b-opd-megatron.sh              | 153 +++++++++++++
 4 files changed, 541 insertions(+), 1 deletion(-)
 create mode 100644 examples/on_policy_distillation/README.md
 create mode 100644 examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
 create mode 100644 examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh

diff --git a/examples/README.md b/examples/README.md
index 1daea83c0..f8ea2cb74 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -11,7 +11,7 @@ These examples provide concrete examples to leverage vime in your own RL workflo
 - **[low_precision](./low_precision)**: Examples of FP8 training and inference for improved throughput and stability.
 - **[mem_agent](./mem_agent)**: MemAgent long-context RL — chunk-wise memory update, HotpotQA GRPO training, and RULER-HQA evaluation.
 - **[multi_agent](./multi_agent)**: Example of running multi-agent RL with `vime`.
-- **[on_policy_distillation](./on_policy_distillation)**: Example implementation for on-policy distillation, extending the reinforcement learning pipeline to support teacher–student distillation directly within on-policy training.
+- **[on_policy_distillation](./on_policy_distillation)**: On-policy distillation (OPD) with an external vLLM teacher or a Megatron-loaded teacher.
 - **[delta_weight_sync](./delta_weight_sync)**: Non-colocated weight sync that ships only the changed bytes over a shared filesystem (training/inference disaggregation), reloading via the vanilla `update_weights_from_disk` path.
 - **[reproducibility](./reproducibility)**: Guides on achieving bitwise experiment reproduction using deterministic modes.
 - **[retool](./retool)**: Demonstrates the retool functionality for tool-enabled language model generation.
diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md
new file mode 100644
index 000000000..b022dd9cf
--- /dev/null
+++ b/examples/on_policy_distillation/README.md
@@ -0,0 +1,176 @@
+# 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
+
+```bash
+# 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:
+
+```bash
+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)
+
+```bash
+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:
+
+```bash
+# 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
diff --git a/examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh b/examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
new file mode 100644
index 000000000..d2d288636
--- /dev/null
+++ b/examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
@@ -0,0 +1,211 @@
+#!/bin/bash
+
+# usage: bash examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
+#
+# 8×GPU colocate OPD: Qwen3-4B student (GPUs 0-3) + Qwen3-32B vLLM teacher (GPUs 4-7).
+# See README.md for checkpoint conversion and data prerequisites.
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+VIME_ROOT=/root/vime
+NUM_TRAIN_GPUS=4
+TEACHER_TP=4
+TEACHER_HOST=127.0.0.1
+TEACHER_PORT=13141
+
+TEACHER_MODEL_PATH=/root/models/Qwen3-32B
+STUDENT_HF=/root/models/Qwen3-4B
+STUDENT_TORCH_DIST=/root/models/Qwen3-4B_torch_dist
+DATA_DIR=/root/datasets
+SAVE_DIR=/root/Qwen3-4B_opd
+LOG_DIR=/root/opd_logs
+
+mkdir -p "${LOG_DIR}" "${SAVE_DIR}"
+
+source "${VIME_ROOT}/scripts/models/qwen3-4B.sh"
+
+echo "=== Step 1: Clean up previous Ray / training processes ==="
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 -f "train.py" 2>/dev/null || true
+sleep 3
+
+echo "=== Step 2: Launch vLLM teacher server (Qwen3-32B, TP=${TEACHER_TP}) ==="
+CUDA_VISIBLE_DEVICES=4,5,6,7 python3 -m vllm.entrypoints.openai.api_server \
+    --model "${TEACHER_MODEL_PATH}" \
+    --host 0.0.0.0 \
+    --port "${TEACHER_PORT}" \
+    --tensor-parallel-size "${TEACHER_TP}" \
+    --gpu-memory-utilization 0.85 \
+    --trust-remote-code \
+    --dtype bfloat16 \
+    --max-model-len 8192 \
+    > "${LOG_DIR}/teacher_vllm.log" 2>&1 &
+TEACHER_PID=$!
+echo "Teacher vLLM server PID: ${TEACHER_PID}"
+
+echo "Waiting for teacher server to be ready..."
+for i in $(seq 1 120); do
+    if ! kill -0 "${TEACHER_PID}" 2>/dev/null; then
+        echo "ERROR: Teacher server process died. Check ${LOG_DIR}/teacher_vllm.log"
+        exit 1
+    fi
+    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" "http://${TEACHER_HOST}:${TEACHER_PORT}/health" 2>/dev/null || true)
+    if [ "${HTTP_CODE}" = "200" ]; then
+        echo "Teacher vLLM server is ready!"
+        break
+    fi
+    if [ "$i" -eq 120 ]; then
+        echo "ERROR: Teacher server failed to start within 10 minutes"
+        kill "${TEACHER_PID}" 2>/dev/null || true
+        exit 1
+    fi
+    sleep 5
+done
+
+echo "=== Step 3: Run OPD training (Qwen3-4B student on GPUs 0-3) ==="
+
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+
+ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${NUM_TRAIN_GPUS}" \
+    --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+CKPT_ARGS=(
+    --hf-checkpoint "${STUDENT_HF}"
+    --ref-load "${STUDENT_TORCH_DIST}"
+    --load "${STUDENT_TORCH_DIST}"
+    --save "${SAVE_DIR}"
+    --save-interval 50
+    --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+    --prompt-data "${DATA_DIR}/gsm8k/train.parquet"
+    --input-key messages
+    --label-key label
+    --apply-chat-template
+    --rollout-shuffle
+    --rm-type math
+    --num-rollout 500
+    --rollout-batch-size 32
+    --n-samples-per-prompt 4
+    --rollout-max-response-len 4096
+    --rollout-temperature 0.8
+    --global-batch-size 64
+)
+
+EVAL_ARGS=(
+    --eval-interval 50
+    --eval-prompt-data gsm8k "${DATA_DIR}/gsm8k/test.parquet"
+    --n-samples-per-eval-prompt 1
+    --eval-max-response-len 4096
+    --eval-top-k 1
+)
+
+PERF_ARGS=(
+    --tensor-model-parallel-size 2
+    --sequence-parallel
+    --pipeline-model-parallel-size 1
+    --context-parallel-size 1
+    --expert-model-parallel-size 1
+    --expert-tensor-parallel-size 1
+    --use-dynamic-batch-size
+    --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+    --advantage-estimator grpo
+    --use-opd
+    --opd-type vllm
+    --opd-kl-coef 1.0
+    --use-kl-loss
+    --kl-loss-coef 0.00
+    --kl-loss-type low_var_kl
+    --entropy-coef 0.00
+    --eps-clip 0.2
+    --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+    --optimizer adam
+    --lr 1e-6
+    --lr-decay-style constant
+    --weight-decay 0.1
+    --adam-beta1 0.9
+    --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+    # --use-wandb
+    # --wandb-project vime-opd
+    # --wandb-group qwen3-4b-32b-opd
+    # --wandb-key ${WANDB_KEY}
+)
+
+VLLM_ARGS=(
+    --rollout-num-gpus-per-engine 2
+    --vllm-gpu-memory-utilization 0.7
+    --vllm-max-num-seqs 32
+    --vllm-max-cudagraph-capture-size 16
+)
+
+RM_ARGS=(
+    --custom-rm-path vime.rollout.on_policy_distillation.reward_func
+    --custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards
+    --rm-url "http://${TEACHER_HOST}:${TEACHER_PORT}/inference/v1/generate"
+)
+
+MISC_ARGS=(
+    --attention-dropout 0.0
+    --hidden-dropout 0.0
+    --accumulate-allreduce-grads-in-fp32
+    --attention-softmax-in-fp32
+    --attention-backend flash
+    --actor-num-nodes 1
+    --actor-num-gpus-per-node "${NUM_TRAIN_GPUS}"
+    --colocate
+    --make-vocab-size-divisible-by 128
+)
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+    --runtime-env-json="${RUNTIME_ENV_JSON}" \
+    -- python3 train.py \
+    --train-backend megatron \
+    ${MODEL_ARGS[@]} \
+    ${CKPT_ARGS[@]} \
+    ${ROLLOUT_ARGS[@]} \
+    ${OPTIMIZER_ARGS[@]} \
+    ${GRPO_ARGS[@]} \
+    ${WANDB_ARGS[@]} \
+    ${PERF_ARGS[@]} \
+    ${EVAL_ARGS[@]} \
+    ${VLLM_ARGS[@]} \
+    ${RM_ARGS[@]} \
+    ${MISC_ARGS[@]} \
+    2>&1 | tee "${LOG_DIR}/opd_training.log"
+
+echo "=== Training complete, stopping teacher server ==="
+kill "${TEACHER_PID}" 2>/dev/null || true
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 -f "train.py" 2>/dev/null || true
+echo "=== Done ==="
diff --git a/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
new file mode 100644
index 000000000..e1c6e75fc
--- /dev/null
+++ b/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
@@ -0,0 +1,153 @@
+#!/bin/bash
+
+# usage: bash examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
+#
+# OPD with a Megatron-loaded teacher (no external vLLM teacher server).
+# Student rollout still uses vLLM. This demo uses the same architecture for
+# student and teacher — replace --opd-teacher-load with a stronger checkpoint
+# in practice.
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+pkill -9 vllm 2>/dev/null || true
+sleep 3
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 python 2>/dev/null || true
+sleep 3
+
+source "/root/vime/scripts/models/qwen3-8B.sh"
+
+CKPT_ARGS=(
+    --hf-checkpoint /root/models/Qwen3-8B
+    --ref-load /root/models/Qwen3-8B_torch_dist
+    --load /root/models/Qwen3-8B_torch_dist
+    --save /root/Qwen3-8B_opd/
+    --save-interval 20
+    --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+    --prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
+    --input-key prompt
+    --apply-chat-template
+    --rollout-shuffle
+    --rm-type math
+    --num-rollout 300
+    --rollout-batch-size 16
+    --n-samples-per-prompt 4
+    --rollout-max-response-len 16384
+    --rollout-temperature 1
+    --global-batch-size 64
+    --balance-data
+)
+
+EVAL_ARGS=(
+    # --eval-interval 20
+    # --eval-prompt-data aime /root/datasets/aime-2024/aime-2024.jsonl
+    # --n-samples-per-eval-prompt 16
+    # --eval-max-response-len 16384
+    # --eval-top-p 1
+)
+
+PERF_ARGS=(
+    --tensor-model-parallel-size 2
+    --sequence-parallel
+    --pipeline-model-parallel-size 1
+    --context-parallel-size 1
+    --expert-model-parallel-size 1
+    --expert-tensor-parallel-size 1
+    --recompute-granularity full
+    --recompute-method uniform
+    --recompute-num-layers 1
+    --use-dynamic-batch-size
+    --max-tokens-per-gpu 16384
+)
+
+GRPO_ARGS=(
+    --advantage-estimator grpo
+    --use-opd
+    --opd-type megatron
+    --opd-kl-coef 1.0
+    --opd-teacher-load /root/models/Qwen3-8B_torch_dist
+    --use-kl-loss
+    --kl-loss-coef 0.00
+    --kl-loss-type low_var_kl
+    --entropy-coef 0.00
+)
+
+OPTIMIZER_ARGS=(
+    --optimizer adam
+    --lr 1e-6
+    --lr-decay-style constant
+    --weight-decay 0.1
+    --adam-beta1 0.9
+    --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+    # --use-wandb
+    # --wandb-project vime-opd
+    # --wandb-group qwen3-8b-opd-megatron
+    # --wandb-key ${WANDB_KEY}
+)
+
+VLLM_ARGS=(
+    --rollout-num-gpus-per-engine 1
+    --vllm-gpu-memory-utilization 0.4
+)
+
+MISC_ARGS=(
+    --attention-dropout 0.0
+    --hidden-dropout 0.0
+    --accumulate-allreduce-grads-in-fp32
+    --attention-softmax-in-fp32
+    --attention-backend flash
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus 8 \
+    --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+    --runtime-env-json="${RUNTIME_ENV_JSON}" \
+    -- python3 train.py \
+    --train-backend megatron \
+    --actor-num-nodes 1 \
+    --actor-num-gpus-per-node 4 \
+    --rollout-num-gpus 4 \
+    --colocate \
+    ${MODEL_ARGS[@]} \
+    ${CKPT_ARGS[@]} \
+    ${ROLLOUT_ARGS[@]} \
+    ${OPTIMIZER_ARGS[@]} \
+    ${GRPO_ARGS[@]} \
+    ${WANDB_ARGS[@]} \
+    ${PERF_ARGS[@]} \
+    ${EVAL_ARGS[@]} \
+    ${VLLM_ARGS[@]} \
+    ${MISC_ARGS[@]}
+
+pkill -9 vllm 2>/dev/null || true
+sleep 3
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 python 2>/dev/null || true

From d66664a1d6648d0ab35afa48acfe4e25b281ad30 Mon Sep 17 00:00:00 2001
From: kaiyuan <kyxiezju@163.com>
Date: Tue, 14 Jul 2026 15:33:50 +0800
Subject: [PATCH 2/4] Harden Megatron OPD example against OOM on 8x80GB.

Signed-off-by: kaiyuan <kyxiezju@163.com>
---
 examples/on_policy_distillation/README.md     | 44 ++++++++++++++++---
 .../run-qwen3-8b-opd-megatron.sh              | 38 +++++++++-------
 2 files changed, 61 insertions(+), 21 deletions(-)

diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md
index b022dd9cf..67a108211 100644
--- a/examples/on_policy_distillation/README.md
+++ b/examples/on_policy_distillation/README.md
@@ -21,10 +21,12 @@ the teacher's log-probabilities.
 
 | 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 |
+| `run-qwen3-4b-32b-opd.sh` | **Recommended.** 8×GPU colocate: Qwen3-4B student + external Qwen3-32B vLLM teacher (GSM8K validated) |
+| `run-qwen3-8b-opd-megatron.sh` | Megatron-loaded teacher (self-distillation demo); includes OOM mitigations for 8×80GB |
 
-## GPU Layout (`run-qwen3-4b-32b-opd.sh`)
+## GPU Layout
+
+### `run-qwen3-4b-32b-opd.sh` (vLLM teacher)
 
 Single node, 8× GPU:
 
@@ -33,6 +35,17 @@ Single node, 8× GPU:
 | 0–3 | Student Megatron train (TP=2) + student vLLM rollout (TP=2), colocate |
 | 4–7 | Teacher vLLM (Qwen3-32B, TP=4) |
 
+### `run-qwen3-8b-opd-megatron.sh` (Megatron teacher)
+
+| GPUs | Role |
+|------|------|
+| 0–3 | Student + teacher Megatron (TP=4), colocate with student vLLM rollout |
+| 4–7 | Student vLLM rollout engines |
+
+Teacher and student share the same architecture (demo uses Qwen3-8B
+self-distillation). Prefer `--opd-type vllm` when the teacher is larger or does
+not fit together with the student in training GPU memory.
+
 ## Key Arguments
 
 | Argument | Description |
@@ -118,7 +131,7 @@ For same-architecture teacher/student pairs that fit in GPU memory together,
 use the Megatron teacher path — no external vLLM teacher server needed:
 
 ```bash
-# Convert teacher (example uses Qwen3-8B; use a stronger checkpoint in practice)
+# Convert student/teacher (demo uses Qwen3-8B self-distillation)
 cd /root/vime
 source scripts/models/qwen3-8B.sh
 PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
@@ -126,11 +139,16 @@ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
   --hf-checkpoint /root/models/Qwen3-8B \
   --save /root/models/Qwen3-8B_torch_dist
 
+# Also prepare rollout data + GSM8K eval set:
+#   /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
+#   /root/datasets/gsm8k/test.parquet
+
 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.
+checkpoint. On 8×A800 80GB, keep the memory defaults (TP=4, full recompute,
+`max-tokens-per-gpu=8192`); see FAQ.
 
 ## Preliminary Results
 
@@ -169,6 +187,22 @@ Training health signal: `rollout/opd_reverse_kl` dropped from 0.216 → 0.110
    `--tensor-model-parallel-size 2` during conversion). Add
    `--make-vocab-size-divisible-by 128` at training time.
 
+5. **Megatron teacher OOM on Qwen3-8B + Qwen3-8B?**
+   Student weights, optimizer state, and teacher weights share the 4 training
+   GPUs. The example script already applies the known mitigations:
+   - `--tensor-model-parallel-size 4`
+   - `--recompute-num-layers 36` (full activation recompute)
+   - `--max-tokens-per-gpu 8192`
+   - `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`
+   - `--save-interval 10` so checkpoints land before a long-run OOM
+   If memory is still tight, reduce `--num-rollout` / `--rollout-max-response-len`,
+   or switch to `--opd-type vllm` with an external teacher server.
+
+6. **Self-distillation: why is `opd_reverse_kl` near 0 at the start?**
+   Teacher and student start from the same weights, so reverse KL is ~0 until
+   the student updates. For a true distillation signal, use a stronger /
+   differently trained teacher checkpoint.
+
 ## References
 
 1. https://thinkingmachines.ai/blog/on-policy-distillation/
diff --git a/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
index e1c6e75fc..d211408ea 100644
--- a/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
+++ b/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
@@ -6,10 +6,16 @@
 # Student rollout still uses vLLM. This demo uses the same architecture for
 # student and teacher — replace --opd-teacher-load with a stronger checkpoint
 # in practice.
+#
+# Memory note (8×A800 80GB): loading student + teacher (both Qwen3-8B) on
+# 4 training GPUs is tight. The defaults below match a validated config:
+# TP=4, full activation recompute, max-tokens-per-gpu=8192, and
+# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. See README FAQ.
 
 set -ex
 
 export PYTHONUNBUFFERED=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 
 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
 if [ "$NVLINK_COUNT" -gt 0 ]; then
@@ -19,11 +25,9 @@ else
 fi
 echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
 
-pkill -9 vllm 2>/dev/null || true
-sleep 3
 ray stop --force 2>/dev/null || true
 pkill -9 ray 2>/dev/null || true
-pkill -9 python 2>/dev/null || true
+pkill -9 -f "train.py" 2>/dev/null || true
 sleep 3
 
 source "/root/vime/scripts/models/qwen3-8B.sh"
@@ -33,7 +37,7 @@ CKPT_ARGS=(
     --ref-load /root/models/Qwen3-8B_torch_dist
     --load /root/models/Qwen3-8B_torch_dist
     --save /root/Qwen3-8B_opd/
-    --save-interval 20
+    --save-interval 10
     --megatron-to-hf-mode bridge
 )
 
@@ -53,15 +57,17 @@ ROLLOUT_ARGS=(
 )
 
 EVAL_ARGS=(
-    # --eval-interval 20
-    # --eval-prompt-data aime /root/datasets/aime-2024/aime-2024.jsonl
-    # --n-samples-per-eval-prompt 16
-    # --eval-max-response-len 16384
-    # --eval-top-p 1
+    --eval-interval 50
+    --eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet
+    --n-samples-per-eval-prompt 1
+    --eval-max-response-len 4096
+    --eval-top-k 1
 )
 
 PERF_ARGS=(
-    --tensor-model-parallel-size 2
+    # TP=4 + full recompute + lower token budget: needed so student+teacher
+    # (both Qwen3-8B) fit on 4×80GB training GPUs.
+    --tensor-model-parallel-size 4
     --sequence-parallel
     --pipeline-model-parallel-size 1
     --context-parallel-size 1
@@ -69,9 +75,9 @@ PERF_ARGS=(
     --expert-tensor-parallel-size 1
     --recompute-granularity full
     --recompute-method uniform
-    --recompute-num-layers 1
+    --recompute-num-layers 36
     --use-dynamic-batch-size
-    --max-tokens-per-gpu 16384
+    --max-tokens-per-gpu 8192
 )
 
 GRPO_ARGS=(
@@ -113,6 +119,7 @@ MISC_ARGS=(
     --accumulate-allreduce-grads-in-fp32
     --attention-softmax-in-fp32
     --attention-backend flash
+    --make-vocab-size-divisible-by 128
 )
 
 export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
@@ -123,7 +130,8 @@ RUNTIME_ENV_JSON="{
   \"env_vars\": {
     \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
     \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
-    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
+    \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\"
   }
 }"
 
@@ -146,8 +154,6 @@ ray job submit --address="http://127.0.0.1:8265" \
     ${VLLM_ARGS[@]} \
     ${MISC_ARGS[@]}
 
-pkill -9 vllm 2>/dev/null || true
-sleep 3
 ray stop --force 2>/dev/null || true
 pkill -9 ray 2>/dev/null || true
-pkill -9 python 2>/dev/null || true
+pkill -9 -f "train.py" 2>/dev/null || true

From 9c9c0f9763d89644e5371ed44e2b8b6e386632a3 Mon Sep 17 00:00:00 2001
From: kaiyuan <kyxiezju@163.com>
Date: Sat, 18 Jul 2026 18:14:26 +0800
Subject: [PATCH 3/4] Align OPD example with slime layout using validated
 8B+32B run.

Signed-off-by: kaiyuan <kyxiezju@163.com>
---
 examples/on_policy_distillation/README.md     | 273 ++++++++----------
 .../run-qwen3-4b-32b-opd.sh                   | 211 --------------
 .../run-qwen3-8B-opd-megatron.sh              | 167 +++++++++++
 .../run-qwen3-8B-opd.sh                       | 211 ++++++++++++++
 .../run-qwen3-8b-opd-megatron.sh              | 159 ----------
 5 files changed, 500 insertions(+), 521 deletions(-)
 delete mode 100644 examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
 create mode 100644 examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
 create mode 100644 examples/on_policy_distillation/run-qwen3-8B-opd.sh
 delete mode 100644 examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh

diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md
index 67a108211..257db6287 100644
--- a/examples/on_policy_distillation/README.md
+++ b/examples/on_policy_distillation/README.md
@@ -1,209 +1,180 @@
 # 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.
+This example shows how to run **on-policy distillation (OPD)** using vime. A
+small student (Qwen3-8B) is aligned to imitate a larger teacher (Qwen3-32B) by
+training only on the student's own rollouts and matching the teacher's
+token-level 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.
+- **OPD is orthogonal to advantage estimators**: OPD works as an additive 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.
+  - **vllm**: Teacher runs on an external vLLM server; teacher log-probs are
+    obtained during rollout.
+  - **megatron**: Teacher is loaded directly 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` | **Recommended.** 8×GPU colocate: Qwen3-4B student + external Qwen3-32B vLLM teacher (GSM8K validated) |
-| `run-qwen3-8b-opd-megatron.sh` | Megatron-loaded teacher (self-distillation demo); includes OOM mitigations for 8×80GB |
-
-## GPU Layout
-
-### `run-qwen3-4b-32b-opd.sh` (vLLM teacher)
-
-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) |
-
-### `run-qwen3-8b-opd-megatron.sh` (Megatron teacher)
-
-| GPUs | Role |
-|------|------|
-| 0–3 | Student + teacher Megatron (TP=4), colocate with student vLLM rollout |
-| 4–7 | Student vLLM rollout engines |
-
-Teacher and student share the same architecture (demo uses Qwen3-8B
-self-distillation). Prefer `--opd-type vllm` when the teacher is larger or does
-not fit together with the student in training GPU memory.
-
 ## Key Arguments
 
 | Argument | Description |
 |----------|-------------|
-| `--use-opd` | Enable on-policy distillation. |
-| `--opd-type` | `vllm` or `megatron`. Required when `--use-opd` is set. |
+| `--use-opd` | Enable on-policy distillation. Required flag to use OPD. |
+| `--opd-type` | Type of OPD: `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` |
+| `--opd-teacher-load` | Path to teacher checkpoint. **Required** when `--opd-type=megatron`, **must not be set** when `--opd-type=vllm`. |
+| `--opd-teacher-ckpt-step` | Optional checkpoint step for teacher model. |
 
-## Components
+## Mode Comparison
 
-- `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.
+| Mode | Teacher Location | When to use |
+|------|------------------|-------------|
+| `vllm` | External vLLM server | Teacher has different architecture or is larger than GPU memory |
+| `megatron` | Loaded into Megatron training | Teacher has same architecture as policy/ref model |
 
-## OPD Data Flow
+## Components
 
-```
-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
-```
+- `vime/rollout/on_policy_distillation.py` implements (for vLLM mode):
+  - `reward_func` calls the teacher server (via `args.rm_url`) with every sample
+    to obtain token-level logprobs.
+  - `post_process_rewards` trims the teacher logprobs to the generated response
+    span and writes the tensors back to each `Sample` to compute advantages.
+- `run-qwen3-8B-opd.sh` launches a vLLM teacher server, then submits a Ray job
+  that runs `train.py`.
+- `run-qwen3-8B-opd-megatron.sh` uses a Megatron-loaded teacher model (no
+  external server needed).
 
-## Running the Example
+## Running the example
 
-### Prerequisites
+### Using vLLM Teacher (External Server)
 
-```bash
-# Models
-hf download Qwen/Qwen3-32B --local-dir /root/models/Qwen3-32B
-hf download Qwen/Qwen3-4B   --local-dir /root/models/Qwen3-4B
+1. Download or prepare the required checkpoints and data.
 
-# 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/
+```bash
+hf download Qwen/Qwen3-32B --local-dir /root/Qwen3-32B
+hf download Qwen/Qwen3-8B --local-dir /root/Qwen3-8B
+hf download --repo-type dataset zhuzilin/dapo-math-17k --local-dir /root/dapo-math-17k
 ```
 
-### 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:
+2. Run the hf to mcore for student model conversion:
 
 ```bash
 cd /root/vime
-source scripts/models/qwen3-4B.sh
+source scripts/models/qwen3-8B.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
+    ${MODEL_ARGS[@]} \
+    --hf-checkpoint /root/Qwen3-8B \
+    --save /root/Qwen3-8B_torch_dist
 ```
 
-### Step 2: Run OPD (vLLM teacher)
+3. Run on-policy distillation:
 
 ```bash
-cd /root/vime
-bash examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
+bash examples/on_policy_distillation/run-qwen3-8B-opd.sh
 ```
 
-The script will:
+GPU layout (8×GPU):
 
-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.
+| GPUs | Role |
+|------|------|
+| 0–3 | Student Megatron train + student vLLM rollout (colocate) |
+| 4–7 | Teacher vLLM (Qwen3-32B, TP=4) |
 
-### Step 3 (optional): Megatron teacher
+### Using Megatron Teacher (No External Server)
 
-For same-architecture teacher/student pairs that fit in GPU memory together,
-use the Megatron teacher path — no external vLLM teacher server needed:
+1. Prepare student checkpoint (same as above).
+
+2. **IMPORTANT**: Convert your teacher model to Megatron format (change the path
+   to your actual teacher):
 
 ```bash
-# Convert student/teacher (demo uses Qwen3-8B self-distillation)
+# This example uses the same model as both student and teacher (for demonstration only)
+# In practice, use a different (stronger) model as the teacher!
 cd /root/vime
-source scripts/models/qwen3-8B.sh
+source scripts/models/qwen3-8B.sh  # Or your teacher model config
+
 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
+    ${MODEL_ARGS[@]} \
+    --hf-checkpoint /root/YourTeacherModel \
+    --save /root/YourTeacherModel_torch_dist
+```
 
-# Also prepare rollout data + GSM8K eval set:
-#   /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
-#   /root/datasets/gsm8k/test.parquet
+3. Edit `run-qwen3-8B-opd-megatron.sh` to update paths:
+   - Change `--opd-teacher-load` to your teacher model path
+   - Adjust `--opd-kl-coef` based on your task
 
-bash examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
-```
+4. Run:
 
-Edit `--opd-teacher-load` in the Megatron script to point at your teacher
-checkpoint. On 8×A800 80GB, keep the memory defaults (TP=4, full recompute,
-`max-tokens-per-gpu=8192`); see FAQ.
+```bash
+bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
+```
 
-## Preliminary Results
+# 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):
+End-to-end run with `run-qwen3-8B-opd.sh` on 8×A800 80GB (dapo-math-17k train,
+GRPO + `--opd-kl-coef 1.0`, ~220 rollouts / iter_0000219). Offline GSM8K greedy
+eval:
 
 | 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?**
+| Qwen3-8B (pre-OPD) | 79.7% (n=300) |
+| Qwen3-8B (post-OPD) | **88.2%** (n=1319, **+8.5 pp**) |
+| Qwen3-32B teacher | 87.0% (n=300) |
+
+Training health signal: `rollout/opd_reverse_kl` dropped from 0.145 → ~0.10
+(−38%). Pure OPD uses `raw_reward=0`; the learning signal is the OPD KL term.
+
+Notes from the validated run:
+
+- Colocate memory is tight for 8B+32B; keep `--rollout-max-response-len 4096`,
+  `--rollout-max-context-len 8192`, `--max-tokens-per-gpu 2048`, and
+  `--vllm-gpu-memory-utilization 0.25` unless you have more headroom.
+- Teacher vLLM needs `--max-model-len 16384` and `--disable-custom-all-reduce`
+  for TP=4 stability.
+- Prefer offline eval after training; in-training eval is currently incompatible
+  with OPD's custom reward payload.
+- Megatron checkpoints are large (~100GB each); plan disk or lower
+  `--save-interval`.
+
+# FAQ
+
+1. **Why are there two OPD modes?**
+   - `vllm` mode: The teacher runs on an independent vLLM server. This is useful
+     when the teacher has a different architecture or is too large to load
+     together with the policy model.
+   - `megatron` mode: The teacher is loaded into Megatron using the same
+     parameter loading mechanism as the reference model. This requires the
+     teacher to have the same architecture as the policy model.
+
+2. **How do I use Megatron-based teacher instead of vLLM server?**
+   Replace your OPD arguments:
+   ```bash
+   # Instead of:
+   --use-opd --opd-type vllm --opd-kl-coef 1.0
+   # Use:
+   --use-opd --opd-type megatron --opd-kl-coef 1.0 --opd-teacher-load /path/to/teacher_checkpoint
+   ```
+
+3. **What happens if I set wrong arguments?**
+   The system will raise clear errors:
+   - `--use-opd` without `--opd-type`: Error asking you to specify type
+   - `--opd-type megatron` without `--opd-teacher-load`: Error asking for teacher checkpoint
+   - `--opd-type vllm` with `--opd-teacher-load`: Error indicating conflict
+
+4. **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.
-
-5. **Megatron teacher OOM on Qwen3-8B + Qwen3-8B?**
-   Student weights, optimizer state, and teacher weights share the 4 training
-   GPUs. The example script already applies the known mitigations:
-   - `--tensor-model-parallel-size 4`
-   - `--recompute-num-layers 36` (full activation recompute)
-   - `--max-tokens-per-gpu 8192`
-   - `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`
-   - `--save-interval 10` so checkpoints land before a long-run OOM
-   If memory is still tight, reduce `--num-rollout` / `--rollout-max-response-len`,
-   or switch to `--opd-type vllm` with an external teacher server.
-
-6. **Self-distillation: why is `opd_reverse_kl` near 0 at the start?**
+5. **Self-distillation: why is `opd_reverse_kl` near 0 at the start?**
    Teacher and student start from the same weights, so reverse KL is ~0 until
    the student updates. For a true distillation signal, use a stronger /
-   differently trained teacher checkpoint.
+   differently trained teacher (or `--opd-type vllm` with Qwen3-32B).
 
-## References
+# References
 
 1. https://thinkingmachines.ai/blog/on-policy-distillation/
 2. https://arxiv.org/abs/2306.13649
diff --git a/examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh b/examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
deleted file mode 100644
index d2d288636..000000000
--- a/examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
+++ /dev/null
@@ -1,211 +0,0 @@
-#!/bin/bash
-
-# usage: bash examples/on_policy_distillation/run-qwen3-4b-32b-opd.sh
-#
-# 8×GPU colocate OPD: Qwen3-4B student (GPUs 0-3) + Qwen3-32B vLLM teacher (GPUs 4-7).
-# See README.md for checkpoint conversion and data prerequisites.
-
-set -ex
-
-export PYTHONUNBUFFERED=1
-
-NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
-if [ "$NVLINK_COUNT" -gt 0 ]; then
-    HAS_NVLINK=1
-else
-    HAS_NVLINK=0
-fi
-echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
-
-VIME_ROOT=/root/vime
-NUM_TRAIN_GPUS=4
-TEACHER_TP=4
-TEACHER_HOST=127.0.0.1
-TEACHER_PORT=13141
-
-TEACHER_MODEL_PATH=/root/models/Qwen3-32B
-STUDENT_HF=/root/models/Qwen3-4B
-STUDENT_TORCH_DIST=/root/models/Qwen3-4B_torch_dist
-DATA_DIR=/root/datasets
-SAVE_DIR=/root/Qwen3-4B_opd
-LOG_DIR=/root/opd_logs
-
-mkdir -p "${LOG_DIR}" "${SAVE_DIR}"
-
-source "${VIME_ROOT}/scripts/models/qwen3-4B.sh"
-
-echo "=== Step 1: Clean up previous Ray / training processes ==="
-ray stop --force 2>/dev/null || true
-pkill -9 ray 2>/dev/null || true
-pkill -9 -f "train.py" 2>/dev/null || true
-sleep 3
-
-echo "=== Step 2: Launch vLLM teacher server (Qwen3-32B, TP=${TEACHER_TP}) ==="
-CUDA_VISIBLE_DEVICES=4,5,6,7 python3 -m vllm.entrypoints.openai.api_server \
-    --model "${TEACHER_MODEL_PATH}" \
-    --host 0.0.0.0 \
-    --port "${TEACHER_PORT}" \
-    --tensor-parallel-size "${TEACHER_TP}" \
-    --gpu-memory-utilization 0.85 \
-    --trust-remote-code \
-    --dtype bfloat16 \
-    --max-model-len 8192 \
-    > "${LOG_DIR}/teacher_vllm.log" 2>&1 &
-TEACHER_PID=$!
-echo "Teacher vLLM server PID: ${TEACHER_PID}"
-
-echo "Waiting for teacher server to be ready..."
-for i in $(seq 1 120); do
-    if ! kill -0 "${TEACHER_PID}" 2>/dev/null; then
-        echo "ERROR: Teacher server process died. Check ${LOG_DIR}/teacher_vllm.log"
-        exit 1
-    fi
-    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" "http://${TEACHER_HOST}:${TEACHER_PORT}/health" 2>/dev/null || true)
-    if [ "${HTTP_CODE}" = "200" ]; then
-        echo "Teacher vLLM server is ready!"
-        break
-    fi
-    if [ "$i" -eq 120 ]; then
-        echo "ERROR: Teacher server failed to start within 10 minutes"
-        kill "${TEACHER_PID}" 2>/dev/null || true
-        exit 1
-    fi
-    sleep 5
-done
-
-echo "=== Step 3: Run OPD training (Qwen3-4B student on GPUs 0-3) ==="
-
-export CUDA_VISIBLE_DEVICES=0,1,2,3
-export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
-
-ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${NUM_TRAIN_GPUS}" \
-    --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-
-CKPT_ARGS=(
-    --hf-checkpoint "${STUDENT_HF}"
-    --ref-load "${STUDENT_TORCH_DIST}"
-    --load "${STUDENT_TORCH_DIST}"
-    --save "${SAVE_DIR}"
-    --save-interval 50
-    --megatron-to-hf-mode bridge
-)
-
-ROLLOUT_ARGS=(
-    --prompt-data "${DATA_DIR}/gsm8k/train.parquet"
-    --input-key messages
-    --label-key label
-    --apply-chat-template
-    --rollout-shuffle
-    --rm-type math
-    --num-rollout 500
-    --rollout-batch-size 32
-    --n-samples-per-prompt 4
-    --rollout-max-response-len 4096
-    --rollout-temperature 0.8
-    --global-batch-size 64
-)
-
-EVAL_ARGS=(
-    --eval-interval 50
-    --eval-prompt-data gsm8k "${DATA_DIR}/gsm8k/test.parquet"
-    --n-samples-per-eval-prompt 1
-    --eval-max-response-len 4096
-    --eval-top-k 1
-)
-
-PERF_ARGS=(
-    --tensor-model-parallel-size 2
-    --sequence-parallel
-    --pipeline-model-parallel-size 1
-    --context-parallel-size 1
-    --expert-model-parallel-size 1
-    --expert-tensor-parallel-size 1
-    --use-dynamic-batch-size
-    --max-tokens-per-gpu 9216
-)
-
-GRPO_ARGS=(
-    --advantage-estimator grpo
-    --use-opd
-    --opd-type vllm
-    --opd-kl-coef 1.0
-    --use-kl-loss
-    --kl-loss-coef 0.00
-    --kl-loss-type low_var_kl
-    --entropy-coef 0.00
-    --eps-clip 0.2
-    --eps-clip-high 0.28
-)
-
-OPTIMIZER_ARGS=(
-    --optimizer adam
-    --lr 1e-6
-    --lr-decay-style constant
-    --weight-decay 0.1
-    --adam-beta1 0.9
-    --adam-beta2 0.98
-)
-
-WANDB_ARGS=(
-    # --use-wandb
-    # --wandb-project vime-opd
-    # --wandb-group qwen3-4b-32b-opd
-    # --wandb-key ${WANDB_KEY}
-)
-
-VLLM_ARGS=(
-    --rollout-num-gpus-per-engine 2
-    --vllm-gpu-memory-utilization 0.7
-    --vllm-max-num-seqs 32
-    --vllm-max-cudagraph-capture-size 16
-)
-
-RM_ARGS=(
-    --custom-rm-path vime.rollout.on_policy_distillation.reward_func
-    --custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards
-    --rm-url "http://${TEACHER_HOST}:${TEACHER_PORT}/inference/v1/generate"
-)
-
-MISC_ARGS=(
-    --attention-dropout 0.0
-    --hidden-dropout 0.0
-    --accumulate-allreduce-grads-in-fp32
-    --attention-softmax-in-fp32
-    --attention-backend flash
-    --actor-num-nodes 1
-    --actor-num-gpus-per-node "${NUM_TRAIN_GPUS}"
-    --colocate
-    --make-vocab-size-divisible-by 128
-)
-
-RUNTIME_ENV_JSON="{
-  \"env_vars\": {
-    \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
-    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
-    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
-  }
-}"
-
-ray job submit --address="http://127.0.0.1:8265" \
-    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-    -- python3 train.py \
-    --train-backend megatron \
-    ${MODEL_ARGS[@]} \
-    ${CKPT_ARGS[@]} \
-    ${ROLLOUT_ARGS[@]} \
-    ${OPTIMIZER_ARGS[@]} \
-    ${GRPO_ARGS[@]} \
-    ${WANDB_ARGS[@]} \
-    ${PERF_ARGS[@]} \
-    ${EVAL_ARGS[@]} \
-    ${VLLM_ARGS[@]} \
-    ${RM_ARGS[@]} \
-    ${MISC_ARGS[@]} \
-    2>&1 | tee "${LOG_DIR}/opd_training.log"
-
-echo "=== Training complete, stopping teacher server ==="
-kill "${TEACHER_PID}" 2>/dev/null || true
-ray stop --force 2>/dev/null || true
-pkill -9 ray 2>/dev/null || true
-pkill -9 -f "train.py" 2>/dev/null || true
-echo "=== Done ==="
diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
new file mode 100644
index 000000000..9deb914f6
--- /dev/null
+++ b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
@@ -0,0 +1,167 @@
+#!/bin/bash
+
+# usage: bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
+#
+# On-Policy Distillation with Megatron-based teacher model.
+# This example uses the original model as the teacher (self-distillation for demonstration).
+#
+# IMPORTANT: This is just an example configuration!
+# In practice, you should:
+# 1. Use a different (stronger) model as the teacher
+# 2. Adjust --opd-kl-coef based on your task
+# 3. Prefer --opd-type vllm (run-qwen3-8B-opd.sh) when the teacher is larger
+#
+# Memory note (8×A800 80GB): student + teacher both Qwen3-8B on 4 training GPUs
+# is tight. Defaults use TP=4, full recompute, and lower token budget.
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+source "/root/vime/scripts/models/qwen3-8B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/Qwen3-8B
+   --ref-load /root/Qwen3-8B_torch_dist
+   --load /root/Qwen3-8B_torch_dist
+   --save /root/Qwen3-8B_vime/
+   --save-interval 10
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+   --num-rollout 300
+   --rollout-batch-size 16
+   --n-samples-per-prompt 4
+   --rollout-max-response-len 16384
+   --rollout-temperature 1
+
+   --global-batch-size 64
+   --balance-data
+)
+
+RM_ARGS=(
+   --rm-type math
+)
+
+EVAL_ARGS=(
+   # --eval-interval 20
+   # --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+   # --n-samples-per-eval-prompt 16
+   # --eval-max-response-len 16384
+   # --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 4
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 36
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   # OPD Configuration
+   --use-opd
+   --opd-type megatron
+   --opd-kl-coef 1.0
+   # CHANGE THIS to a stronger teacher checkpoint in practice
+   --opd-teacher-load /root/Qwen3-8B_torch_dist
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+   #--use-wandb
+   # --wandb-project vime-opd
+   # --wandb-group qwen3-8B-opd-megatron
+   # --wandb-key ${WANDB_KEY}
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --vllm-gpu-memory-utilization 0.4
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --make-vocab-size-divisible-by 128
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
+    \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --working-dir /root/vime \
+   -- python3 train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 4 \
+   --rollout-num-gpus 4 \
+   --colocate \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${RM_ARGS[@]}
+
+#### clear after training
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 -f "train.py" || true
+sleep 3
diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd.sh b/examples/on_policy_distillation/run-qwen3-8B-opd.sh
new file mode 100644
index 000000000..e9bca96b8
--- /dev/null
+++ b/examples/on_policy_distillation/run-qwen3-8B-opd.sh
@@ -0,0 +1,211 @@
+#!/bin/bash
+
+# usage: bash examples/on_policy_distillation/run-qwen3-8B-opd.sh
+#
+# OPD with external vLLM teacher: Qwen3-8B student + Qwen3-32B teacher.
+# GPU layout (8×GPU): 0-3 student train+rollout (colocate), 4-7 teacher vLLM (TP=4).
+# Hyperparams below were validated on 8×A800 80GB (GSM8K +8.5 pp).
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+# Start the teacher model server
+TEACHER_IP="127.0.0.1"
+TEACHER_PORT=13141
+LOG_FILE="/tmp/vllm_teacher_$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 6).log"
+
+## Launch the teacher model server in the background (GPUs 4-7)
+CUDA_VISIBLE_DEVICES=4,5,6,7 python3 -m vllm.entrypoints.openai.api_server \
+    --model /root/Qwen3-32B \
+    --host 0.0.0.0 \
+    --port ${TEACHER_PORT} \
+    --tensor-parallel-size 4 \
+    --gpu-memory-utilization 0.85 \
+    --trust-remote-code \
+    --dtype bfloat16 \
+    --max-model-len 16384 \
+    --disable-custom-all-reduce \
+    > "${LOG_FILE}" 2>&1 &
+TEACHER_PID=$!
+
+echo "Starting teacher model server (pid=${TEACHER_PID})..."
+
+## Wait for the teacher model server to be ready (/health returns empty body)
+for i in $(seq 1 120); do
+    if ! kill -0 "${TEACHER_PID}" 2>/dev/null; then
+        echo "ERROR: Teacher server process died. Check ${LOG_FILE}"
+        tail -n 20 "${LOG_FILE}"
+        exit 1
+    fi
+    HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" "http://${TEACHER_IP}:${TEACHER_PORT}/health" 2>/dev/null || true)
+    if [ "${HTTP_CODE}" = "200" ]; then
+        echo "Teacher model server is up and running at ${TEACHER_IP}:${TEACHER_PORT}."
+        break
+    fi
+    if [ "$i" -eq 120 ]; then
+        echo "ERROR: Teacher server failed to start within 10 minutes"
+        tail -n 20 "${LOG_FILE}"
+        kill "${TEACHER_PID}" 2>/dev/null || true
+        exit 1
+    fi
+    echo "Waiting for the teacher model server to start..."
+    sleep 5
+done
+sleep 5
+
+source "/root/vime/scripts/models/qwen3-8B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/Qwen3-8B
+   --ref-load /root/Qwen3-8B_torch_dist
+   # First run: load from converted torch_dist. For resume, set --load to --save path.
+   --load /root/Qwen3-8B_torch_dist
+   --save /root/Qwen3-8B_vime/
+   --save-interval 20
+   --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type math
+   --num-rollout 300
+   --rollout-batch-size 16
+   --n-samples-per-prompt 4
+   # 4096 response / 8192 context: validated for 8B colocate on 80GB GPUs.
+   # Longer responses need more memory (raise carefully with max-tokens-per-gpu).
+   --rollout-max-response-len 4096
+   --rollout-max-context-len 8192
+   --rollout-temperature 1
+
+   --global-batch-size 64
+   --balance-data
+)
+
+RM_ARGS=(
+   --custom-rm-path vime.rollout.on_policy_distillation.reward_func
+   --custom-reward-post-process-path vime.rollout.on_policy_distillation.post_process_rewards
+   --rm-url http://${TEACHER_IP}:${TEACHER_PORT}/inference/v1/generate
+)
+
+EVAL_ARGS=(
+   # In-training eval is currently incompatible with OPD reward_func (returns a
+   # dict; eval logging expects scalar rewards). Prefer offline eval after training.
+   # --eval-interval 50
+   # --eval-prompt-data gsm8k /root/gsm8k/test.parquet
+   # --eval-input-key messages
+   # --n-samples-per-eval-prompt 1
+   # --eval-max-response-len 4096
+   # --eval-top-k 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 2
+   --sequence-parallel
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 2048
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-opd
+   --opd-type vllm
+   --opd-kl-coef 1.0
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+WANDB_ARGS=(
+   #--use-wandb
+   # --wandb-project vime-opd
+   # --wandb-group qwen3-8B-opd
+   # --wandb-key ${WANDB_KEY}
+)
+
+VLLM_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --vllm-gpu-memory-utilization 0.25
+)
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+   --make-vocab-size-divisible-by 128
+)
+
+# Student uses GPUs 0-3 (teacher already occupies 4-7)
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --working-dir /root/vime \
+   -- python3 train.py \
+   --train-backend megatron \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 4 \
+   --colocate \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${VLLM_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${RM_ARGS[@]}
+
+#### clear after training
+kill ${TEACHER_PID} 2>/dev/null || true
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 -f "train.py" || true
+sleep 3
diff --git a/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
deleted file mode 100644
index d211408ea..000000000
--- a/examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
+++ /dev/null
@@ -1,159 +0,0 @@
-#!/bin/bash
-
-# usage: bash examples/on_policy_distillation/run-qwen3-8b-opd-megatron.sh
-#
-# OPD with a Megatron-loaded teacher (no external vLLM teacher server).
-# Student rollout still uses vLLM. This demo uses the same architecture for
-# student and teacher — replace --opd-teacher-load with a stronger checkpoint
-# in practice.
-#
-# Memory note (8×A800 80GB): loading student + teacher (both Qwen3-8B) on
-# 4 training GPUs is tight. The defaults below match a validated config:
-# TP=4, full activation recompute, max-tokens-per-gpu=8192, and
-# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. See README FAQ.
-
-set -ex
-
-export PYTHONUNBUFFERED=1
-export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
-
-NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
-if [ "$NVLINK_COUNT" -gt 0 ]; then
-    HAS_NVLINK=1
-else
-    HAS_NVLINK=0
-fi
-echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
-
-ray stop --force 2>/dev/null || true
-pkill -9 ray 2>/dev/null || true
-pkill -9 -f "train.py" 2>/dev/null || true
-sleep 3
-
-source "/root/vime/scripts/models/qwen3-8B.sh"
-
-CKPT_ARGS=(
-    --hf-checkpoint /root/models/Qwen3-8B
-    --ref-load /root/models/Qwen3-8B_torch_dist
-    --load /root/models/Qwen3-8B_torch_dist
-    --save /root/Qwen3-8B_opd/
-    --save-interval 10
-    --megatron-to-hf-mode bridge
-)
-
-ROLLOUT_ARGS=(
-    --prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl
-    --input-key prompt
-    --apply-chat-template
-    --rollout-shuffle
-    --rm-type math
-    --num-rollout 300
-    --rollout-batch-size 16
-    --n-samples-per-prompt 4
-    --rollout-max-response-len 16384
-    --rollout-temperature 1
-    --global-batch-size 64
-    --balance-data
-)
-
-EVAL_ARGS=(
-    --eval-interval 50
-    --eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet
-    --n-samples-per-eval-prompt 1
-    --eval-max-response-len 4096
-    --eval-top-k 1
-)
-
-PERF_ARGS=(
-    # TP=4 + full recompute + lower token budget: needed so student+teacher
-    # (both Qwen3-8B) fit on 4×80GB training GPUs.
-    --tensor-model-parallel-size 4
-    --sequence-parallel
-    --pipeline-model-parallel-size 1
-    --context-parallel-size 1
-    --expert-model-parallel-size 1
-    --expert-tensor-parallel-size 1
-    --recompute-granularity full
-    --recompute-method uniform
-    --recompute-num-layers 36
-    --use-dynamic-batch-size
-    --max-tokens-per-gpu 8192
-)
-
-GRPO_ARGS=(
-    --advantage-estimator grpo
-    --use-opd
-    --opd-type megatron
-    --opd-kl-coef 1.0
-    --opd-teacher-load /root/models/Qwen3-8B_torch_dist
-    --use-kl-loss
-    --kl-loss-coef 0.00
-    --kl-loss-type low_var_kl
-    --entropy-coef 0.00
-)
-
-OPTIMIZER_ARGS=(
-    --optimizer adam
-    --lr 1e-6
-    --lr-decay-style constant
-    --weight-decay 0.1
-    --adam-beta1 0.9
-    --adam-beta2 0.98
-)
-
-WANDB_ARGS=(
-    # --use-wandb
-    # --wandb-project vime-opd
-    # --wandb-group qwen3-8b-opd-megatron
-    # --wandb-key ${WANDB_KEY}
-)
-
-VLLM_ARGS=(
-    --rollout-num-gpus-per-engine 1
-    --vllm-gpu-memory-utilization 0.4
-)
-
-MISC_ARGS=(
-    --attention-dropout 0.0
-    --hidden-dropout 0.0
-    --accumulate-allreduce-grads-in-fp32
-    --attention-softmax-in-fp32
-    --attention-backend flash
-    --make-vocab-size-divisible-by 128
-)
-
-export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
-ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus 8 \
-    --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-
-RUNTIME_ENV_JSON="{
-  \"env_vars\": {
-    \"PYTHONPATH\": \"/root/vime:/root/Megatron-LM/\",
-    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
-    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
-    \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\"
-  }
-}"
-
-ray job submit --address="http://127.0.0.1:8265" \
-    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-    -- python3 train.py \
-    --train-backend megatron \
-    --actor-num-nodes 1 \
-    --actor-num-gpus-per-node 4 \
-    --rollout-num-gpus 4 \
-    --colocate \
-    ${MODEL_ARGS[@]} \
-    ${CKPT_ARGS[@]} \
-    ${ROLLOUT_ARGS[@]} \
-    ${OPTIMIZER_ARGS[@]} \
-    ${GRPO_ARGS[@]} \
-    ${WANDB_ARGS[@]} \
-    ${PERF_ARGS[@]} \
-    ${EVAL_ARGS[@]} \
-    ${VLLM_ARGS[@]} \
-    ${MISC_ARGS[@]}
-
-ray stop --force 2>/dev/null || true
-pkill -9 ray 2>/dev/null || true
-pkill -9 -f "train.py" 2>/dev/null || true

From 058a41c6b7265dfa6d549a3f078ba4541d76dc10 Mon Sep 17 00:00:00 2001
From: kaiyuan <kyxiezju@163.com>
Date: Sat, 18 Jul 2026 19:19:54 +0800
Subject: [PATCH 4/4] Trim OPD script comments to match slime style.

Signed-off-by: kaiyuan <kyxiezju@163.com>
---
 examples/on_policy_distillation/README.md     | 19 +++----------------
 .../run-qwen3-8B-opd-megatron.sh              | 11 +++--------
 .../run-qwen3-8B-opd.sh                       | 15 +++------------
 3 files changed, 9 insertions(+), 36 deletions(-)

diff --git a/examples/on_policy_distillation/README.md b/examples/on_policy_distillation/README.md
index 257db6287..9d85d3e69 100644
--- a/examples/on_policy_distillation/README.md
+++ b/examples/on_policy_distillation/README.md
@@ -77,7 +77,7 @@ PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
 bash examples/on_policy_distillation/run-qwen3-8B-opd.sh
 ```
 
-GPU layout (8×GPU):
+GPU layout:
 
 | GPUs | Role |
 |------|------|
@@ -115,9 +115,8 @@ bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
 
 # Preliminary Results
 
-End-to-end run with `run-qwen3-8B-opd.sh` on 8×A800 80GB (dapo-math-17k train,
-GRPO + `--opd-kl-coef 1.0`, ~220 rollouts / iter_0000219). Offline GSM8K greedy
-eval:
+End-to-end run with `run-qwen3-8B-opd.sh` (dapo-math-17k train, GRPO +
+`--opd-kl-coef 1.0`, ~220 rollouts / iter_0000219). Offline GSM8K greedy eval:
 
 | Model | GSM8K Accuracy |
 |-------|----------------|
@@ -128,18 +127,6 @@ eval:
 Training health signal: `rollout/opd_reverse_kl` dropped from 0.145 → ~0.10
 (−38%). Pure OPD uses `raw_reward=0`; the learning signal is the OPD KL term.
 
-Notes from the validated run:
-
-- Colocate memory is tight for 8B+32B; keep `--rollout-max-response-len 4096`,
-  `--rollout-max-context-len 8192`, `--max-tokens-per-gpu 2048`, and
-  `--vllm-gpu-memory-utilization 0.25` unless you have more headroom.
-- Teacher vLLM needs `--max-model-len 16384` and `--disable-custom-all-reduce`
-  for TP=4 stability.
-- Prefer offline eval after training; in-training eval is currently incompatible
-  with OPD's custom reward payload.
-- Megatron checkpoints are large (~100GB each); plan disk or lower
-  `--save-interval`.
-
 # FAQ
 
 1. **Why are there two OPD modes?**
diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
index 9deb914f6..22a0c1b2e 100644
--- a/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
+++ b/examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
@@ -1,18 +1,13 @@
 #!/bin/bash
 
-# usage: bash examples/on_policy_distillation/run-qwen3-8B-opd-megatron.sh
-#
-# On-Policy Distillation with Megatron-based teacher model.
-# This example uses the original model as the teacher (self-distillation for demonstration).
+# On-Policy Distillation with Megatron-based teacher model
+# This example uses the original model as the teacher (self-distillation for demonstration)
 #
 # IMPORTANT: This is just an example configuration!
 # In practice, you should:
 # 1. Use a different (stronger) model as the teacher
 # 2. Adjust --opd-kl-coef based on your task
-# 3. Prefer --opd-type vllm (run-qwen3-8B-opd.sh) when the teacher is larger
-#
-# Memory note (8×A800 80GB): student + teacher both Qwen3-8B on 4 training GPUs
-# is tight. Defaults use TP=4, full recompute, and lower token budget.
+# 3. Configure proper evaluation metrics
 
 set -ex
 
diff --git a/examples/on_policy_distillation/run-qwen3-8B-opd.sh b/examples/on_policy_distillation/run-qwen3-8B-opd.sh
index e9bca96b8..c02f6bc79 100644
--- a/examples/on_policy_distillation/run-qwen3-8B-opd.sh
+++ b/examples/on_policy_distillation/run-qwen3-8B-opd.sh
@@ -1,10 +1,6 @@
 #!/bin/bash
 
 # usage: bash examples/on_policy_distillation/run-qwen3-8B-opd.sh
-#
-# OPD with external vLLM teacher: Qwen3-8B student + Qwen3-32B teacher.
-# GPU layout (8×GPU): 0-3 student train+rollout (colocate), 4-7 teacher vLLM (TP=4).
-# Hyperparams below were validated on 8×A800 80GB (GSM8K +8.5 pp).
 
 set -ex
 
@@ -23,7 +19,7 @@ TEACHER_IP="127.0.0.1"
 TEACHER_PORT=13141
 LOG_FILE="/tmp/vllm_teacher_$(head /dev/urandom | tr -dc A-Za-z0-9 | head -c 6).log"
 
-## Launch the teacher model server in the background (GPUs 4-7)
+## Launch the teacher model server in the background
 CUDA_VISIBLE_DEVICES=4,5,6,7 python3 -m vllm.entrypoints.openai.api_server \
     --model /root/Qwen3-32B \
     --host 0.0.0.0 \
@@ -39,7 +35,7 @@ TEACHER_PID=$!
 
 echo "Starting teacher model server (pid=${TEACHER_PID})..."
 
-## Wait for the teacher model server to be ready (/health returns empty body)
+## Wait for the teacher model server to be ready
 for i in $(seq 1 120); do
     if ! kill -0 "${TEACHER_PID}" 2>/dev/null; then
         echo "ERROR: Teacher server process died. Check ${LOG_FILE}"
@@ -67,7 +63,6 @@ source "/root/vime/scripts/models/qwen3-8B.sh"
 CKPT_ARGS=(
    --hf-checkpoint /root/Qwen3-8B
    --ref-load /root/Qwen3-8B_torch_dist
-   # First run: load from converted torch_dist. For resume, set --load to --save path.
    --load /root/Qwen3-8B_torch_dist
    --save /root/Qwen3-8B_vime/
    --save-interval 20
@@ -83,8 +78,6 @@ ROLLOUT_ARGS=(
    --num-rollout 300
    --rollout-batch-size 16
    --n-samples-per-prompt 4
-   # 4096 response / 8192 context: validated for 8B colocate on 80GB GPUs.
-   # Longer responses need more memory (raise carefully with max-tokens-per-gpu).
    --rollout-max-response-len 4096
    --rollout-max-context-len 8192
    --rollout-temperature 1
@@ -100,8 +93,6 @@ RM_ARGS=(
 )
 
 EVAL_ARGS=(
-   # In-training eval is currently incompatible with OPD reward_func (returns a
-   # dict; eval logging expects scalar rewards). Prefer offline eval after training.
    # --eval-interval 50
    # --eval-prompt-data gsm8k /root/gsm8k/test.parquet
    # --eval-input-key messages
@@ -169,7 +160,7 @@ MISC_ARGS=(
    --make-vocab-size-divisible-by 128
 )
 
-# Student uses GPUs 0-3 (teacher already occupies 4-7)
+# launch the master node of ray in container
 export CUDA_VISIBLE_DEVICES=0,1,2,3
 export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
 ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 4 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265