From 043dd9d085d93b6139ab8712bafc902c59605f0a Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Sat, 17 Oct 2026 18:48:20 -0700 Subject: [PATCH 1/2] router replay: a routeless fetched batch falls back to native routing Links 4585549-51 died at the first prev-logprob fetch: the batch mixed rows restored from the pre-R3 step_102 checkpoint and a token-free group, so the routed_experts column was absent and the worker raised. Treat that as the documented missing-route fallback for the whole batch (log it) rather than a schema failure. Test updated. Co-Authored-By: Claude Fable 5.1 (cherry picked from commit 568ca7f5a5f24323492d4d1716579d23451e4864) Signed-off-by: Lin Jia --- .../policy/workers/megatron_policy_worker.py | 16 ++++++++++------ tests/unit/models/megatron/test_router_replay.py | 15 ++++++++------- 2 files changed, 18 insertions(+), 13 deletions(-) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 31ba43761dd..1c4e5a1c68e 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -192,13 +192,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: diff --git a/tests/unit/models/megatron/test_router_replay.py b/tests/unit/models/megatron/test_router_replay.py index 250c91a5b00..3f289cad9ad 100644 --- a/tests/unit/models/megatron/test_router_replay.py +++ b/tests/unit/models/megatron/test_router_replay.py @@ -649,13 +649,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( From 0ea9540c14a5de755585a5d7bde925d3a578b4ec Mon Sep 17 00:00:00 2001 From: Lin Jia Date: Sat, 24 Oct 2026 00:26:19 -0700 Subject: [PATCH 2/2] fix(router-replay): token-free groups and routeless batches route natively instead of raising With policy.router_replay.enabled=true three guards raised RuntimeError for a payload without routed_experts and took the run down: - TQReplayBuffer.commit: a prompt group whose every completion is token-free or masked (every sandbox failed before the model spoke) has no generation to route. It now commits with a one-line log; groups with trainable tokens still raise, listing the payload fields. BatchedDataDict is tensordict-backed, so the guard reads fields with .get, not 'in'. - SyncRolloutActor: the same check on the synchronous path, keyed on token_loss_mask. - The Megatron worker's _should_use_router_replay raised for a fetched batch without routes (rows restored from a checkpoint written before Router Replay was enabled); it now logs and uses the native router for that batch, as the documented per-token missing-route fallback does. Tests: test_tq_replay_buffer.py (token-free group commits), test_router_replay.py (worker fallback). Signed-off-by: Lin Jia --- .../algorithms/async_utils/replay_buffer.py | 34 ++++++++++++++++--- nemo_rl/experience/sync_rollout_actor.py | 9 ++++- .../test_tq_replay_buffer.py | 27 +++++++++++++++ 3 files changed, 64 insertions(+), 6 deletions(-) diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 54567cd6055..e0d9e4da506 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/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 @@ -1259,11 +1260,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: diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index 6496d3b1101..7ff5ad8e9c5 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/nemo_rl/experience/sync_rollout_actor.py @@ -300,7 +300,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 " diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 1e7f9ddc84f..77cce7e1f68 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -721,6 +721,33 @@ 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): + del record, pad_value_dict, include_message_violation_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)