diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index 6f4659a72a..37a111edd4 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 @@ -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: diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index 33c7914aa6..c7040ae1ac 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/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 " diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index dd530f22e6..8c06eeba04 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/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: diff --git a/tests/unit/models/megatron/test_router_replay.py b/tests/unit/models/megatron/test_router_replay.py index 5cacefe963..724fcbd3fb 100644 --- a/tests/unit/models/megatron/test_router_replay.py +++ b/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( diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index a84267ede4..a28dee60b9 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/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)