diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index e579e837960..8f39d5a2279 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -1564,6 +1564,36 @@ def _should_use_async_rollouts(master_config: MasterConfig) -> bool: return False +def _preserve_router_replay_routed_experts( + target: BatchedDataDict, + flat_messages: BatchedDataDict, + policy_config: PolicyConfig, +) -> None: + """Carry rollout-recorded routes into policy worker inputs when R3 is enabled.""" + if router_replay_enabled(policy_config) and "routed_experts" in flat_messages: + target["routed_experts"] = flat_messages["routed_experts"] + + +def _build_async_grpo_train_data( + flat_messages: BatchedDataDict, + input_lengths: torch.Tensor, + repeated_batch: BatchedDataDict, + policy_config: PolicyConfig, +) -> BatchedDataDict[ClippedPGLossDataDict]: + """Build the async no-TQ policy train batch from flattened rollout messages.""" + train_data = BatchedDataDict[ClippedPGLossDataDict]( + { + "input_ids": flat_messages["token_ids"], + "input_lengths": input_lengths, + "generation_logprobs": flat_messages["generation_logprobs"], + "token_mask": flat_messages["token_loss_mask"], + "sample_mask": repeated_batch["loss_multiplier"], + } + ) + _preserve_router_replay_routed_experts(train_data, flat_messages, policy_config) + return train_data + + def _should_use_nemo_gym(master_config: MasterConfig) -> bool: """Determine if NeMo-Gym should be used for rollouts and validation based on the configuration.""" env_config = master_config.env @@ -2383,11 +2413,9 @@ def grpo_train( # but the train_data whitelist above drops it. Copy it back so # the Megatron worker's train-stage router-replay guard finds # it. Mirrors the TQ producer (sync_rollout_actor.py). - if ( - router_replay_enabled(master_config.policy) - and "routed_experts" in flat_messages - ): - train_data["routed_experts"] = flat_messages["routed_experts"] + _preserve_router_replay_routed_experts( + train_data, flat_messages, master_config.policy + ) train_data.to("cpu") metrics_logging_data["content"] = flat_messages["content"] @@ -2438,11 +2466,9 @@ def grpo_train( # intentionally ignores routed_experts (require_router_replay # =False short-circuits before the field is read), so a # present-but-unused field here is safe. - if ( - router_replay_enabled(master_config.policy) - and "routed_experts" in flat_messages - ): - logprob_data["routed_experts"] = flat_messages["routed_experts"] + _preserve_router_replay_routed_experts( + logprob_data, flat_messages, master_config.policy + ) if not skip_prev_logprobs: train_data["prev_logprobs"] = policy.get_logprobs( @@ -3185,6 +3211,14 @@ def async_grpo_train( assert master_config.loss_fn.use_importance_sampling_correction, ( "Importance sampling correction must be enabled for async GRPO for good convergence due to off-policy samples!" ) + if router_replay_enabled(master_config.policy) and ( + master_config.data_plane or {} + ).get("enabled", False): + raise NotImplementedError( + "policy.router_replay.enabled=true with async GRPO is currently " + "supported only when data_plane.enabled=false. Async + TQ support " + "has not been merged yet." + ) if master_config.grpo["async_grpo"]["max_trajectory_age_steps"] > 1: if not master_config.grpo["async_grpo"].get("in_flight_weight_updates", False): @@ -3600,16 +3634,12 @@ def async_grpo_train( ], ) - # Create training data - # Note: advantages will be computed and added after logprobs are available - train_data = BatchedDataDict[ClippedPGLossDataDict]( - { - "input_ids": flat_messages["token_ids"], - "input_lengths": input_lengths, - "generation_logprobs": flat_messages["generation_logprobs"], - "token_mask": flat_messages["token_loss_mask"], - "sample_mask": repeated_batch["loss_multiplier"], - } + # Create training data. Advantages are added after logprobs. + train_data = _build_async_grpo_train_data( + flat_messages, + input_lengths, + repeated_batch, + master_config.policy, ) train_data.to("cpu") diff --git a/tests/unit/algorithms/test_grpo_router_replay_async.py b/tests/unit/algorithms/test_grpo_router_replay_async.py new file mode 100644 index 00000000000..42d1d030d37 --- /dev/null +++ b/tests/unit/algorithms/test_grpo_router_replay_async.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from nemo_rl.algorithms.grpo import ( + MasterConfig, + _build_async_grpo_train_data, + _default_grpo_save_state, + async_grpo_train, +) +from nemo_rl.distributed.batched_data_dict import BatchedDataDict + + +@pytest.fixture(scope="session", autouse=True) +def init_ray_cluster(): + yield + + +@pytest.fixture(scope="session", autouse=True) +def ray_gpu_monitor(): + class _NoopGpuMonitor: + def _collect_metrics(self): + return {} + + def stop(self): + pass + + yield _NoopGpuMonitor() + + +@pytest.fixture(scope="session", autouse=True) +def session_data(_unit_test_data): + yield _unit_test_data + + +def _make_async_master_config(data_plane=None) -> MasterConfig: + return MasterConfig.model_construct( + **{ + "policy": { + "router_replay": {"enabled": True}, + "generation": { + "backend": "vllm", + "vllm_cfg": {"async_engine": True}, + }, + }, + "loss_fn": SimpleNamespace(use_importance_sampling_correction=True), + "data_plane": data_plane, + } + ) + + +# Keep this focused on async no-TQ batch construction instead of full Ray orchestration. +@pytest.mark.parametrize( + ("policy_config", "expect_routed_experts"), + [ + ({"router_replay": {"enabled": True}}, True), + ({"router_replay": {"enabled": False}}, False), + ], +) +def test_build_async_grpo_train_data_preserves_routed_experts_for_r3( + policy_config, expect_routed_experts +): + routes = torch.arange(1 * 3 * 2 * 4, dtype=torch.int32).reshape(1, 3, 2, 4) + flat_messages = BatchedDataDict( + { + "token_ids": torch.tensor([[1, 2, 3]]), + "generation_logprobs": torch.zeros(1, 3), + "token_loss_mask": torch.tensor([[0, 1, 1]]), + "routed_experts": routes, + } + ) + input_lengths = torch.tensor([3]) + repeated_batch = BatchedDataDict({"loss_multiplier": torch.tensor([1.0])}) + + train_data = _build_async_grpo_train_data( + flat_messages, + input_lengths, + repeated_batch, + policy_config, + ) + + assert torch.equal(train_data["input_ids"], flat_messages["token_ids"]) + assert torch.equal(train_data["input_lengths"], input_lengths) + assert torch.equal( + train_data["generation_logprobs"], flat_messages["generation_logprobs"] + ) + assert torch.equal(train_data["token_mask"], flat_messages["token_loss_mask"]) + assert torch.equal(train_data["sample_mask"], repeated_batch["loss_multiplier"]) + + if expect_routed_experts: + assert torch.equal(train_data["routed_experts"], routes) + else: + assert "routed_experts" not in train_data + + +def test_async_grpo_r3_rejects_data_plane_until_async_tq_exists(): + master_config = _make_async_master_config(data_plane={"enabled": True}) + + with pytest.raises(NotImplementedError, match="data_plane.enabled=false"): + async_grpo_train( + MagicMock(), + MagicMock(), + MagicMock(), + None, + MagicMock(), + MagicMock(), + {"math": MagicMock()}, + None, + MagicMock(), + MagicMock(), + _default_grpo_save_state(), + master_config, + )