feat(moe): add Rollout Routing Replay (R3) for MoE RL training

[khazic]

What

Adds opt-in Rollout Routing Replay (R3) to the MoE subsystem.

In on-policy RL on a Mixture-of-Experts model, the rollout (inference) engine and the training engine compute the router's top-k expert selection independently. Small numerical differences flip a fraction of routing decisions per layer, and that compounds across layers until most tokens train under a different expert assignment than they were rolled out with. The mismatch breaks the importance-sampling assumption behind GRPO/GSPO and destabilizes training.

R3 removes the mismatch: capture the top-k expert selection during the rollout-equivalent forward, then replay that exact selection during the training forward. Only the discrete selection is replayed. The router logits and their softmax/sigmoid are still recomputed from the live router weights, so the gradient continues to flow into the router. This mirrors Megatron-LM's moe_enable_routing_replay.

Changes

• components/moe/router_replay.py (new): a RouterReplay handle (RECORD / REPLAY modes) with a process-global per-layer registry plus record() / replay() context managers, and the replay_selection helper.
• MoEConfig.enable_routing_replay flag (default False).
• Hooks in the shared Gate (every score_func branch) and the custom Gemma4Gate, so both record and replay their final expert selection.
• Unit tests at 100% coverage of router_replay.py.

Default path stays a no-op

When enable_routing_replay=False (the default), no RouterReplay handle is created and the gate forward is unchanged: replay_selection returns the original index tensor. The two branches that previously consumed topk's returned values directly re-gather only when replay actually swapped the selection, guarded by an identity check, so the replay-off and record paths add no extra tensor work.

Usage

from nemo_automodel.components.moe.router_replay import RouterReplay

# Capture the selection on the rollout-equivalent forward.
with RouterReplay.record():
    model(batch)
captured = RouterReplay.collect()      # one tensor per MoE layer, in layer order

# Replay it on the training forward over the same tokens.
with RouterReplay.replay(captured):
    loss = model(batch)
loss.backward()

Testing

Unit tests: tests/unit_tests/moe/test_router_replay.py, 27 tests, 100% coverage of router_replay.py. They cover the record/replay roundtrip across every score_func branch, the default-path no-op, gradient flow to the router under replay, multi-layer distribute/collect, error handling, and the Gemma4Gate path.

Beyond the unit tests, I exercised the real Automodel model classes on GPU (record on batch A, replay on a different batch B):

Check	Qwen3-MoE (Qwen3MoeForCausalLM, 3 MoE layers)	Gemma4-MoE (Gemma4MoE + Gemma4Gate)
Natural routing of A and B differs	yes	yes
Replay reproduces A's selection across all layers	yes	yes
Replay overrides B's natural selection	yes	yes
Router weight receives a gradient under replay	yes (CE-loss backward)	yes

Test environment

• 1x NVIDIA A800-80GB
• Python 3.12, PyTorch 2.12.0+cu130, Transformers 5.13.0.dev0
• ruff format and ruff check clean

Addresses #2454.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

/ok to test 20552e3

[HuiyingLi]

/claude review

[claude]

LGTM — clean design, well-tested, no issues found. The record/replay separation, identity-check optimization, and context-manager cleanup are all correct.

[HuiyingLi]

I tested PR #2797 locally on 8x H100s.

Validation done:

• tests/unit_tests/moe/test_router_replay.py -q: 27 passed.
• Qwen3-MoE 30B-A3B EP8 training smoke test with the existing 8k mock config: 5/5 steps passed.
• Qwen3-MoE EP8 with real Gate and enable_routing_replay=true: 1/1 step passed.
• Added a temporary hook-based EP8 replay check on MoE + ExpertParallel:
  • recorded routing from rollout input
  • verified train input naturally routed differently
  • replayed rollout routing on train input
  • verified in-forward gate indices matched recorded rollout indices
  • verified replay changed the routed output
  • backward completed with finite grads