Add On-Policy Self-Distillation (OPSD) via --opd-type self

.github/workflows/pr-test.yml

@@ -123,7 +123,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}]
+        info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_opsd.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -506,7 +506,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}]
+        info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_opsd.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}

.github/workflows/pr-test.yml.j2

@@ -14,6 +14,7 @@
     {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8},
     {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8},
     {'test_file': 'test_qwen2.5_0.5B_opd_sglang.py', 'num_gpus': 8},
+    {'test_file': 'test_qwen2.5_0.5B_opsd.py', 'num_gpus': 4},
     {'test_file': 'test_qwen3_4B_external_pd.py', 'num_gpus': 6},
     {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4},
     {'test_file': 'test_qwen2.5_0.5B_fanout_short.py', 'num_gpus': 4},

docs/en/advanced/on-policy-distillation.md

@@ -69,6 +69,49 @@ The teacher model is loaded directly into Megatron via `--opd-teacher-load`. Tea
 
 > **Note**: The teacher checkpoint must be in Megatron format (`torch_dist` or `torch`). You can convert from HuggingFace format using `tools/convert_hf_to_torch_dist.py`.
 
+### Self Mode (`--opd-type self`) — On-Policy Self-Distillation (OPSD)
+
+This mode implements **On-Policy Self-Distillation** ([Self-Distilled Reasoner](https://github.com/siyan-zhao/OPSD)). A single model acts as both student and teacher, differing only by *context*:
+
+- **Student**: conditioned on the problem only; generates the on-policy rollout (the trainable current policy).
+- **Teacher**: the **same model, frozen at the initial-policy checkpoint** (`--opd-teacher-load`), conditioned on the problem **plus privileged information** (the ground-truth solution). The teacher does not generate — it scores the student's response tokens in a single forward pass.
+
+Unlike the other modes (which fold a sampled-token reverse-KL into the advantage), OPSD uses a **direct, full-vocab token-level Jensen-Shannon divergence (JSD)** as the loss, with **no task reward** (pure distillation):
+
+$$
+\mathcal{L}_{\text{OPSD}} = \mathbb{E}_{t}\Big[\min\big(\text{JSD}_\beta\big(P_{\text{teacher}}(\cdot|x, s, y_{<t})\,\|\,P_{\text{student}}(\cdot|x, y_{<t})\big),\ c\big)\Big]
+$$
+
+where $x$ is the problem, $s$ the privileged solution, $\beta$ the JSD interpolation weight (`--opsd-beta`: 0 → forward KL, 1 → reverse KL, 0.5 → symmetric), and $c$ the per-token clip (`--opsd-jsd-clip`, default 0.05) that prevents high-divergence style tokens (e.g. "wait", "think") from dominating the gradient.
+
+**Key arguments**:
+
+| Argument | Description |
+|----------|-------------|
+| `--opsd-beta` | JSD interpolation weight (default 0.5). |
+| `--opsd-jsd-clip` | Per-token JSD clamp (default 0.05). |
+| `--opsd-privileged-info-key` | Dataset field with the privileged information (ground-truth solution) shown only to the teacher. **Required**. |
+
+**Data format**: the reference release uses [`siyanzhao/Openthoughts_math_30k_opsd`](https://huggingface.co/datasets/siyanzhao/Openthoughts_math_30k_opsd) with `problem` (the question, `--input-key problem`) and `solution` (the ground-truth solution, `--opsd-privileged-info-key solution`).
+
+**How it works**:
+1. The teacher (frozen init checkpoint) is loaded as an additional Megatron model.
+2. During rollout conversion, the privileged teacher sequence `[prompt + privileged_info + response]` is built (the response is a verbatim copy of the student's rollout).
+3. Before training, the teacher is forwarded on that privileged sequence to obtain response-position logits.
+4. During the student training forward, the per-token full-vocab JSD between student and teacher logits is computed, clamped, and reduced into the loss.
+
+**Configuration**:
+```bash
+--use-opd
+--opd-type self
+--opd-teacher-load /path/to/init_torch_dist   # frozen teacher = initial policy
+--opsd-beta 0.5
+--opsd-jsd-clip 0.05
+--opsd-privileged-info-key solution           # dataset field with the ground-truth solution
+```
+
+> **Limitations (current MVP)**: OPSD requires `--context-parallel-size 1`. The teacher response logits are held over the full vocabulary between the teacher and student forwards, so memory scales with `response_length × vocab_size`; pass `--opsd-offload-teacher-logits` to offload them to CPU (trading GPU memory for host↔device transfer), and prefer smaller models / shorter responses for now. A future chunked JSD would avoid materializing them. The teacher's micro-batches are re-packed by the (longer) privileged sequence lengths so the teacher forward respects `--max-tokens-per-gpu`. The example script is `examples/on_policy_distillation/run-qwen3-8B-opsd.sh`.
+
 ## Running the Examples
 
 Complete example scripts are provided in `examples/on_policy_distillation/`:

examples/on_policy_distillation/run-qwen3-8B-opsd.sh

@@ -0,0 +1,172 @@
+#!/bin/bash
+
+# On-Policy Self-Distillation (OPSD), following "Self-Distilled Reasoner".
+#
+# A single model acts as both student and teacher:
+#   - Student: conditioned on the problem only; generates the on-policy rollout.
+#   - Teacher: the SAME model, frozen at the init checkpoint (--opd-teacher-load),
+#     conditioned on the problem + privileged information (the ground-truth solution).
+# The teacher scores the student's response tokens in one forward pass, and the loss
+# is a token-level full-vocab JSD between teacher and student over the student's tokens.
+#
+# Requirements:
+#   - --opd-type self   (Megatron-only; context-parallel-size must be 1)
+#   - --opsd-privileged-info-key names the dataset field holding the ground-truth solution
+#   - --opd-teacher-load points at the initial-policy checkpoint (frozen teacher)
+#   - No task reward is used (pure distillation).
+
+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)"
+
+source "/root/slime/scripts/models/qwen3-8B.sh"
+
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/Qwen3-8B
+   --ref-load /root/Qwen3-8B_torch_dist
+   --load /root/Qwen3-8B_slime/
+   --save /root/Qwen3-8B_slime/
+   --save-interval 20
+)
+
+# Data: the OPSD release uses siyanzhao/Openthoughts_math_30k_opsd, with fields
+#   "problem"  -> the question (student + teacher context)
+#   "solution" -> the ground-truth solution (privileged info, teacher-only)
+# Download and convert to jsonl, e.g.:
+#   hf download --repo-type dataset siyanzhao/Openthoughts_math_30k_opsd --local-dir /root/opsd-math-30k
+ROLLOUT_ARGS=(
+   --prompt-data /root/opsd-math-30k/train.jsonl
+   --input-key problem
+   --apply-chat-template
+   --rollout-shuffle
+   --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
+)
+
+# OPSD is pure distillation; the reward is computed but unused (no advantages).
+RM_ARGS=(
+   --rm-type math
+)
+
+EVAL_ARGS=(
+   # --eval-interval 20
+   # --eval-prompt-data aime ${DATA_DIR}/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                          # OPSD requires context-parallel-size 1
+   --expert-model-parallel-size 1
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   # --micro-batch-size 1
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 16384
+)
+
+OPSD_ARGS=(
+   # On-Policy Self-Distillation
+   --use-opd                                           # Enable on-policy distillation
+   --opd-type self                                     # OPSD: single model is both student and teacher
+   --opd-teacher-load /root/Qwen3-8B_torch_dist        # Frozen teacher = initial-policy checkpoint
+   --opsd-beta 0.5                                      # JSD interpolation (0=fwd KL, 1=rev KL, 0.5=symmetric)
+   --opsd-jsd-clip 0.05                                # Per-token JSD clamp (stops style tokens dominating)
+   --opsd-privileged-info-key solution                 # Dataset field with the ground-truth solution (teacher-only)
+)
+
+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 slime-dev
+   # --wandb-group qwen3-8B-opsd
+   # --wandb-key ${WANDB_KEY}
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --sglang-mem-fraction-static 0.4
+)
+
+
+MISC_ARGS=(
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   --attention-backend flash
+)
+
+
+
+
+# launch the master node of ray in container
+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
+
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json='{
+     "env_vars": {
+        "PYTHONPATH": "/root/Megatron-LM/",
+        "CUDA_DEVICE_MAX_CONNECTIONS": "1"
+     }
+   }' \
+   -- python3 train.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 2 \
+   --rollout-num-gpus 4 \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${OPSD_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${RM_ARGS[@]}
+
+
+
+####clear after training
+pkill -9 sglang
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python