MoE post-training with Rollout Routing Replay (R3), validated on Qwen3-30B-A3B (MATH-500 86.95 -> 92.35)

[iamziyuzhao]

What this adds

AstraFlow could not post-train Mixture-of-Experts models: the trainer's routers re-decide expert assignment, so gradients are computed against different experts than the rollout used, and MoE RL destabilises. This PR adds Rollout Routing Replay (R3, arXiv:2510.11370) end to end, plus the loop-control and weight-sync changes a validated Qwen3-30B-A3B run turned out to require. It is deliberately limited to what runs when you train a MoE model with the recipe below, and written as the minimum change to the existing code: one contract per hop instead of guard code, and a loud failure at the first place a contract violation reaches.

• R3 capture (raas/) — gconfig.return_routed_experts asks sglang (>= 0.5.13) for per-token top-k expert ids; the backend decodes the base64 blob into a flat int32 array, the engine accumulates the chunks across interrupt/resume iterations and reshapes them to [seq_len - 1, num_hidden_layers * top_k] (numpy raises if the payload does not cover every forwarded position). SGLangConfig.build_args refuses the two launches that capture nothing: an implicit or bypassing moe_runner_backend (sglang's auto picks a non-capturing kernel on Blackwell) and a disabled chunked prefill.
• R3 transport (core/workflow/, dataflow/utils.py, train_worker/utils/data.py) — the mask rides the trajectory as int32 [1, S, num_hidden_layers * top_k], aligned to input_ids (the never-forwarded final position repeats the last row); padding, micro-batch splitting and packing generalise over trailing dims.
• R3 replay (train_worker/utils/mcore/routing_replay.py, Megatron engine) — actor.megatron.moe_router_replay patches TopKRouter.forward to softmax(logits.masked_fill(~mask, -inf)) in fp32, which equals top-k renormalisation and keeps router gradients. Records are installed per micro-batch and per model chunk right before that chunk's forward and read through a forward and a backward cursor (activation recompute, 1F1B); every pass asserts all records were consumed. Router configurations the masked softmax does not reproduce are rejected at init; context/sequence parallelism raise NotImplementedError. The reference policy inherits the actor's setting.
• transformers 5.x — mbridge 0.1.0 reads hf_config.rope_theta, which transformers 5 moved into rope_parameters; the engine backfills it right after AutoBridge.from_pretrained (two lines).
• Closed-loop rollout buffer — dataflow.buffer.max_buffered_samples pauses prompt submission while enough samples are already buffered (one _submit_budget method). Without it the fresh heap fills, put() evicts the oldest, and edf hands the trainer the oldest survivor: two open-loop 30B runs trained on data ~22 versions old and eroded/collapsed once staleness passed ~17. buffer/evicted is exported; the eval block resumes the flow on every exit path.
• Weight-sync budget — one WEIGHT_SYNC_TIMEOUT_SEC = 300 replaces waits tuned for 8B models (a 57 GiB sync is 76 s on the loopback, 110-130 s next to a checkpoint write). RaaS labels the engine with the version the sender actually served, so queued notify_version requests hit the existing "already loaded" check and coalesce instead of pulling every version in turn. A sender-agent startup failure raises with its own traceback instead of a TypeError.
• Recipe — examples/math/qwen3-30b-a3b-m2po/: one 8xB200 node, trainer TP1/PP1/DP4/EP4 on 4 GPUs + SGLang DP2 on 2, GRPO (clip 0.2/0.28, lr 1e-6, 256 samples/step, no KL).

Validation

• Qwen3-30B-A3B, 800 steps (previous revision of this branch; MATH prompts, 4x B200 trainer + 2 rollout GPUs, ~14 h wall clock including one GPU-contention crash and an HF-checkpoint resume at step 299):

  global step	0	100	200	300	400	500	600	700	750	800
  MATH-500 avg@4	86.95	87.75	88.30	89.40	89.10	90.90	90.80	91.85	92.10	92.35

  +5.40 over baseline (pass@4 97.0), still rising at the final eval. Sample staleness held at 2-5 versions throughout; zero buffer evictions; zero staleness drops.

• This revision, Qwen3-30B-A3B, 40 steps on the same node and recipe (2026-09-09, 4x B200 trainer + 1 rollout GPU): MATH-500 avg@4 87.45 at step 0 and 86.55 at step 25, against 86.95 and 86.70 for the previous revision's run at the same steps (eval noise is about +-0.7). Importance-weight abs delta 0.0074-0.0102 (previous: ~0.009), sample staleness 2.0-4.2 avg / <= 5 max (previous: 2.8-4.9 / <= 5), zero evictions and zero stale drops, entropy 0.23-0.34, grad norm 0.0035-0.0088, clip ratio 0.04-0.06 %, replay armed on all 48 layers in every rank, weight-update coalescing fired 7 times, no tracebacks in any log, clean shutdown.

• This revision, 800 steps, same node and recipe: in progress (started 2026-09-09; saves every 100 steps, evals every 25). The full curve will be added here next to the previous revision's when it finishes.

• This revision vs the previous one, same GPU smoke (tiny 4-layer Qwen3-MoE, sglang rollout with capture -> dataflow closed loop (cap 16) -> Megatron actor + ref with replay -> TCP weight sync, 8 steps, same seed): both runs report replay armed on 4/4 layers, importance weights within [0.992, 1.008] and approx-KL max ~0.0078 on every step, entropy 11.88, zero evictions, staleness 1-2, one queued weight update coalesced. The same ids reach the router: the flat int32 layout is a reshape of the previous int16 [rows, L, K], checked bit-for-bit on CPU.

• Radix cache: a GPU probe (sglang 0.5.13.post1) showed routing rows and tokens bit-identical between radix cache on and off, on full and partial prefix hits, and after a weight reload (which flushed the cache). The earlier refusal of the radix cache was wrong and is gone; the recipe still leaves it disabled.

• Replay equivalence (previous revision, same math): record-vs-replay max logit delta 0.0078 at Qwen3-30B-A3B shapes, bitwise 0.0 at 4L/top-8 and 48L/top-2.

• Unit tests: CPU tests live on iamziyuzhao/astraflow@feat/moe-r3-v2-extras, stacked on this branch; not part of this PR.

Using it

bash examples/math/qwen3-30b-a3b-m2po/scripts/1_astraflow.sh        # CPU
bash examples/math/qwen3-30b-a3b-m2po/scripts/2_raas.sh             # GPUs 4,5
bash examples/math/qwen3-30b-a3b-m2po/scripts/3_trainer_model0.sh   # GPUs 0-3

yaml/experiment_b200_1node.yaml + yaml/raas_b200_1node.yaml; the README lists the prerequisites and what to watch during bring-up.

Scope

Rewritten on 2026-09-08 from +3004/-93 (35 files) to +904/-38 (30 files) with the same behaviour for the recipe: the rollout side no longer needs the model config (the trainer views the flat mask with its own layer count), the replay context lost its debug stage machine and memory-release bookkeeping, validation of internally produced tensors and duplicated guards were removed in favour of the contract above, and the mbridge compatibility module collapsed to two lines. The previous revision is preserved as feat/moe-r3-full / feat/moe-r3 (d46edba) on the fork.