[maocheng23]

Stacked on #4459 (base branch opd-teacher-scoring-waste); PR 6 of 6. Review and merge bottom-up:

1. fix(megatron): pass FA4 unbounded window sides as None from TE (wrong causal attention on Blackwell) #4465: TE → FA4 causal window on Blackwell (wrong attention) (fix/te-fa4-causal-window → main)
2. [OPD] Keep a scalar reward: store teacher log-probs when scoring returns (supersedes #3266) #4466: OPD keeps a scalar reward (zero-std metrics crash with dict rewards) (opd-scalar-reward → fix/te-fa4-causal-window)
3. [Megatron] Split TE FusedAdam multi-tensor calls under the NVTETensor handle cap #4467: TE FusedAdam under the NVTETensor handle cap (per-expert tensors) (te-fused-adam-chunking → opd-scalar-reward)
4. [OPD] Allow --skip-actor-forward-only with sampled-token OPD #4468: --skip-actor-forward-only with sampled-token OPD (opd-skip-actor-forward-only → te-fused-adam-chunking)
5. [OPD] Cut wasted teacher-scoring work: pooled session, response-only teacher logprobs, skip stale fully-async groups #4459: Teacher-scoring waste: pooled session, logprob_start_len, stale-group skip (opd-teacher-scoring-waste → opd-skip-actor-forward-only)
6. [OPD] Add the Qwen3.6-35B-A3B multi-teacher OPD launcher (reasoning_gym, B200/GB300) #4493: Qwen3.6-35B-A3B multi-teacher OPD launcher (reasoning_gym) (mopd-reasoning-gym-recipe → opd-teacher-scoring-waste) ← this PR

New on 2026-10-27, on top of the rebased stack.
Tests at this PR's head (macOS, CPU): see Tests below. pre-commit run on this PR's files passes.
GPU-validated on 8×GB300 (2 trays × 4 GPUs, radixark/miles:latest with this stack checked out at bac0b1a344; c571c4b's Ray-node check was rerun against the same live cluster). Three runs of the Qwen3.6-35B-A3B multi-teacher OPD recipe with no recipe-side workarounds: sync with --skip-actor-forward-only (5 updates), fully async with --max-weight-staleness 1 (6 updates), and one-step-overlap async with MTP decoding (20 updates, held-out eval at 0/10/20). Not run on B200/B300, or with TP, CP or PP > 1.

Adds an official launcher for multi-teacher on-policy distillation of Qwen3.6-35B-A3B on two reasoning_gym tasks, ported from a standalone recipe validated on 8×GB300. The stack below it lets the launcher drop that recipe's workarounds.

What it trains

• Student Qwen/Qwen3.6-35B-A3B. Two frozen GRPO teachers: semianalysisai/Qwen3.6-35B-A3B-caesar-cipher-GRPO-20260924 @ dfedf677 and semianalysisai/Qwen3.6-35B-A3B-simple-geometry-GRPO-20260925 @ 679d4b66.
• Each rollout is scored only by its task's teacher (--opd-teacher-urls, --opd-teacher-key task).
• reasoning_gym 0.1.25, train seed 20000, held-out seed 10000.
• Sampled-token reverse-KL OPD, no task reward, PPO clip 0.2 / 0.28, Adam lr 1e-6.
• 20 updates × 128 prompts × 1 response, 30,720-token responses, temperature 1.
• Held-out eval: avg@4, exact match on the last <answer> after </think>.

Changes

• scripts/run_qwen3_6_35b_a3b_mopd.py: ScriptArgs on ExecuteTrainConfig, with prepare and execute entrypoints and --layout from a frozen table. It ships two first-class setups:
  • Synchronous (train.py --colocate), the matched baseline when run without --mtp-spec-decode:
    • b200-sync (the default): 8 B200 shared by the trainer (EP8) and 8 TP1 engines; each teacher DP4 on node B.
    • gb300-sync: the same on 2 × 4 GB300.
  • One-step overlap (train_async.py without --fully-async, nothing dropped), the recommended setup with --mtp-spec-decode:
    • b200-async: trainer EP8 on node A; 6 engines and 1 GPU per teacher on node B.
    • gb300-async: the same on 2 × 4 GB300.
  • prepare downloads the student and both teachers (pinned revisions), builds the prompts, and converts the student.
  • Async layouts require MILES_SCRIPT_EXTERNAL_RAY=1. Before submitting, execute checks that the trainer node sorts first. It lists the cluster's GPU nodes with a one-line ray.nodes() call through the backend (launchers reach Ray only through command_utils) and sorts them with Miles' own placement sort_key.
  • Defaults: token-mean loss (--calculate-per-token-loss), --log-probs-backend fused, no CPU offload. --mtp-spec-decode turns on EAGLE with Qwen3.6's MTP head. The eval config is passed inline (encode_pseudo_file).
• examples/on_policy_distillation/mopd_reasoning_gym/:
  • prepare_data.py: builds the train pool and held-out sets.
  • hooks.py: calls miles.rollout.on_policy_distillation.reward_func and adds timing stamps. The reward is the reasoning_gym verifier score, logged as rollout/raw_reward. post_process_rewards delegates to the stock hook, so the objective stays pure OPD.
  • serve_teachers.py: SGLang DP teachers, each on a fixed per-task port, with staggered starts.
  • analyze.py, make_report.py: optional tooling that turns a --dump-dir run into tables, plots and an HTML page. They need matplotlib.
  • README.md: one section per setup, each with the GPU split on B200 and GB300, the exact teacher, Ray and launch commands (trainer node at the lower IP; NCCL_IB_GID_INDEX=3 on RoCE clusters across subnets), and the expected results. It is mirrored to docs/examples/on-policy-distillation/mopd-reasoning-gym.md by scripts/tools/sync_example_docs.py, with the nav entries regenerated.
• docs/models/qwen/qwen3-6-moe.md: points to the launcher.
• Tests:
  • tests/fast/launch_scripts/test_run_qwen3_6_35b_a3b_mopd.py: sync vs async flags, the external-Ray requirement, teacher routing, MTP, and the node-order check (numeric IP order, rejection).
  • tests/fast/examples/on_policy_distillation/test_mopd_reasoning_gym_hooks.py: the verifier, routed teacher scoring under [OPD] Keep a scalar reward: store teacher log-probs when scoring returns (supersedes #3266) #4466's scalar-reward contract, the eval path, and post-processing.
  • Launch-script snapshots for prepare and execute.

What the stack removes from the validated recipe

I recorded the validated standalone recipe's launcher and this launcher under the snapshot harness and compared the train.py argv as a flag → values multiset, plus the Ray runtime env, for all four layouts with and without MTP. These are every difference:

Difference	Why
--custom-megatron-init-path mopd_rg.init_trainer removed	Its FA4 causal-window patch is replaced by #4465, which wraps TE's FA4 entry points in initialize.init(). Qwen3.6's full-attention layers go through TE DotProductAttention (get_gpt_decoder_block_spec(use_transformer_engine=True)), and TE looks the FA4 functions up as module globals at call time, so #4465 covers this model. Its FusedAdam chunking is replaced by #4467.
--custom-rm-path / --custom-reward-post-process-path / --custom-generate-function-path now point at examples.on_policy_distillation.mopd_reasoning_gym.hooks.*	Module moved. The reward hook no longer extracts teacher_log_probs from a returned dict, because #4466's reward_func stores them and returns a scalar. Timing stamps and the verifier reward are unchanged.
--skip-actor-forward-only added (sync layouts only)	Allowed by #4468 for sampled-token OPD with one optimizer step per rollout. Async layouts already use --use-rollout-logprobs and run no forward-only pass, so they are unchanged. --no-skip-actor-forward-only restores the validated sync configuration.
--eval-config inline instead of a JSON file written into the data dir	Same datasets and values; per-dataset fields moved to defaults.
--attention-dropout/--hidden-dropout 0 → 0.0; TB/W&B project names; recipe dir no longer on PYTHONPATH	Cosmetic / module move.

GPU validation

On 8×GB300 (2 trays × 4), radixark/miles:latest (TE 2.17.0+miles, flash-attn-4 4.0.0b35, SGLang 0.5.23) with this branch checked out at bac0b1a344, no recipe-side workarounds. c571c4b changes only the launcher's Ray-node listing; that check was rerun against the same live two-node cluster (passes for gb300-async, refuses a mismatched layout).

run	layout and flags	result
sync, 5 updates	gb300-sync --mtp-spec-decode (default --skip-actor-forward-only)	old-policy forward 0 s; train phase median 43 s (two-pass 55 s); ppo_kl = pg_clipfrac = 0; training accuracy 0.44 → 0.71
fully async, 6 updates	gb300-async --extra-args "--fully-async --max-weight-staleness 1 --async-max-concurrent-samples 256"	413 stale groups vs 768 trained; 294 skipped their teacher call; step 50–195 s
one-step async, 20 updates	gb300-async --mtp-spec-decode, eval at 0/10/20	held-out caesar_cipher 26.6% → 76.6% → 77.2%, simple_geometry 52.5% → 99.6% → 99.4%; median step 171 s excluding eval

The one-step async result matches the pre-stack run of the same recipe with its workarounds (78.5% / 99.6%, 173 s) within eval noise (128 prompts × 4 samples, about ±2–3 points). train/train_rollout_logprob_abs_diff was 0.014–0.020 in every run.

Not run at this head: the 20-update sync layout, the B200 layouts, and any layout with TP, CP or PP > 1. In async layouts Ray may place the rollout manager on node B, and --dump-dir must then be on storage both nodes share for analyze.py to find rollout_data/.

The README and docs page report only our own GB300 runs and the teacher cards' public accuracies; no third-party comparison numbers are in the repo.

Tests

All runs on macOS, CPU, Python 3.12. Nothing here can run on GPU.

• New tests:
  • test_run_qwen3_6_35b_a3b_mopd.py: 11 passed.
  • test_mopd_reasoning_gym_hooks.py: 11 passed. Putting back the validated recipe's dict-extraction workaround makes 2 of them fail.
• Real Ray check: I ran the node-listing command against a real local Ray head with 8 fake GPUs. It parses, and the check rejects that single-node cluster for both async layouts.
• pytest tests/fast/launch_scripts tests/manual/launch_scripts: 525 passed, 42 failed, 86 errors. The 128 failures and errors are identical on origin/main, which has 510 passing.
  • 108 come from the shell harness needing /proc.
  • The rest are goldens that don't match on main (run_deepseek_v4_1, run_qwen3_0_6b_torchtitan, amd/run_gpt_oss_20b_lora, run_glm5_3_flash_lora, model-args snapshots).
  • macOS /root is read-only, so a local-only pytest plugin stubbed convert_checkpoint's lock and stamp writes there. CI runs as root on Linux and needs no stub. Without the stub, 46 more errors come from that lock, on main as well.
• Snapshots: prepare.txt and execute.txt were recorded with MILES_UPDATE_LAUNCH_SCRIPT_SNAPSHOTS=1 and verified stable.
• tests/fast/backends at this head: 3,469 passed, 16 failed. The same 16 fail on origin/main (3,436 passed there): loss-snapshot tolerances on macOS, torchtitan on-disk optimizer state, actor-init checkpoint scope, checkpoint transfer, and one order-dependent weight-update test that passes alone.
• tests/fast/{rollout,ray,utils,examples,doc}, test_train_async.py, test_megatron_cli_flags.py: 12,810 passed, 35 failed, 4 errors.
  • 36 of these fail identically on origin/main: agentic and session benchmarks, mock SGLang servers, disk-backed tensors, helm snapshots, git state.
  • test_hot_restart::...raises_what_it_last_saw is timing-flaky; it fails on origin/main in the same environment too.
  • test_transport::...proxy_stop_closes_idle_connections passed when rerun.
  • test_ray_communication_boundary caught the launcher importing ray; fixed in c571c4b and now passing.
• pre-commit run: clean on every changed file.

🤖 Generated with Claude Code