fix(router-replay): allow routeless batches and token-free groups

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -41,6 +41,7 @@
 )
 
 import ray
+import torch
 
 from nemo_rl.algorithms.async_utils.interfaces import ReplayBufferProtocol
 from nemo_rl.data_plane import KVBatchMeta
@@ -1264,11 +1265,34 @@ async def commit(
             prompt_idx=record.prompt_idx,
         )
         if self._require_routed_experts and ROUTED_EXPERTS_FIELD not in fields:
-            raise RuntimeError(
-                "policy.router_replay.enabled=true requires routed_experts in "
-                "the SingleController rollout payload, but payload packing did "
-                "not produce that field. Check vLLM routed-expert capture and "
-                "the async message-log flattening path."
+            # A group whose every completion is token-free (every sandbox failed
+            # before the model spoke and the rollouts were masked) has no
+            # generation to route: it is legitimately routeless and must not take
+            # the run down. Masked placeholders may still carry assistant tokens,
+            # so "every completion is masked" counts as token-free too. Only
+            # groups with trainable tokens are required to carry routes.
+            # BatchedDataDict is tensordict-backed: `in` would coerce it to bool.
+            token_mask = train_batch.get("token_mask", None)
+            mask_sample = train_batch.get("mask_sample", None)
+            all_masked = mask_sample is not None and bool(
+                torch.as_tensor(mask_sample).bool().all().item()
+            )
+            has_generation_tokens = (
+                token_mask is None or bool(torch.as_tensor(token_mask).sum().item() > 0)
+            ) and not all_masked
+            if has_generation_tokens:
+                raise RuntimeError(
+                    "policy.router_replay.enabled=true requires routed_experts in "
+                    "the SingleController rollout payload, but payload packing did "
+                    "not produce that field. Check vLLM routed-expert capture and "
+                    "the async message-log flattening path."
+                )
+            print(
+                "router_replay: committing a token-free group without "
+                f"routed_experts (group_id={group_id}, prompt_idx={record.prompt_idx}, "
+                f"completions={len(record.completions)}); Megatron has nothing to "
+                "route for it",
+                flush=True,
             )
         trace_rollout_payload(keys=sample_ids, data=train_batch)
         async with self._data_plane_checkpoint_barrier.mutation("group_commits") as cut:

nemo_rl/experience/sync_rollout_actor.py

@@ -311,7 +311,14 @@ def rollout_to_tq(
         router_replay_enabled = bool(
             (cfg.policy.get("router_replay") or {}).get("enabled", False)
         )
-        if router_replay_enabled and ROUTED_EXPERTS_FIELD not in flat:
+        # A token-free batch (every rollout masked before the model spoke) has
+        # nothing to route; only batches with trainable tokens must carry routes.
+        has_generation_tokens = bool(flat["token_loss_mask"].sum().item() > 0)
+        if (
+            router_replay_enabled
+            and ROUTED_EXPERTS_FIELD not in flat
+            and has_generation_tokens
+        ):
             raise RuntimeError(
                 "policy.router_replay.enabled=true requires routed_experts in "
                 "the rollout bulk payload, but rollout flattening did not "

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -194,13 +194,17 @@ def _should_use_router_replay(
         return False
     if "routed_experts" in data:
         return True
-    raise RuntimeError(
-        "policy.router_replay.enabled=true requires routed_experts for "
-        f"{stage}, but the fetched batch does not contain that field. This "
-        "usually means the TQ schema, field selection, or rollout write path "
-        "stopped carrying routed_experts. Reference-logprob intentionally skips "
-        "routed_experts; prev-logprob and train must not."
+    # A fetched batch can legitimately lack routes: rows restored from a
+    # checkpoint written before router replay was enabled, or token-free
+    # (masked) prompt groups committed without routes. Megatron then routes
+    # those tokens natively, exactly like the documented missing-route
+    # fallback does per token; it is not a schema failure.
+    print(
+        "router_replay: fetched batch has no routed_experts for "
+        f"{stage}; using Megatron's native router for this batch",
+        flush=True,
     )
+    return False
 
 
 def _model_self_packs_for_cp(model: Any) -> bool:

tests/unit/models/megatron/test_router_replay.py

@@ -771,13 +771,14 @@ def test_router_replay_worker_guard_requires_routed_experts_when_expected():
         stage="reference-logprob",
         require=False,
     )
-    with pytest.raises(RuntimeError, match="requires routed_experts for prev-logprob"):
-        _should_use_router_replay(
-            enabled=True,
-            data=data,
-            stage="prev-logprob",
-            require=True,
-        )
+    # A routeless fetched batch (pre-R3 checkpoint rows, token-free groups) falls
+    # back to native routing instead of failing the run.
+    assert not _should_use_router_replay(
+        enabled=True,
+        data=data,
+        stage="prev-logprob",
+        require=True,
+    )
 
     data["routed_experts"] = torch.zeros(2, 4, 3, 2, dtype=torch.int32)
     assert _should_use_router_replay(

tests/unit/single_controller/test_tq_replay_buffer.py

@@ -730,6 +730,44 @@ def test_commit_requires_routed_experts_before_tq_write(self):
         assert dp.depth() == 0
         assert buf.ready_list == [False]
 
+    def test_commit_allows_a_token_free_group_without_routed_experts(self, monkeypatch):
+        """All completions masked/token-free: nothing to route, so R3 must not fail the run."""
+
+        def _token_free(
+            record,
+            *,
+            pad_value_dict,
+            include_message_violation_fields,
+            per_token_multimodal_fields: tuple[str, ...] = (),
+        ):
+            del (
+                record,
+                pad_value_dict,
+                include_message_violation_fields,
+                per_token_multimodal_fields,
+            )
+            return BatchedDataDict[Any](
+                {
+                    "input_ids": torch.ones((_N_GENS, 3), dtype=torch.long),
+                    "input_lengths": torch.full((_N_GENS,), 3, dtype=torch.long),
+                    "token_mask": torch.ones((_N_GENS, 3), dtype=torch.float32),
+                    "mask_sample": torch.ones(_N_GENS, dtype=torch.bool),
+                    "total_reward": torch.zeros(_N_GENS, dtype=torch.float32),
+                }
+            )
+
+        monkeypatch.setattr(_replay_buffer_module, "record_to_train_batch", _token_free)
+        dp = FakeDataPlaneClient()
+        buf = _make_buffer(dp, require_routed_experts=True)
+        group_id = buf.reserve(weight_version=3)
+        _run(
+            buf.commit(
+                group_id, _make_record(), start_weight_version=3, end_weight_version=3
+            )
+        )
+        assert len(dp.put_calls) == 1
+        assert buf.ready_list == [True]
+
     def test_commit_raises_for_unknown_group_id(self):
         dp = FakeDataPlaneClient()
         buf = _make_buffer(dp)