fix(genrm): propagate prompt cohort retry identity

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -100,6 +100,7 @@ def collect_trajectories(
             generation_config=generation_config,
             # This utility consumes the Tables below to write its trajectory JSONL.
             log_full_result_tables=True,
+            num_generations_per_prompt=1,
             max_rollout_turns=None,
             greedy=False,
         )

nemo_rl/algorithms/distillation.py

@@ -803,6 +803,9 @@ def _distillation_train_impl(
                             task_to_env=task_to_env,
                             max_seq_len=None,
                             generation_config=generation_config,
+                            num_generations_per_prompt=(
+                                master_config.distillation.num_generations_per_prompt
+                            ),
                             log_full_result_tables=should_log_nemo_gym_full_result_tables(
                                 wandb_enabled=master_config.logger["wandb_enabled"],
                                 wandb_config=master_config.logger["wandb"],
@@ -1289,6 +1292,7 @@ def validate(
                     task_to_env=val_task_to_env,
                     max_seq_len=None,
                     generation_config=generation_config,
+                    num_generations_per_prompt=1,
                     log_full_result_tables=should_log_nemo_gym_full_result_tables(
                         wandb_enabled=master_config.logger["wandb_enabled"],
                         wandb_config=master_config.logger["wandb"],

nemo_rl/algorithms/grpo.py

@@ -3176,6 +3176,9 @@ def _grpo_train_impl(
                                 "max_total_sequence_length"
                             ],
                             generation_config=generation_config,
+                            num_generations_per_prompt=(
+                                master_config.grpo.num_generations_per_prompt
+                            ),
                             log_full_result_tables=should_log_nemo_gym_full_result_tables(
                                 wandb_enabled=master_config.logger["wandb_enabled"],
                                 wandb_config=master_config.logger["wandb"],
@@ -4195,6 +4198,7 @@ def validate(
                     task_to_env=val_task_to_env,
                     max_seq_len=master_config.policy["max_total_sequence_length"],
                     generation_config=generation_config,
+                    num_generations_per_prompt=val_num_generations_per_prompt,
                     sampling_params=val_sampling_params,
                     log_full_result_tables=should_log_nemo_gym_full_result_tables(
                         wandb_enabled=master_config.logger["wandb_enabled"],

nemo_rl/algorithms/ppo.py

@@ -1463,6 +1463,9 @@ def ppo_train(
                             task_to_env=task_to_env,
                             max_seq_len=None,
                             generation_config=generation_config,
+                            num_generations_per_prompt=(
+                                master_config.ppo.num_generations_per_prompt
+                            ),
                             log_full_result_tables=should_log_nemo_gym_full_result_tables(
                                 wandb_enabled=master_config.logger["wandb_enabled"],
                                 wandb_config=master_config.logger["wandb"],

nemo_rl/experience/rollouts.py

@@ -2457,6 +2457,7 @@ async def run_async_nemo_gym_rollout(
     sampling_params: Optional[GenerationSamplingParams] = None,
     deduplicate_multimodal_data: bool = False,
     debug_payload_metrics: bool = False,
+    identity_num_generations: Optional[int] = None,
 ) -> AsyncGenerator[NemoGymRolloutResult, None]:
     """Stream complete NeMo-Gym prompt groups in group-completion order.
 
@@ -2495,6 +2496,10 @@ async def run_async_nemo_gym_rollout(
             remote Gym return and restore the exact original payload locally.
         debug_payload_metrics: Emit logical, physical, and serialized media
             payload metrics at the Gym Ray boundary.
+        identity_num_generations: Number of contiguous rows sharing one logical
+            prompt-group identity. Defaults to ``num_generations``. Synchronous
+            callers set this independently because they collect the full batch as
+            one result while preserving per-prompt GenRM cohort identities.
 
     Yields:
         ``NemoGymRolloutResult`` objects in prompt-group completion order. Rows
@@ -2505,10 +2510,11 @@ async def run_async_nemo_gym_rollout(
         AssertionError: If an unsupported generation option is requested.
         TypeError: If a row lacks a valid ``responses_create_params`` dictionary or
             the actor returns a non-integer row index.
-        ValueError: If ``num_generations`` is not positive, the batch is empty or
-            not divisible by ``num_generations``, ``returns_entire_batch`` has an
-            incompatible size, a streamed row index is out of range or duplicated,
-            a prompt group mixes agents, or its task indices disagree.
+        ValueError: If either generation count is not positive, the batch is empty
+            or not divisible by the relevant generation count,
+            ``returns_entire_batch`` has an incompatible size, a streamed row index
+            is out of range or duplicated, a prompt group mixes agents, or its task
+            indices disagree.
         RuntimeError: If the actor fails, returns NaN generation logprobs, ends the
             stream before all expected rows arrive, or produces no final group.
     """
@@ -2555,12 +2561,20 @@ async def run_async_nemo_gym_rollout(
     )
     if num_generations <= 0:
         raise ValueError("num_generations must be greater than zero")
+    if identity_num_generations is None:
+        identity_num_generations = num_generations
+    if identity_num_generations <= 0:
+        raise ValueError("identity_num_generations must be greater than zero")
     if not nemo_gym_rows:
         raise ValueError("NeMo-Gym rollout batch must not be empty")
     if len(nemo_gym_rows) % num_generations != 0:
         raise ValueError(
             "NeMo-Gym rollout batch size must be divisible by num_generations"
         )
+    if len(nemo_gym_rows) % identity_num_generations != 0:
+        raise ValueError(
+            "NeMo-Gym rollout batch size must be divisible by identity_num_generations"
+        )
     if returns_entire_batch and len(nemo_gym_rows) != num_generations:
         raise ValueError(
             "returns_entire_batch requires num_generations to equal the batch size"
@@ -2584,7 +2598,7 @@ async def run_async_nemo_gym_rollout(
             nemo_gym_rows,
             generation_config,
             sampling_params,
-            num_generations,
+            identity_num_generations,
         )
         accumulator = _NemoGymStreamAccumulator(
             rows=nemo_gym_rows,
@@ -2696,6 +2710,8 @@ def run_nemo_gym_rollout_sync(
     task_to_env: dict[str, EnvironmentInterface],
     generation_config: GenerationConfig,
     log_full_result_tables: bool,
+    *,
+    num_generations_per_prompt: int,
     max_seq_len: Optional[int] = None,
     max_rollout_turns: Optional[int] = None,
     greedy: bool = False,
@@ -2723,6 +2739,9 @@ def run_nemo_gym_rollout_sync(
         generation_config: Sampling parameters forwarded to every NeMo-Gym row.
         log_full_result_tables: Whether to include complete per-agent result
             payloads as W&B Tables in the rollout metrics.
+        num_generations_per_prompt: Number of contiguous rows belonging to each
+            logical prompt group. This controls Gym/GenRM cohort identity only;
+            the synchronous API still collects and returns the entire input batch.
         max_seq_len: Policy sequence-length limit used for compatibility validation.
         max_rollout_turns: Must be ``None`` because NeMo-Gym owns turn limits.
         greedy: Must be ``False`` because this path does not support greedy mode.
@@ -2759,6 +2778,7 @@ async def _consume_rollout() -> NemoGymRolloutResult:
             task_to_env=task_to_env,
             generation_config=generation_config,
             num_generations=input_batch.size,
+            identity_num_generations=num_generations_per_prompt,
             log_full_result_tables=log_full_result_tables,
             max_seq_len=max_seq_len,
             max_rollout_turns=max_rollout_turns,

nemo_rl/experience/sync_rollout_actor.py

@@ -255,6 +255,7 @@ def rollout_to_tq(
                 max_seq_len=None,
                 max_rollout_turns=None,
                 generation_config=cfg.policy["generation"],
+                num_generations_per_prompt=group_size,
                 log_full_result_tables=should_log_nemo_gym_full_result_tables(
                     wandb_enabled=cfg.logger["wandb_enabled"],
                     wandb_config=cfg.logger["wandb"],

tests/unit/environments/test_nemo_gym_rollout.py

@@ -589,6 +589,7 @@ def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepte
         max_seq_len=_GENERATION_CONFIG["max_total_sequence_length"],
         generation_config=deepcopy(_GENERATION_CONFIG),
         log_full_result_tables=True,
+        num_generations_per_prompt=1,
     )
 
     final_batch = result.final_batch

tests/unit/experience/test_rollouts.py

@@ -2131,8 +2131,10 @@ def test_postprocess_nemo_gym_group_returns_task_index(log_full_result_tables):
     ) is log_full_result_tables
 
 
-def test_run_nemo_gym_rollout_sync_drains_entire_batch(monkeypatch):
-    input_batch = BatchedDataDict({"loss_multiplier": torch.ones(3)})
+def test_run_nemo_gym_rollout_sync_separates_collection_and_identity_groups(
+    monkeypatch,
+):
+    input_batch = BatchedDataDict({"loss_multiplier": torch.ones(4)})
     expected = rollouts_mod.NemoGymRolloutResult(
         input_ids=torch.empty(0),
         final_batch=input_batch,
@@ -2142,6 +2144,7 @@ def test_run_nemo_gym_rollout_sync_drains_entire_batch(monkeypatch):
 
     async def fake_stream(**kwargs):
         assert kwargs["num_generations"] == input_batch.size
+        assert kwargs["identity_num_generations"] == 2
         assert kwargs["returns_entire_batch"] is True
         assert kwargs["log_full_result_tables"] is False
         assert kwargs["deduplicate_multimodal_data"] is True
@@ -2157,13 +2160,60 @@ async def fake_stream(**kwargs):
         task_to_env={},
         generation_config={},
         log_full_result_tables=False,
+        num_generations_per_prompt=2,
         deduplicate_multimodal_data=True,
         debug_payload_metrics=True,
     )
 
     assert actual is expected
 
 
+@pytest.mark.parametrize(
+    ("row_count", "identity_num_generations", "error"),
+    [
+        (2, 0, "identity_num_generations must be greater than zero"),
+        (
+            3,
+            2,
+            "NeMo-Gym rollout batch size must be divisible by identity_num_generations",
+        ),
+    ],
+)
+def test_run_async_nemo_gym_rollout_validates_identity_group_size(
+    row_count,
+    identity_num_generations,
+    error,
+):
+    input_batch = BatchedDataDict(
+        {"extra_env_info": [{"responses_create_params": {}} for _ in range(row_count)]}
+    )
+
+    async def collect():
+        return [
+            result
+            async for result in run_async_nemo_gym_rollout(
+                policy_generation=SimpleNamespace(
+                    cfg={"max_total_sequence_length": 128}
+                ),
+                input_batch=input_batch,
+                tokenizer=None,
+                task_to_env={},
+                generation_config={
+                    "max_new_tokens": 16,
+                    "stop_strings": [],
+                    "stop_token_ids": [],
+                },
+                num_generations=row_count,
+                identity_num_generations=identity_num_generations,
+                log_full_result_tables=False,
+                sampling_params=SimpleNamespace(top_k=0),
+            )
+        ]
+
+    with pytest.raises(ValueError, match=error):
+        asyncio.run(collect())
+
+
 def test_rollout_manager_consumes_stream_and_restores_input_order():
     class _ReadyRef:
         def __init__(self, value):
@@ -2434,6 +2484,7 @@ def test_run_async_nemo_gym_rollout(
         max_seq_len=nemo_gym_vllm_generation.cfg["vllm_cfg"]["max_model_len"],
         generation_config=nemo_gym_vllm_generation.cfg,
         log_full_result_tables=True,
+        num_generations_per_prompt=1,
         max_rollout_turns=None,
         debug_payload_metrics=True,
     )