[OPD] Add Qwen3.5-35B-A3B single-node self-distillation example

examples/on_policy_distillation/qwen3_5_35b_selfdistill/README.md

@@ -0,0 +1,210 @@
+# Qwen3.5-35B-A3B Self-Distillation on a Single Node (RLVR teacher → OPD)
+
+A reproducible two-phase on-policy-distillation (OPD) example for the
+**Qwen3.5-35B-A3B** MoE on a **single 8×H200 node**, using the **in-process
+Megatron teacher** (`--opd-type megatron`, no separate teacher server).
+
+It differs from the sibling examples in three ways:
+
+1. **Real MoE at scale on one node.** The 2-node/16-GPU recipe is re-tiled to 8 GPUs.
+2. **A genuinely diverged teacher.** `run-qwen3-8B-opd-megatron.sh` uses `teacher == base`
+   (a mechanism demo where the reverse-KL is ~0). Here Phase 1 *trains* the teacher
+   with RLVR so it is measurably better and more concise than the base — the
+   prerequisite for OPD to actually move the student.
+3. **Self-distillation is the only valid option here.** Qwen3.5 has its own tokenizer
+   (vocab 248320); the smaller Qwen3 models (vocab 151936) are not token-compatible,
+   so a cross-model teacher would be invalid. Teacher and student are the same family.
+
+## Pipeline
+
+```
+Phase 1  (phase1_rlvr_teacher.sh)         Phase 2  (phase2_opd_selfdistill.sh)
+base 35B  --RLVR (GRPO, lr 1e-5)-->  teacher        base 35B (student)
+          better + concise                     |  <-- reverse-KL (--opd-type megatron)
+          (eval 0.83 -> 0.89)                  teacher (Phase-1 ckpt, in-process)
+```
+
+## Single-node parallelism (world = 8)
+
+The original recipe was 2 nodes × 8 GPUs (`TP2 PP1 CP2 EP8 ETP1`, DP4). On one node
+we keep the same dims and only halve DP:
+
+| dim | value | check |
+|-----|-------|-------|
+| TP  | 2 | decoder `TP*PP*CP = 2` ; `8 % 2 = 0` → DP = 4 |
+| PP  | 1 | |
+| CP  | 2 | shards the long (~17k) sequence so 24k context fits |
+| EP  | 8 | `num_experts 256 % 8 = 0` ; expert `ETP*EP*PP = 8` → expert_dp = 1 |
+| ETP | 1 | expert_dp(1) ≠ dp(4) is allowed (miles rank order ends in `pp`) |
+
+`--colocate` time-shares the train and rollout phases (each fits 143 GB separately,
+not summed); `--optimizer-cpu-offload` puts Adam state on host RAM; the model is a
+hybrid linear-attention MoE so the KV cache is small. Peak ≈ 124 GB / 143 GB per GPU.
+
+## Reproduce
+
+**0. Prereqs** — model + torch_dist checkpoint, and the train/eval split:
+
+```bash
+# model (and mcore conversion, see ../README.md for convert_hf_to_torch_dist usage)
+#   ${MODEL_DIR}/Qwen3.5-35B-A3B  and  ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist
+# disjoint, seeded train/eval split (eval is held out from BOTH phases):
+python make_split.py --src /path/to/dapo-math-17k.jsonl --out-dir ${DATA_DIR}
+#   -> ${DATA_DIR}/dapo_train.jsonl (16886)  ${DATA_DIR}/dapo_eval.jsonl (512)
+```
+
+**1. Phase 1 — train the teacher** (watch `rollout/raw_reward` climb and
+`eval/dapo_heldout` rise above the base ~0.83):
+
+```bash
+MODEL_DIR=... DATA_DIR=... OUTPUT_DIR=/persistent/ckpt-teacher \
+  bash phase1_rlvr_teacher.sh
+```
+
+**2. Phase 2 — distill the teacher into the base student**:
+
+```bash
+# pure OPD (default): training reward = 0, only the teacher reverse-KL drives learning
+TEACHER_LOAD=/persistent/ckpt-teacher DATA_DIR=... \
+  bash phase2_opd_selfdistill.sh
+
+# grounded OPD: correctness reward (raw_reward == accuracy, climbs) + teacher reverse-KL
+MODE=grounded TEACHER_LOAD=/persistent/ckpt-teacher DATA_DIR=... \
+  bash phase2_opd_selfdistill.sh
+```
+
+`OUTPUT_DIR` / the teacher checkpoint must live on **persistent** storage. On a
+KubeRay pod the head can be recreated and wipe the container overlay (`/root`); a
+node-local disk (e.g. `/node_public`) survives and makes runs resumable.
+
+## Run on GB200 / GB300 (CUDA 13, Blackwell) — `phase2_gb200.sh`
+
+The recipe above targets a single **8×H200** node. Blackwell nodes (GB200/GB300)
+have **4 GPUs/node**, so `world = 8` becomes **2 nodes × 4 GPUs** — same parallel
+dims (`TP2 PP1 CP2 EP8 ETP1`, DP4), only the node tiling changes. `phase2_gb200.sh`
+is the GB200 variant of `phase2_opd_selfdistill.sh`; the deltas (all validated on
+2× GB200, reproducing the base eval `0.84` / `~14k`) are:
+
+- **Tiling** — `--actor-num-nodes 2 --num-gpus-per-node 4` (override via
+  `ACTOR_NUM_NODES` / `GPUS_PER_NODE`). Pin both nodes to one NVLink (MNNVL) domain
+  so the EP8 all-to-all stays on the NVLink fabric.
+- **sglang backends** (cf. `scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py`) —
+  `--sglang-moe-runner-backend flashinfer_cutlass`, `--sglang-attention-backend
+  trtllm_mha`, and `--moe-token-dispatcher-type flex`. The default triton fused-MoE
+  mis-shards routed experts on the megatron→sglang weight sync
+  (`fused_moe_triton ... _load_w13`: `tensor a (64) vs b (2048)`), and FA3 is SM≤90
+  only (Blackwell is SM 10.x).
+- **NCCL** — `NCCL_NVLS_ENABLE=0` (multi-node Blackwell NVLS bind fails
+  `ncclCommInitRank`); keep `NCCL_MNNVL_ENABLE=1`.
+- **k8s** — if a `prometheus` Service exists in the namespace, set
+  `PROMETHEUS_PORT=9090` (kube injects a `tcp://…:9090` URL that breaks miles'
+  `int(PROMETHEUS_PORT)`).
+
+`phase2_gb200.sh` already sets the sglang/MoE backends and folds
+`NCCL_NVLS_ENABLE=0` + `PROMETHEUS_PORT=9090` into the Ray runtime env. Run it on the
+Ray head in the CUDA-13 ARM64 miles image, with `MILES_DIR` pointing at the repo:
+
+```bash
+ACTOR_NUM_NODES=2 GPUS_PER_NODE=4 MILES_DIR=/workspace/miles \
+MODEL_DIR=... DATA_DIR=... TEACHER_LOAD=/persistent/ckpt-teacher OUTPUT_DIR=/persistent/ckpt-opd-pure \
+  bash phase2_gb200.sh          # MODE=pure (default) | MODE=grounded
+```
+
+## Run Phase 2 only (skip Phase 1)
+
+If you already have a teacher checkpoint, skip Phase 1 and run Phase 2 directly —
+point `--opd-teacher-load` (`TEACHER_LOAD`) at the teacher's **torch_dist parent
+dir** (the one containing `latest_checkpointed_iteration.txt`). You still need the
+base model (`--hf-checkpoint` + the `--ref-load` torch_dist) and the data split, but
+no Phase-1 run.
+
+If your teacher is in **HuggingFace** format, convert it to torch_dist first with
+`convert_gb200.sh` (a thin wrapper over `tools/convert_hf_to_torch_dist.py` carrying
+the Qwen3.5 `MODEL_ARGS`):
+
+```bash
+# teacher: HF safetensors -> Megatron torch_dist parent dir
+bash convert_gb200.sh /path/to/teacher-hf /persistent/ckpt-teacher
+# (and the base, if you don't have Qwen3.5-35B-A3B_torch_dist yet)
+bash convert_gb200.sh ${MODEL_DIR}/Qwen3.5-35B-A3B ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist
+
+TEACHER_LOAD=/persistent/ckpt-teacher MODEL_DIR=... DATA_DIR=... \
+  bash phase2_gb200.sh          # or phase2_opd_selfdistill.sh on 8×H200
+```
+
+> **Teacher expert layout.** The public `Qwen/Qwen3.5-35B-A3B` ships *fused* experts
+> (`mlp.experts.gate_up_proj`); a teacher round-tripped through
+> `convert_torch_dist_to_hf` may ship *unfused* per-expert weights
+> (`mlp.experts.{i}.gate_proj.weight`). `miles_plugins/mbridge/qwen3_5.py` now
+> autodetects both for the main layers (mirroring the existing MTP-expert
+> autodetect), so either layout converts without manual re-fusing.
+
+## Results (DAPO-math, held-out 512, eval @ 24k cap, temp 0.6)
+
+**Phase 1 — RLVR teacher** (lr 1e-5):
+
+| step | eval/dapo_heldout | eval response length |
+|------|-------------------|----------------------|
+| 0 (base) | 0.828 | 14,070 |
+| 5        | **0.887** | **6,248** |
+
+The teacher becomes both more accurate **and** ~2× more concise. This Phase-1
+teacher checkpoint is published at
+[**cm00cm/Qwen3.5-35B-A3B-DAPO-RLVR-teacher**](https://huggingface.co/cm00cm/Qwen3.5-35B-A3B-DAPO-RLVR-teacher)
+(weights only) and can be used directly as the Phase-2 teacher via
+`--opd-teacher-load` after `convert_hf_to_torch_dist.py`.
+
+**Phase 2 — pure OPD** (student = base, teacher = Phase-1 step-5 ckpt; reward = 0):
+
+| step | eval/dapo_heldout | eval response length | opd_reverse_kl |
+|------|-------------------|----------------------|----------------|
+| 0 (base) | 0.840 | 14,070 | — |
+| 5        | 0.852 | **6,132** | 0.045 → 0.013 |
+
+With **zero task reward**, pure reverse-KL distillation transfers the teacher's
+concise behavior to the base student — eval length **−57%** with accuracy
+preserved/slightly up (the +1.2 pt is within the ~1.6 pt eval SE; the robust,
+headline effect is the efficiency transfer). A nonzero, shrinking `opd_reverse_kl`
+confirms the teacher genuinely differs from the student and the student is
+converging onto it.
+
+**Phase 2 — grounded OPD** (correctness reward + teacher reverse-KL):
+
+| step | rollout/raw_reward | train length | opd_reverse_kl |
+|------|--------------------|--------------|----------------|
+| 1 | 0.637 | 18,778 | 0.045 |
+| 2 | **0.910** | **7,665** | 0.014 |
+
+With the correctness reward kept, `rollout/raw_reward` (== accuracy) climbs while
+the student simultaneously adopts the teacher's concise responses (18.8k → 7.7k).
+The shrinking `opd_reverse_kl` (0.045 → 0.014) shows the student converging onto
+the teacher. (At lr 1e-5 the RLVR reward alone also drives accuracy up — Phase 1
+is the controlled view of that — so grounded OPD's `raw_reward` climb reflects
+RLVR + the teacher pull combined; the pure-OPD run above isolates OPD's effect.)
+
+## Gotchas (each cost a wasted run to find)
+
+- **Reward grader.** `--rm-type deepscaler` requires a `</think>` tag and returns 0
+  otherwise; Qwen3.5 reasons inline (no tag) → every reward 0. `--rm-type math`
+  only reads `\boxed{}`; `--rm-type dapo` only `Answer:`. Use the format-agnostic
+  `rm.reward_func` (accepts either). Always pass `--label-key label` for the
+  `{prompt, label}` DAPO jsonl, or `Sample.label` is `None` and reward reads 0.
+- **Context length.** The 35B's DAPO chain-of-thought is ~14–17k tokens. An 8k
+  response cap truncates ~95% of rollouts mid-reasoning → reward ~0. Use ≥24k
+  (CP2 makes 24–32k feasible).
+- **`--opd-teacher-load` path.** Point at the checkpoint **parent** dir (contains
+  `latest_checkpointed_iteration.txt`), not an `iter_XXXXXXX` subdir. The subdir
+  has no metadata → silent fallback to base → teacher == student → `opd_reverse_kl ≈ 0`.
+  Sanity check: in the rollout log, `teacher_log_probs` should differ from
+  `rollout/log_probs`.
+- **Teacher must diverge.** A few RLVR steps at lr 1e-6 barely move the weights, so
+  the teacher ≈ base and OPD is inert (`opd_reverse_kl ≈ 5e-4`). lr 1e-5 diverges it
+  fast (`opd_reverse_kl ≈ 5e-2`). `--opd-kl-coef` cannot amplify a ~0 KL.
+- **Memory.** `with_ref = (--use-kl-loss or --kl-coef≠0)`. Dropping `--use-kl-loss`
+  keeps only student + teacher (2×35B) in memory; the teacher reverse-KL is the
+  regularizer. Adding it loads a 3rd model and risks OOM.
+
+## References
+- Phase-1 teacher checkpoint: https://huggingface.co/cm00cm/Qwen3.5-35B-A3B-DAPO-RLVR-teacher
+- ../README.md (served-teacher OPD), ../run-qwen3-8B-opd-megatron.sh (in-process teacher)
+- https://thinkingmachines.ai/blog/on-policy-distillation/

examples/on_policy_distillation/qwen3_5_35b_selfdistill/convert_gb200.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+# =============================================================================
+# Convert a HF Qwen3.5-35B-A3B checkpoint -> Megatron torch_dist.
+# Used to stage both the base (--ref-load) and the teacher (--opd-teacher-load)
+# for phase2_gb200.sh, since neither is pre-staged on /cluster_public.
+#
+# Usage: convert_gb200.sh <hf_checkpoint_dir> <torch_dist_save_dir>
+# =============================================================================
+set -ex
+HF_IN=${1:?hf checkpoint dir}
+SAVE_OUT=${2:?torch_dist save dir}
+MILES_DIR=${MILES_DIR:-/workspace/miles}
+MEGATRON_PATH=${MEGATRON_PATH:-/root/Megatron-LM}
+
+# Identical architecture spec to phase2_gb200.sh's MODEL_ARGS.
+MODEL_ARGS=(
+   --spec miles_plugins.models.qwen3_5 get_qwen3_5_spec
+   --disable-bias-linear --qk-layernorm --group-query-attention
+   --num-attention-heads 16 --num-query-groups 2 --kv-channels 256
+   --num-layers 40 --hidden-size 2048 --ffn-hidden-size 512
+   --normalization RMSNorm --apply-layernorm-1p --position-embedding-type rope
+   --norm-epsilon 1e-6 --rotary-percent 0.25 --swiglu
+   --untie-embeddings-and-output-weights --vocab-size 248320 --rotary-base 10000000
+   --moe-ffn-hidden-size 512 --moe-shared-expert-intermediate-size 512
+   --moe-router-score-function softmax --moe-token-dispatcher-type alltoall
+   --moe-router-topk 8
+   --moe-layer-freq "[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]"
+   --num-experts 256 --moe-grouped-gemm --moe-token-drop-policy probs --moe-router-dtype fp32
+   --moe-permute-fusion --moe-aux-loss-coeff 0 --attention-output-gate --moe-shared-expert-gate
+   --mtp-num-layers 1
+)
+
+cd "${MILES_DIR}"
+PYTHONPATH="${MILES_DIR}:${MEGATRON_PATH}" python3 "${MILES_DIR}/tools/convert_hf_to_torch_dist.py" \
+   "${MODEL_ARGS[@]}" \
+   --hf-checkpoint "${HF_IN}" \
+   --save "${SAVE_OUT}"
+echo "CONVERTED ${HF_IN} -> ${SAVE_OUT}"
+ls -la "${SAVE_OUT}"; cat "${SAVE_OUT}/latest_checkpointed_iteration.txt" 2>/dev/null || echo "(no latest_checkpointed_iteration.txt yet)"

examples/on_policy_distillation/qwen3_5_35b_selfdistill/eval_dapo_heldout.yaml

@@ -0,0 +1,17 @@
+# Held-out DAPO eval (disjoint from the training split, see README "Data split").
+# Scored by the example's format-agnostic reward (--custom-rm-path ...rm.reward_func),
+# which reports accuracy. max_response_len 24576 must exceed the model's reasoning
+# length (~14-17k for the base 35B) or accuracy is suppressed by truncation.
+eval:
+  defaults:
+    temperature: 0.6
+    top_p: 0.95
+  datasets:
+    - name: dapo_heldout
+      path: ${DATA_DIR}/dapo_eval.jsonl     # rendered by the launch scripts via envsubst
+      input_key: prompt
+      label_key: label                       # REQUIRED: without it Sample.label=None -> eval reads 0
+      n_samples_per_eval_prompt: 1
+      max_response_len: 24576
+      metadata_overrides:
+        opd_reward_mode: eval_math            # tags eval samples for reward_func_pure_opd

examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py

@@ -0,0 +1,61 @@
+"""Carve a disjoint train/eval split from dapo-math-17k for the self-distillation example.
+
+The file is ordered by difficulty, so we shuffle with a FIXED SEED before splitting
+(a contiguous tail-N split would be systematically easier and bias the eval). The
+512-problem eval split is held out from BOTH phases. Dedup is on prompt text (labels
+are not unique). Usage:
+
+    python make_split.py --src /path/dapo-math-17k.jsonl --out-dir /path/split
+"""
+
+import argparse
+import hashlib
+import json
+import os
+import random
+
+
+def prompt_text(d):
+    p = d["prompt"]
+    return "\n".join(m.get("content", "") for m in p) if isinstance(p, list) else str(p)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--src", required=True, help="dapo-math-17k.jsonl")
+    ap.add_argument("--out-dir", required=True)
+    ap.add_argument("--eval-n", type=int, default=512)
+    ap.add_argument("--seed", type=int, default=42)
+    args = ap.parse_args()
+    os.makedirs(args.out_dir, exist_ok=True)
+
+    rows, seen = [], set()
+    with open(args.src) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            key = prompt_text(json.loads(line))
+            if key in seen:
+                continue
+            seen.add(key)
+            rows.append(line)
+
+    random.Random(args.seed).shuffle(rows)  # fixed seed: reproducible, unbiased split
+    eval_rows, train_rows = rows[-args.eval_n :], rows[: -args.eval_n]
+
+    ek = {prompt_text(json.loads(r)) for r in eval_rows}
+    tk = {prompt_text(json.loads(r)) for r in train_rows}
+    assert ek.isdisjoint(tk), "LEAK: eval prompt found in train split"
+
+    with open(os.path.join(args.out_dir, "dapo_train.jsonl"), "w") as f:
+        f.write("\n".join(train_rows) + "\n")
+    with open(os.path.join(args.out_dir, "dapo_eval.jsonl"), "w") as f:
+        f.write("\n".join(eval_rows) + "\n")
+
+    md5 = hashlib.md5("\n".join(eval_rows).encode()).hexdigest()
+    print(f"train={len(train_rows)} eval={len(eval_rows)} seed={args.seed} eval_md5={md5}")
+
+
+if __name__ == "__main__":
+    main()