feat(trainer): measure how often replayed routes agree with the trainer's routing

src/prime_rl/trainer/model.py

@@ -243,19 +243,40 @@ def get_load_balance_stats(
     per_layer_routing_confidence = []
     block_mlps = list(iter_moe_blocks(model))
     if not block_mlps:
-        return {"max_vio": None, "routing_confidence": None, "tokens_per_expert": torch.empty(0, 0)}
+        return {
+            "max_vio": None,
+            "routing_confidence": None,
+            "router_replay/agreement": None,
+            "router_replay/zero_overlap": None,
+            "tokens_per_expert": torch.empty(0, 0),
+        }
 
     local_tokens_per_expert = torch.stack([block_mlp.tokens_per_expert for block_mlp in block_mlps])
-    layer_stats = [(block_mlp.tokens_per_expert, block_mlp.routing_confidence_sum) for block_mlp in block_mlps]
+    layer_stats = [
+        (block_mlp.tokens_per_expert, block_mlp.routing_confidence_sum, block_mlp.replay_agreement)
+        for block_mlp in block_mlps
+    ]
     if group is not None:
-        sizes = [tokens_per_expert.numel() + 1 for tokens_per_expert, _ in layer_stats]
+        sizes = [tokens_per_expert.numel() + 5 for tokens_per_expert, _, _ in layer_stats]
         packed_stats = torch.cat(
-            [torch.cat((tokens_per_expert, confidence.reshape(1))) for tokens_per_expert, confidence in layer_stats]
+            [
+                torch.cat((tokens_per_expert, confidence.reshape(1), replay))
+                for tokens_per_expert, confidence, replay in layer_stats
+            ]
         )
         dist.all_reduce(packed_stats, op=dist.ReduceOp.SUM, group=group)
-        layer_stats = [(stats[:-1], stats[-1]) for stats in packed_stats.split(sizes)]
-
-    for block_mlp, (tokens_per_expert, routing_confidence_sum) in zip(block_mlps, layer_stats):
+        layer_stats = [(stats[:-5], stats[-5], stats[-4:]) for stats in packed_stats.split(sizes)]
+
+    # Per layer: replayed tokens, checked tokens, checked tokens with no shared expert, summed shared fraction.
+    # Hash-routed layers have nothing to check, so they are left out.
+    learned_replay = [
+        replay for block_mlp, (_, _, replay) in zip(block_mlps, layer_stats) if block_mlp.router.learned_selection
+    ]
+    replay = torch.stack(learned_replay) if learned_replay else None
+    for block_mlp in block_mlps:
+        block_mlp.replay_agreement.zero_()
+
+    for block_mlp, (tokens_per_expert, routing_confidence_sum, _) in zip(block_mlps, layer_stats):
         num_routed_tokens = tokens_per_expert.sum() / block_mlp.router.top_k
         tokens_per_expert = tokens_per_expert.sort(dim=0, descending=True).values[block_mlp.router.top_k :]
         balanced_load = tokens_per_expert.mean()
@@ -267,9 +288,15 @@ def get_load_balance_stats(
 
         block_mlp.tokens_per_expert.zero_()
         block_mlp.routing_confidence_sum.zero_()
+    # Every rank replays or none does (the trainer requires routes on every micro batch when
+    # replay is on), so all ranks agree on whether these stats exist.
+    replaying = replay is not None and bool(replay[:, 0].sum() > 0)
+    checked = replay[:, 1].clamp(min=1) if replaying else None
     return {
         "max_vio": torch.stack(per_layer_max_vio),
         "routing_confidence": torch.stack(per_layer_routing_confidence),
+        "router_replay/agreement": replay[:, 3] / checked if replaying else None,
+        "router_replay/zero_overlap": replay[:, 2] / checked if replaying else None,
         "tokens_per_expert": local_tokens_per_expert,
     }
 
@@ -289,8 +316,8 @@ def get_global_moe_stats(
     for name, values in load_balance_stats.items():
         if values is None:
             continue
-        value = values.max() if name == "max_vio" else values.mean()
-        if name == "max_vio":
+        value = values.max() if name in ("max_vio", "router_replay/zero_overlap") else values.mean()
+        if name in ("max_vio", "router_replay/zero_overlap"):
             dist.all_reduce(value, op=dist.ReduceOp.MAX, group=dp_cp_group)
         else:
             dist.all_reduce(value, op=dist.ReduceOp.SUM, group=dp_cp_group)
@@ -905,6 +932,7 @@ def _reset_runtime_moe_buffers(model: nn.Module) -> None:
         if isinstance(module, MoE) and module.tokens_per_expert.device.type != "meta":
             module.tokens_per_expert.zero_()
             module.routing_confidence_sum.zero_()
+            module.replay_agreement.zero_()
 
 
 def _validate_flash_attn_4_installed() -> None:

src/prime_rl/trainer/models/deepseek_v4/moe.py

@@ -24,6 +24,8 @@ class DeepseekV4HashRouter(TokenChoiceTopKRouter):
     unbuilt, keeping the state dict aligned with HF's.
     """
 
+    learned_selection = False
+
     def __init__(self, *, vocab_size: int, **router_kwargs) -> None:
         super().__init__(**router_kwargs, selection_bias=False)
         self.register_buffer("tid2eid", torch.zeros(vocab_size, self.top_k, dtype=torch.long), persistent=True)

src/prime_rl/trainer/models/layers/moe.py

@@ -44,6 +44,16 @@ def _record_moe_routing_statistics_fake(
     return None
 
 
+@torch.library.custom_op("prime_rl::record_router_replay_agreement", mutates_args=("replay_agreement",))
+def record_router_replay_agreement(replay_agreement: torch.Tensor, counts: torch.Tensor) -> None:
+    replay_agreement.add_(counts)
+
+
+@record_router_replay_agreement.register_fake
+def _record_router_replay_agreement_fake(replay_agreement: torch.Tensor, counts: torch.Tensor) -> None:
+    return None
+
+
 @dataclass
 class MoEArgs:
     num_experts: int = 8
@@ -146,6 +156,10 @@ class TokenChoiceTopKRouter(nn.Module):
         topk_sorted (bool): Whether selected experts are returned in descending score order.
     """
 
+    learned_selection = True
+    """Whether inference picks the experts as the top-k of this router's scores (unlike a hash
+    router's token-id lookup), so replayed routes can be checked against the router's own."""
+
     def __init__(
         self,
         dim: int,
@@ -180,14 +194,14 @@ def forward(
         self,
         x: torch.Tensor,
         routed_experts: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
         """
         Args:
             x (torch.Tensor): Input tensor with shape ``(bs*slen, dim)``.
             routed_experts (torch.Tensor | None, optional): Optional tensor with shape ``(bs * slen, top_k)``.
 
         Returns:
-            tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+            tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
                 - top_scores (torch.Tensor):
                     Routing scores for selected experts with shape ``(bs*slen, top_k)``.
                 - selected_experts_indices (torch.Tensor):
@@ -196,6 +210,9 @@ def forward(
                     Number of tokens assigned to each expert with shape ``(num_experts,)``.
                 - routing_confidence_sum (torch.Tensor):
                     Sum over tokens of the selected-expert probability mass before route normalization/scaling.
+                - own_experts (torch.Tensor | None):
+                    With replayed routes and a learned selection, the router's own top-k with shape
+                    ``(bs*slen, top_k)``; otherwise None.
         """
         # scores shape (bs*slen, num_experts)
         assert routed_experts is None or routed_experts.shape[-1] == self.top_k, (
@@ -223,9 +240,14 @@ def forward(
         # NOTE: selection biases are only used for routing. The gating value
         #       top_scores is still derived from the original scores/logits.
 
+        own_experts = None
         if routed_experts is not None:
             top_scores = scores.gather(dim=1, index=routed_experts)
             selected_experts_indices = routed_experts
+            if self.learned_selection:
+                with torch.no_grad():
+                    selection_scores = scores if self.selection_bias is None else scores + self.selection_bias
+                    own_experts = selection_scores.topk(self.top_k, dim=1, sorted=False).indices
         elif self.force_balanced:
             num_tokens = scores.shape[0]
             arange = torch.arange(num_tokens * self.top_k, device=scores.device)
@@ -266,7 +288,7 @@ def forward(
             max=self.num_experts,
         ).to(torch.int64)
 
-        return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum
+        return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum, own_experts
 
     def init_weights(self, init_std: float):
         nn.init.trunc_normal_(self.gate.weight, mean=0.0, std=init_std)
@@ -344,6 +366,10 @@ def __init__(
             persistent=False,
         )
         self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False)
+        # Router replay check: [replayed tokens, of them checked (see ``_replay_agreement``),
+        # checked tokens with no expert in common with the router's own top-k, summed fraction
+        # of shared experts].
+        self.register_buffer("replay_agreement", torch.zeros(4, dtype=torch.float32), persistent=False)
 
     def set_token_dispatcher(self, token_dispatcher: TokenDispatcher) -> None:
         self.token_dispatcher = token_dispatcher
@@ -383,6 +409,7 @@ def forward(
             selected_experts_indices,
             num_tokens_per_expert,
             routing_confidence_sum,
+            own_experts,
         ) = self.router(x, routed_experts=routed_experts)
 
         # Accumulate expert usage for selection-bias updates and metrics.
@@ -393,6 +420,10 @@ def forward(
                 num_tokens_per_expert,
                 routing_confidence_sum,
             )
+            if own_experts is not None:
+                record_router_replay_agreement(
+                    self.replay_agreement, self._replay_agreement(own_experts, routed_experts)
+                )
 
         routed_output = self.token_dispatcher.run(
             self.prepare_expert_input(x),
@@ -415,6 +446,27 @@ def forward(
 
         return routed_output.reshape(bs, slen, dim)
 
+    def _replay_agreement(self, own: torch.Tensor, routed_experts: torch.Tensor) -> torch.Tensor:
+        """Compare replayed routes with the router's own top-k. Trainer and sampler weights
+        differ slightly, so most tokens share most experts; tokens sharing none point at
+        routes misaligned with their tokens. Rows with a repeated expert id are padding or
+        capture holes (a real top-k has distinct ids) and are skipped. So are rows equal to
+        the previous row: inference never forwards a sample's final token, so its row is a
+        copy of the previous one, which shares no expert with the trainer's in most early layers."""
+        shared = (own.unsqueeze(-1) == routed_experts.unsqueeze(-2)).any(dim=-1).sum(dim=-1)
+        sorted_ids = routed_experts.sort(dim=-1).values
+        valid = (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all(dim=-1)
+        valid[1:] &= (routed_experts[1:] != routed_experts[:-1]).any(dim=-1)
+        top_k = routed_experts.shape[-1]
+        return torch.stack(
+            (
+                torch.full((), valid.numel(), device=valid.device, dtype=torch.float32),
+                valid.sum(),
+                (valid & (shared == 0)).sum(),
+                (shared * valid).sum() / top_k,
+            )
+        ).float()
+
     def init_weights(
         self,
         init_std: float,
@@ -428,5 +480,6 @@ def init_weights(
         with torch.device(buffer_device):
             self.tokens_per_expert = torch.zeros(self.experts.num_experts, dtype=torch.float32)
             self.routing_confidence_sum = torch.tensor(0.0, dtype=torch.float32)
+            self.replay_agreement = torch.zeros(4, dtype=torch.float32)
             if self.router.selection_bias is not None:
                 self.router.selection_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32)

src/prime_rl/trainer/rl/train.py

@@ -75,6 +75,9 @@
 from prime_rl.utils.pathing import resolve_latest_ckpt_step
 from prime_rl.utils.utils import clean_exit
 
+# Warn when more than this fraction of a layer's replayed tokens shares no expert with the trainer's routing
+MAX_REPLAY_ZERO_OVERLAP = 0.01
+
 
 @clean_exit
 def train(config: TrainerConfig):
@@ -660,7 +663,17 @@ def train(config: TrainerConfig):
             step_message += f" | Max Vio {tensor_stats['max_vio/mean']:.4f}"
         if "routing_confidence/mean" in tensor_stats:
             step_message += f" | Routing Conf. {tensor_stats['routing_confidence/mean']:.4f}"
+        if "router_replay/agreement/mean" in tensor_stats:
+            step_message += f" | Replay Agreement {tensor_stats['router_replay/agreement/mean']:.1%}"
         logger.success(step_message)
+        # Trainer and sampler weights differ slightly, so a replayed route rarely shares no
+        # expert with the router's own choice. Many such tokens mean the routes are misaligned.
+        zero_overlap = tensor_stats.get("router_replay/zero_overlap/max", 0.0)
+        if zero_overlap > MAX_REPLAY_ZERO_OVERLAP:
+            logger.warning(
+                f"{zero_overlap:.1%} of replayed tokens in some MoE layer share no expert with the trainer's own "
+                "routing. The replayed routes are likely misaligned with their tokens."
+            )
 
         # Log performance metrics
         perf_metrics = {

tests/unit/train/models/test_moe.py

@@ -4,7 +4,7 @@
 
 from prime_rl.configs.trainer import ModelConfig
 from prime_rl.trainer.distributed.token_dispatcher import LocalTokenDispatcher
-from prime_rl.trainer.model import compute_expert_load_stats, is_tt_moe_model
+from prime_rl.trainer.model import compute_expert_load_stats, get_load_balance_stats, is_tt_moe_model
 from prime_rl.trainer.models.fusions import fuse_gate_up_projections
 from prime_rl.trainer.models.layers.activations import ActivationDispatch, ClampedSilu
 from prime_rl.trainer.models.layers.expert_compute import BF16ExpertCompute, GroupedGemmExpertCompute
@@ -251,3 +251,33 @@ def test_expert_load_stats():
     assert stats.keys() == expected.keys()
     for name, value in expected.items():
         assert stats[name].item() == pytest.approx(value)
+
+
+@pytest.mark.parametrize("learned_selection", [True, False])
+def test_router_replay_agreement_stats(learned_selection):
+    moe = MoE.from_args(MoEArgs(num_experts=4), dim=4, hidden_dim=8, shared_expert=None)
+    moe.router.learned_selection = learned_selection
+    moe.tokens_per_expert.fill_(1.0)
+    # 4 replayed tokens, all checked, 1 shares no expert with the router's own top-k, 3 shared in total
+    moe.replay_agreement.copy_(torch.tensor([4.0, 4.0, 1.0, 3.0]))
+    model = torch.nn.Module()
+    model.model = torch.nn.Module()
+    model.model.layers = torch.nn.ModuleList([torch.nn.Module()])
+    model.model.layers[0].mlp = moe
+
+    stats = get_load_balance_stats(model)
+
+    if learned_selection:
+        assert stats["router_replay/agreement"].tolist() == [0.75]
+        assert stats["router_replay/zero_overlap"].tolist() == [0.25]
+    else:
+        assert stats["router_replay/agreement"] is None
+    assert not moe.replay_agreement.any()
+
+
+def test_router_replay_agreement_skips_repeated_rows():
+    moe = MoE.from_args(MoEArgs(num_experts=4, top_k=2), dim=4, hidden_dim=8, shared_expert=None)
+    own = torch.tensor([[0, 1], [2, 3], [2, 3]])
+    # last row: the inference placeholder for a sample's final token, a copy of the previous row
+    routed = torch.tensor([[0, 1], [1, 0], [1, 0]])
+    assert moe._replay_agreement(own, routed).tolist() == [3.0, 2.0, 1.0, 1.0]