[faresobeid]

Adds a router-replay check: each MoE layer compares the replayed routes with the router's own top-k and logs how often they agree. A content check (no holes, ids in range) cannot see routes that are misaligned with their tokens, for example a one-token shift in packing, CP sharding or concatenated multi-turn payloads. miles caught such a bug this way: miles#3146, where under BSHD 3 of 4 tokens got the wrong experts.

What

• TokenChoiceTopKRouter.forward, when it gets routed_experts, also returns its own top-k, taken from the scores it already computed (under no_grad; one extra top-k per layer, no extra gate GEMM or router all-gather). MoE.forward accumulates, per layer:
  • replayed tokens;
  • checked tokens;
  • checked tokens that share no expert with the router's own top-k;
  • the summed fraction of shared experts.
• Rows with a repeated expert id are skipped, because a real top-k has distinct ids. These are padding and capture holes.
• Rows equal to the previous row are skipped too. Inference never forwards a sample's final token, so verifiers fills its row with a copy of the previous one. Without this skip, that copy shares no expert with the trainer's top-k in early layers: 94% of final tokens in layer 0 on Qwen3-30B-A3B. That is about 1% of all tokens, so the warning fired on a healthy run. Real adjacent rows are rarely identical: 0.04% of rows in the worst layer.
• The counts go through a custom op, like the existing routing statistics, and are built on device (no host sync).
• Hash-routed DeepSeek-V4 layers return no own top-k (learned_selection = False), because inference does not pick their experts from the scores.
• get_load_balance_stats reduces the counts with the other MoE stats:
  • router_replay/agreement: mean shared fraction, averaged over layers and ranks;
  • router_replay/zero_overlap: fraction of tokens sharing no expert, max over layers and ranks.
  • Both leave out hash-routed layers, so they do not pull the mean down.
  • The step log shows Replay Agreement.
  • The trainer warns when zero_overlap exceeds 1 % (MAX_REPLAY_ZERO_OVERLAP, as miles suggests).
• Without replay nothing changes. The stats only exist when routes are passed, which with enable_router_replay is every micro batch on every rank, so all ranks agree on the collective.

Checks

• pytest tests/unit/train/models/test_moe.py tests/unit/orchestrator tests/unit/test_configs.py -m "not gpu": 361 passed. ruff is clean.

• CPU check of the counts after moving the own top-k into the router: 6 replayed rows, one with a repeated id, gives [6, 5, 0, 5].

• GPU, 1 node, configs/ci/integration/reverse-text-moe (samsja/mini-glm-moe) with --trainer.enable-router-replay, 10 steps:

  • Replay Agreement 99.4% on every step, and no warning.
  • Same run with a throwaway patch that rolls the replayed routes by one token: Replay Agreement 59.3%. Every step logged 3.7% of replayed tokens in some MoE layer share no expert ..., and Mismatch KL went from 0.0001 to 0.027.

• GPU, Qwen3-30B-A3B reverse-text, 1 node (EP 4, compile on), max_off_policy_steps = 0, 6 steps: Replay Agreement is 98.9-99.0% and zero_overlap is 0 on every step. Before the placeholder skip, the same setup showed 98.4% and 1.0-1.2%, and the warning fired on every step. A per-token dump of steps 1-8 found all zero-overlap tokens at the final token of a sample; that row equals the previous row in all 48 layers (605/605).

Likely conflicts

• fix(trainer): distribute routing replay padding across experts #3350 (spread replay padding across experts): padding rows would no longer repeat an id, so they would be counted as checked tokens.
• feat: pass router-replay ids and sampling masks by handle on multi-node runs #3902 (by-handle payloads): it touches the replay data path, not this code.

---

Note

Medium Risk
Extends the MoE forward and distributed stat packing/reduction paths when router replay is enabled, but only adds metrics and logging without changing routing or loss computation.

Overview
Adds router replay agreement metrics so training can detect when replayed expert routes are misaligned with the tokens they are applied to (e.g. packing or sharding bugs), beyond simple validity checks on route tensors.

When routed_experts is supplied, learned-selection routers also compute their own top-k under no_grad and MoE accumulates per-layer counts (replayed/checked tokens, zero-overlap tokens, shared-expert fraction), skipping padding rows and hash-routed DeepSeek-V4 layers (learned_selection = False). These roll into existing MoE load-balance reduction as router_replay/agreement (mean over layers/ranks) and router_replay/zero_overlap (max over layers/ranks). The step log shows replay agreement and warns if any layer exceeds 1% zero overlap.

No behavior change when replay is off; overhead is one extra top-k per layer only during replay.

^{Reviewed by Cursor Bugbot for commit 798da15. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Stale Bugbot comment from a previous run.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 04d7d8b. Configure here.}