diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/README.md b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/README.md new file mode 100644 index 00000000000..569ea741906 --- /dev/null +++ b/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 `` 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/ diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/convert_gb200.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/convert_gb200.sh new file mode 100644 index 00000000000..8cd85d2e047 --- /dev/null +++ b/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 +# ============================================================================= +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)" diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/eval_dapo_heldout.yaml b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/eval_dapo_heldout.yaml new file mode 100644 index 00000000000..217b9aa4256 --- /dev/null +++ b/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 diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/make_split.py new file mode 100644 index 00000000000..15f9264ff38 --- /dev/null +++ b/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() diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase1_rlvr_teacher.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase1_rlvr_teacher.sh new file mode 100755 index 00000000000..b04c3d5a5af --- /dev/null +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase1_rlvr_teacher.sh @@ -0,0 +1,105 @@ +#!/bin/bash +# ============================================================================= +# Phase 1: RLVR-train a Qwen3.5-35B-A3B *teacher* on DAPO-math (single 8xH200 node) +# ============================================================================= +# Trains the base model with GRPO + format-agnostic correctness reward so it becomes +# measurably BETTER and MORE CONCISE than the base. This trained checkpoint is the +# teacher distilled into the base student in Phase 2. +# +# Single node, 8 GPUs: TP2 / PP1 / CP2 / EP8 / ETP1 (world=8, DP2) +# lr 1e-5 is intentional: it makes the teacher DIVERGE from base fast enough to +# matter within ~5-10 steps. lr 1e-6 (few steps) barely moves the weights, leaving +# the teacher ~= base, which makes the Phase-2 OPD reverse-KL ~0 (inert). +# +# Usage: bash phase1_rlvr_teacher.sh +# Env (override as needed): +# MODEL_DIR dir holding Qwen3.5-35B-A3B and Qwen3.5-35B-A3B_torch_dist +# DATA_DIR dir holding dapo_train.jsonl / dapo_eval.jsonl (see make_split.py) +# OUTPUT_DIR writable, *node-local-persistent* dir for checkpoints (see README) +# EXAMPLE_DIR this directory (for rm.py + eval config on PYTHONPATH) +# ============================================================================= +set -ex +export PYTHONUNBUFFERED=16 + +MODEL_DIR=${MODEL_DIR:-/cluster_public/miles_data/models} +DATA_DIR=${DATA_DIR:-/node_public/maocheng-qwen35/data} +OUTPUT_DIR=${OUTPUT_DIR:-/node_public/maocheng-qwen35/ckpt-teacher} +EXAMPLE_DIR=${EXAMPLE_DIR:-$(cd "$(dirname "$0")" && pwd)} +MILES_DIR=${MILES_DIR:-/root/miles} +RAY_ADDRESS=${RAY_ADDRESS:-http://127.0.0.1:8265} +mkdir -p "${OUTPUT_DIR}" + +# Render the eval config (substitutes ${DATA_DIR}). +EVAL_CONFIG="${OUTPUT_DIR}/eval_dapo_heldout.yaml" +DATA_DIR="${DATA_DIR}" envsubst < "${EXAMPLE_DIR}/eval_dapo_heldout.yaml" > "${EVAL_CONFIG}" + +# Qwen3.5-35B-A3B architecture (no scripts/models/*.sh ships for it). +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 +) +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B + --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist + --load ${OUTPUT_DIR} --save ${OUTPUT_DIR} --save-interval 5 +) +ROLLOUT_ARGS=( + --prompt-data ${DATA_DIR}/dapo_train.jsonl --input-key prompt --label-key label + --apply-chat-template --rollout-shuffle + --num-rollout 20 --rollout-batch-size 32 --n-samples-per-prompt 8 + --rollout-max-response-len 24576 --rollout-temperature 1 --num-steps-per-rollout 1 + --over-sampling-batch-size 32 --global-batch-size 256 --balance-data +) +# Format-agnostic correctness reward -> rollout/raw_reward == accuracy. +RM_ARGS=( --custom-rm-path examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func ) +EVAL_ARGS=( --eval-interval 5 --eval-config ${EVAL_CONFIG} ) +GRPO_ARGS=( + --advantage-estimator grpo --kl-loss-type low_var_kl --entropy-coef 0.00 + --eps-clip 0.2 --eps-clip-high 0.28 --use-tis +) +OPTIMIZER_ARGS=( + --optimizer adam --lr 1e-5 --lr-decay-style constant --weight-decay 0.1 + --adam-beta1 0.9 --adam-beta2 0.98 + --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer +) +PERF_ARGS=( + --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1 + --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 + --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 + --use-dynamic-batch-size --max-tokens-per-gpu 16384 --log-probs-chunk-size 4096 +) +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 8 --sglang-mem-fraction-static 0.8 --sglang-ep-size 8 + --sglang-watchdog-timeout 1800 --sglang-enable-metrics --sglang-attention-backend fa3 + --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay + --sglang-mamba-scheduler-strategy extra_buffer +) +MISC_ARGS=( + --attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 --attention-backend flash +) +WANDB_ARGS=( --use-wandb --wandb-project miles-opd --wandb-group qwen3.5-35b-rlvr-teacher ) + +RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\"}}" + +cd "${MILES_DIR}" +ray job submit --address="${RAY_ADDRESS}" --submission-id qwen3.5-rlvr-teacher --no-wait \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 ${MILES_DIR}/train.py \ + --actor-num-nodes 1 --actor-num-gpus-per-node 8 --num-gpus-per-node 8 --colocate \ + ${MODEL_ARGS[@]} ${CKPT_ARGS[@]} ${ROLLOUT_ARGS[@]} ${OPTIMIZER_ARGS[@]} ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} ${PERF_ARGS[@]} ${EVAL_ARGS[@]} ${SGLANG_ARGS[@]} ${MISC_ARGS[@]} ${RM_ARGS[@]} +echo "Submitted Phase-1 RLVR teacher training (submission-id qwen3.5-rlvr-teacher)." +echo "Watch rollout/raw_reward climb and eval/dapo_heldout rise above the base ~0.83." diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh new file mode 100755 index 00000000000..512abe575c3 --- /dev/null +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_gb200.sh @@ -0,0 +1,140 @@ +#!/bin/bash +# ============================================================================= +# Phase 2: On-policy distillation of the Phase-1 teacher into the BASE student +# ============================================================================= +# Student = base Qwen3.5-35B-A3B. Teacher = the Phase-1 RLVR checkpoint (in-process, +# --opd-type megatron). The student trains only on its own rollouts; the teacher's +# token-level reverse-KL is folded into the GRPO advantages. +# +# Two modes (MODE env var): +# pure (default) — training task reward = 0; ONLY the teacher reverse-KL drives +# learning. Cleanest attribution: any change is distillation. +# eval/dapo_heldout still measures real accuracy. +# grounded — keep the correctness reward (rollout/raw_reward == accuracy, climbs) +# AND add the teacher reverse-KL on top. +# +# Memory note: we do NOT pass --use-kl-loss, so the reference model is NOT loaded +# (with_ref = use_kl_loss or kl_coef!=0). That keeps only student + teacher in memory +# (2 x 35B, ~124/143 GB per GPU). Adding --use-kl-loss would load a 3rd model and risk OOM. +# The teacher's reverse-KL is the regularizer. +# +# Usage: TEACHER_LOAD=/path/to/teacher_ckpt_parent bash phase2_opd_selfdistill.sh +# IMPORTANT: --opd-teacher-load must point at the checkpoint PARENT directory (the one +# containing latest_checkpointed_iteration.txt), NOT an iter_XXXXXXX subdir. Pointing at +# the subdir prints "could not find metadata file" and silently falls back to base -> +# teacher == student -> opd_reverse_kl ~= 0 (inert). +# ============================================================================= +set -ex +export PYTHONUNBUFFERED=16 + +MODE=${MODE:-pure} +# GB200 tiling: this cluster is 4 GPUs/node, so world=8 = 2 nodes x 4 GPUs. +# (Upstream phase2_opd_selfdistill.sh assumes a single 8xH200 node.) The parallel +# dims (TP2/PP1/CP2/EP8/ETP1, DP4) are unchanged; only the node tiling differs. +# Both nodes must share one NVLink (MNNVL) domain -- launch with MSC_RACK pinned. +ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-2} +GPUS_PER_NODE=${GPUS_PER_NODE:-4} +MODEL_DIR=${MODEL_DIR:-/cluster_public/miles_data/models} +DATA_DIR=${DATA_DIR:-/node_public/maocheng-qwen35/data} +OUTPUT_DIR=${OUTPUT_DIR:-/node_public/maocheng-qwen35/ckpt-opd-${MODE}} +TEACHER_LOAD=${TEACHER_LOAD:-/node_public/maocheng-qwen35/ckpt-teacher} # parent dir! +EXAMPLE_DIR=${EXAMPLE_DIR:-$(cd "$(dirname "$0")" && pwd)} +MILES_DIR=${MILES_DIR:-/workspace/miles} +RAY_ADDRESS=${RAY_ADDRESS:-http://127.0.0.1:8265} +OPD_KL_COEF=${OPD_KL_COEF:-0.2} +mkdir -p "${OUTPUT_DIR}" + +EVAL_CONFIG="${OUTPUT_DIR}/eval_dapo_heldout.yaml" +DATA_DIR="${DATA_DIR}" envsubst < "${EXAMPLE_DIR}/eval_dapo_heldout.yaml" > "${EVAL_CONFIG}" + +if [ "${MODE}" = "pure" ]; then + RM_FUNC=examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func_pure_opd +else + RM_FUNC=examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func +fi + +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 flex + --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 +) +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B + --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist + --load ${OUTPUT_DIR} --save ${OUTPUT_DIR} --save-interval 5 +) +OPD_ARGS=( + --use-opd --opd-type megatron --opd-teacher-load ${TEACHER_LOAD} --opd-kl-coef ${OPD_KL_COEF} +) +ROLLOUT_ARGS=( + --prompt-data ${DATA_DIR}/dapo_train.jsonl --input-key prompt --label-key label + --apply-chat-template --rollout-shuffle + --num-rollout 12 --rollout-batch-size 32 --n-samples-per-prompt 8 + --rollout-max-response-len 24576 --rollout-temperature 1 --num-steps-per-rollout 1 + --over-sampling-batch-size 32 --global-batch-size 256 --balance-data +) +RM_ARGS=( --custom-rm-path ${RM_FUNC} ) +EVAL_ARGS=( --eval-interval 5 --eval-config ${EVAL_CONFIG} ) +GRPO_ARGS=( + --advantage-estimator grpo --kl-loss-type low_var_kl --entropy-coef 0.00 + --eps-clip 0.2 --eps-clip-high 0.28 --use-tis +) +OPTIMIZER_ARGS=( + --optimizer adam --lr 1e-5 --lr-decay-style constant --weight-decay 0.1 + --adam-beta1 0.9 --adam-beta2 0.98 + --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer +) +PERF_ARGS=( + --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1 + --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 + --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 + --use-dynamic-batch-size --max-tokens-per-gpu 16384 --log-probs-chunk-size 4096 +) +# Blackwell (GB200 sm100 / B300 sm103) backends, per scripts/run_qwen3_5_35b_a3b_mtp_cp2_ep8.py: +# - moe-runner-backend flashinfer_cutlass: the default triton fused-MoE mis-shards +# routed experts on the megatron->sglang weight sync (w13 reshape 64 vs 2048). +# - attention-backend trtllm_mha: FA3 is SM<=90 only; flashinfer/fa3 don't fit here. +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 8 --sglang-mem-fraction-static 0.7 --sglang-ep-size 8 + --sglang-watchdog-timeout 1800 --sglang-enable-metrics + --sglang-moe-runner-backend flashinfer_cutlass --sglang-attention-backend trtllm_mha + --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay + --sglang-mamba-scheduler-strategy extra_buffer +) +MISC_ARGS=( + --attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 --attention-backend flash +) +WANDB_ARGS=( --use-wandb --wandb-project miles-opd --wandb-group qwen3.5-35b-opd-${MODE} ) + +# PROMETHEUS_PORT override: the contractor k8s namespace has a `prometheus` +# Service, so kube auto-injects PROMETHEUS_PORT=tcp://:9090 into every pod. +# miles' add_prometheus_arguments does int(os.environ["PROMETHEUS_PORT"]) and +# crashes on the URL; force the plain port back via the ray runtime env. +# NCCL_NVLS_ENABLE=0: on multi-node GB200 the NVLS (NVLink SHARP) bind fails +# during cross-node ncclCommInitRank ("unhandled cuda error" / CUDA 999), which +# kills the EP8 sglang engine spanning both nodes. MNNVL stays enabled (pod env) +# for the NVLink fabric itself; only NVLS multicast is disabled. +RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\", \"PROMETHEUS_PORT\": \"9090\", \"NCCL_NVLS_ENABLE\": \"0\"}}" + +cd "${MILES_DIR}" +ray job submit --address="${RAY_ADDRESS}" --submission-id qwen3.5-opd-${MODE} --no-wait \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 ${MILES_DIR}/train.py \ + --actor-num-nodes ${ACTOR_NUM_NODES} --actor-num-gpus-per-node ${GPUS_PER_NODE} --num-gpus-per-node ${GPUS_PER_NODE} --colocate \ + ${MODEL_ARGS[@]} ${CKPT_ARGS[@]} ${OPD_ARGS[@]} ${ROLLOUT_ARGS[@]} ${OPTIMIZER_ARGS[@]} ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} ${PERF_ARGS[@]} ${EVAL_ARGS[@]} ${SGLANG_ARGS[@]} ${MISC_ARGS[@]} ${RM_ARGS[@]} +echo "Submitted Phase-2 OPD (${MODE}) self-distillation (submission-id qwen3.5-opd-${MODE})." +echo "Watch opd_reverse_kl (>>0 means the teacher differs from the student) and" +echo "eval/dapo_heldout (student moving toward the teacher's accuracy, with shorter responses)." diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_opd_selfdistill.sh b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_opd_selfdistill.sh new file mode 100755 index 00000000000..0c3109a1744 --- /dev/null +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/phase2_opd_selfdistill.sh @@ -0,0 +1,121 @@ +#!/bin/bash +# ============================================================================= +# Phase 2: On-policy distillation of the Phase-1 teacher into the BASE student +# ============================================================================= +# Student = base Qwen3.5-35B-A3B. Teacher = the Phase-1 RLVR checkpoint (in-process, +# --opd-type megatron). The student trains only on its own rollouts; the teacher's +# token-level reverse-KL is folded into the GRPO advantages. +# +# Two modes (MODE env var): +# pure (default) — training task reward = 0; ONLY the teacher reverse-KL drives +# learning. Cleanest attribution: any change is distillation. +# eval/dapo_heldout still measures real accuracy. +# grounded — keep the correctness reward (rollout/raw_reward == accuracy, climbs) +# AND add the teacher reverse-KL on top. +# +# Memory note: we do NOT pass --use-kl-loss, so the reference model is NOT loaded +# (with_ref = use_kl_loss or kl_coef!=0). That keeps only student + teacher in memory +# (2 x 35B, ~124/143 GB per GPU). Adding --use-kl-loss would load a 3rd model and risk OOM. +# The teacher's reverse-KL is the regularizer. +# +# Usage: TEACHER_LOAD=/path/to/teacher_ckpt_parent bash phase2_opd_selfdistill.sh +# IMPORTANT: --opd-teacher-load must point at the checkpoint PARENT directory (the one +# containing latest_checkpointed_iteration.txt), NOT an iter_XXXXXXX subdir. Pointing at +# the subdir prints "could not find metadata file" and silently falls back to base -> +# teacher == student -> opd_reverse_kl ~= 0 (inert). +# ============================================================================= +set -ex +export PYTHONUNBUFFERED=16 + +MODE=${MODE:-pure} +MODEL_DIR=${MODEL_DIR:-/cluster_public/miles_data/models} +DATA_DIR=${DATA_DIR:-/node_public/maocheng-qwen35/data} +OUTPUT_DIR=${OUTPUT_DIR:-/node_public/maocheng-qwen35/ckpt-opd-${MODE}} +TEACHER_LOAD=${TEACHER_LOAD:-/node_public/maocheng-qwen35/ckpt-teacher} # parent dir! +EXAMPLE_DIR=${EXAMPLE_DIR:-$(cd "$(dirname "$0")" && pwd)} +MILES_DIR=${MILES_DIR:-/root/miles} +RAY_ADDRESS=${RAY_ADDRESS:-http://127.0.0.1:8265} +OPD_KL_COEF=${OPD_KL_COEF:-0.2} +mkdir -p "${OUTPUT_DIR}" + +EVAL_CONFIG="${OUTPUT_DIR}/eval_dapo_heldout.yaml" +DATA_DIR="${DATA_DIR}" envsubst < "${EXAMPLE_DIR}/eval_dapo_heldout.yaml" > "${EVAL_CONFIG}" + +if [ "${MODE}" = "pure" ]; then + RM_FUNC=examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func_pure_opd +else + RM_FUNC=examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func +fi + +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 +) +CKPT_ARGS=( + --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B + --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B_torch_dist + --load ${OUTPUT_DIR} --save ${OUTPUT_DIR} --save-interval 5 +) +OPD_ARGS=( + --use-opd --opd-type megatron --opd-teacher-load ${TEACHER_LOAD} --opd-kl-coef ${OPD_KL_COEF} +) +ROLLOUT_ARGS=( + --prompt-data ${DATA_DIR}/dapo_train.jsonl --input-key prompt --label-key label + --apply-chat-template --rollout-shuffle + --num-rollout 12 --rollout-batch-size 32 --n-samples-per-prompt 8 + --rollout-max-response-len 24576 --rollout-temperature 1 --num-steps-per-rollout 1 + --over-sampling-batch-size 32 --global-batch-size 256 --balance-data +) +RM_ARGS=( --custom-rm-path ${RM_FUNC} ) +EVAL_ARGS=( --eval-interval 5 --eval-config ${EVAL_CONFIG} ) +GRPO_ARGS=( + --advantage-estimator grpo --kl-loss-type low_var_kl --entropy-coef 0.00 + --eps-clip 0.2 --eps-clip-high 0.28 --use-tis +) +OPTIMIZER_ARGS=( + --optimizer adam --lr 1e-5 --lr-decay-style constant --weight-decay 0.1 + --adam-beta1 0.9 --adam-beta2 0.98 + --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer +) +PERF_ARGS=( + --tensor-model-parallel-size 2 --sequence-parallel --pipeline-model-parallel-size 1 + --context-parallel-size 2 --expert-model-parallel-size 8 --expert-tensor-parallel-size 1 + --recompute-granularity full --recompute-method uniform --recompute-num-layers 1 + --use-dynamic-batch-size --max-tokens-per-gpu 16384 --log-probs-chunk-size 4096 +) +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 8 --sglang-mem-fraction-static 0.8 --sglang-ep-size 8 + --sglang-watchdog-timeout 1800 --sglang-enable-metrics --sglang-attention-backend fa3 + --sglang-cuda-graph-bs 1 2 4 8 16 32 --use-rollout-routing-replay + --sglang-mamba-scheduler-strategy extra_buffer +) +MISC_ARGS=( + --attention-dropout 0.0 --hidden-dropout 0.0 --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 --attention-backend flash +) +WANDB_ARGS=( --use-wandb --wandb-project miles-opd --wandb-group qwen3.5-35b-opd-${MODE} ) + +RUNTIME_ENV_JSON="{\"env_vars\": {\"PYTHONPATH\": \"${MILES_DIR}:/root/Megatron-LM/\", \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", \"WANDB_API_KEY\": \"${WANDB_API_KEY}\"}}" + +cd "${MILES_DIR}" +ray job submit --address="${RAY_ADDRESS}" --submission-id qwen3.5-opd-${MODE} --no-wait \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 ${MILES_DIR}/train.py \ + --actor-num-nodes 1 --actor-num-gpus-per-node 8 --num-gpus-per-node 8 --colocate \ + ${MODEL_ARGS[@]} ${CKPT_ARGS[@]} ${OPD_ARGS[@]} ${ROLLOUT_ARGS[@]} ${OPTIMIZER_ARGS[@]} ${GRPO_ARGS[@]} \ + ${WANDB_ARGS[@]} ${PERF_ARGS[@]} ${EVAL_ARGS[@]} ${SGLANG_ARGS[@]} ${MISC_ARGS[@]} ${RM_ARGS[@]} +echo "Submitted Phase-2 OPD (${MODE}) self-distillation (submission-id qwen3.5-opd-${MODE})." +echo "Watch opd_reverse_kl (>>0 means the teacher differs from the student) and" +echo "eval/dapo_heldout (student moving toward the teacher's accuracy, with shorter responses)." diff --git a/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py new file mode 100644 index 00000000000..88b8cbd572e --- /dev/null +++ b/examples/on_policy_distillation/qwen3_5_35b_selfdistill/rm.py @@ -0,0 +1,73 @@ +"""Reward functions for the Qwen3.5-35B-A3B self-distillation example. + +Two reward functions are provided: + +* ``reward_func`` — format-agnostic boxed-answer correctness (0/1). Used as the + task reward for Phase 1 (RLVR teacher training) AND as the eval scorer for both + phases. It accepts EITHER answer format the DAPO prompt may elicit: + ``\\boxed{ANS}`` (via ``grade_answer_verl``) OR ``Answer: ANS`` (via the DAPO + ``compute_score``). A clean 0/1 == true accuracy, so ``rollout/raw_reward`` and + ``eval/`` are directly interpretable. + +* ``reward_func_pure_opd`` — for *pure* on-policy distillation (Phase 2, pure + variant). EVAL samples (tagged ``opd_reward_mode=eval_math`` via the eval + config's ``metadata_overrides``) are scored for accuracy so ``eval/`` + still measures real held-out accuracy; TRAINING samples return a constant + ``0.0`` so the GRPO task advantage is ~0 and the ONLY learning signal is the + teacher's reverse-KL. This isolates OPD's effect for clean attribution. + +Why a custom reward instead of ``--rm-type``: + - ``--rm-type math`` only extracts ``\\boxed{}``; ``--rm-type dapo`` only the + ``Answer:`` form; ``--rm-type deepscaler`` additionally requires a + ```` tag in the response. Qwen3.5 reasons inline (no ````), + so ``deepscaler`` would score every response 0. This format-agnostic reward + avoids all three traps. + +Place ``examples/on_policy_distillation/qwen3_5_35b_selfdistill`` on PYTHONPATH (or +run from the miles repo root) and pass e.g. +``--custom-rm-path examples.on_policy_distillation.qwen3_5_35b_selfdistill.rm.reward_func``. +""" + +from miles.rollout.rm_hub.math_dapo_utils import compute_score as _dapo_score +from miles.rollout.rm_hub.math_utils import grade_answer_verl + + +def _is_eval_sample(sample) -> bool: + md = sample.metadata if isinstance(getattr(sample, "metadata", None), dict) else {} + return md.get("opd_reward_mode") == "eval_math" + + +def _is_correct(sample) -> bool: + response, label = sample.response or "", sample.label + if label is None: + return False + # boxed form: \boxed{ANS} + try: + if grade_answer_verl(response, label): + return True + except Exception: + pass + # DAPO prompt form: "Answer: ANS" + try: + result = _dapo_score(response, label) + if isinstance(result, dict) and result.get("acc"): + return True + except Exception: + pass + return False + + +async def reward_func(args, sample, **kwargs): + """Format-agnostic correctness (1.0 / 0.0). Phase-1 task reward + eval scorer.""" + return 1.0 if _is_correct(sample) else 0.0 + + +async def reward_func_pure_opd(args, sample, **kwargs): + """Pure-OPD reward: accuracy for eval samples, constant 0.0 for training. + + Training reward is constant so the GRPO task advantage vanishes and the only + learning signal is the OPD reverse-KL toward the teacher. + """ + if _is_eval_sample(sample): + return 1.0 if _is_correct(sample) else 0.0 + return 0.0 diff --git a/miles_plugins/mbridge/qwen3_5.py b/miles_plugins/mbridge/qwen3_5.py index fe1556ef7d9..67d1dcb35f7 100644 --- a/miles_plugins/mbridge/qwen3_5.py +++ b/miles_plugins/mbridge/qwen3_5.py @@ -111,6 +111,21 @@ class Qwen3_5Bridge(Qwen2MoEBridge): "mlp.experts.linear_fc2": ["mtp.layers.{layer_number}.mlp.experts.down_proj"], } + # Main-layer experts can ALSO ship unfused (per-expert .weight files) — e.g. a + # checkpoint round-tripped out through convert_torch_dist_to_hf, whose experts + # are stored split rather than as the fused 3-D ``gate_up_proj``/``down_proj`` + # tensors the public Qwen3.5-35B-A3B release uses. ``_experts_fused()`` + # autodetects from the safetensor index, mirroring ``_mtp_experts_fused()``. + _MLP_EXPERTS_MAPPING_UNFUSED = { + "mlp.experts.linear_fc1": [ + "model.language_model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight", + "model.language_model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight", + ], + "mlp.experts.linear_fc2": [ + "model.language_model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight" + ], + } + # Override to make ffn_hidden_size optional (Qwen3.5 MoE has no intermediate_size) _CONFIG_MAPPING = { "num_layers": "num_hidden_layers", @@ -180,11 +195,34 @@ def _get_gptmodel_args(self) -> dict: ret["mtp_block_spec"] = mtp_block_spec return ret + def _experts_fused(self) -> bool: + """Detect whether MAIN-layer MoE expert weights are fused 3-D tensors. + + Fused (public Qwen3.5-35B-A3B): ``...mlp.experts.gate_up_proj``. + Unfused (round-tripped checkpoints): ``...mlp.experts.{i}.gate_proj.weight``. + Same lazy/cached resolution from ``safetensor_io.index`` as + ``_mtp_experts_fused()``; defaults to fused when the index is unavailable + so pre-init access keeps the historical behaviour. + """ + cached = getattr(self, "_experts_fused_cached", None) + if cached is not None: + return cached + io = getattr(self, "safetensor_io", None) + index = getattr(io, "index", None) if io is not None else None + if not index: + return True + fused = any("model.language_model.layers." in k and k.endswith("mlp.experts.gate_up_proj") for k in index) + self._experts_fused_cached = fused + return fused + def _weight_name_mapping_mlp(self, name: str) -> list[str]: - """Override to handle fused expert weights.""" + """Override to handle fused (default) or unfused per-expert weights.""" layer_number = name.split(".")[2] + mapping = self._MLP_MAPPING + if "mlp.experts.linear_fc" in name and not self._experts_fused(): + mapping = {**self._MLP_MAPPING, **self._MLP_EXPERTS_MAPPING_UNFUSED} convert_names = [] - for keyword, mapping_names in self._MLP_MAPPING.items(): + for keyword, mapping_names in mapping.items(): if keyword in name: if "{expert_id}" in mapping_names[0]: expert_id = name.split("weight")[-1]