fix(grpo): preserve prompt groups across rollout paths

nemo_rl/algorithms/grpo.py

@@ -2370,6 +2370,42 @@ def _apply_mask_sample_filter(repeated_batch: BatchedDataDict[DatumSpec]) -> int
     return num_masked
 
 
+def _prompt_grouping_ids(
+    prompt_token_ids: torch.Tensor,
+    *,
+    num_prompts: int,
+    num_generations_per_prompt: int,
+    use_nemo_gym: bool,
+    group_offset: int = 0,
+) -> torch.Tensor:
+    """Return identities for generations originating from the same prompt."""
+    if not use_nemo_gym:
+        return prompt_token_ids
+
+    return (
+        torch.arange(group_offset, group_offset + num_prompts)
+        .repeat_interleave(num_generations_per_prompt)
+        .unsqueeze(1)
+    )
+
+
+def _replay_prompt_batches(
+    trajectories: list[dict[str, Any]],
+    *,
+    use_explicit_grouping: bool,
+) -> list[BatchedDataDict]:
+    """Prepare replay batches with stable identities for each prompt group."""
+    prompt_batches = []
+    for group_index, trajectory in enumerate(trajectories):
+        prompt_batch = trajectory["batch"]
+        if use_explicit_grouping:
+            prompt_batch["prompt_grouping_ids"] = torch.full(
+                (prompt_batch.size, 1), group_index, dtype=torch.int64
+            )
+        prompt_batches.append(prompt_batch)
+    return prompt_batches
+
+
 def _should_log_nemo_gym_responses(master_config: MasterConfig) -> bool:
     """Whether NeMo Gym is responsible for full response logging.
 
@@ -3005,6 +3041,7 @@ def grpo_train(
         batch_cache: BatchedDataDict[DatumSpec] = None
         # This is the number of batches we processed so far at each step to generate responses whose std is non-zero. Maximum threshold is set by dynamic_sampling_max_gen_batches. Used in the case of dynamic sampling.
         dynamic_sampling_num_gen_batches = 0
+        prompt_grouping_id_offset = 0
 
         # Run grpo/dapo training loop (single-turn)
         for batch in wrapped_dataloader:
@@ -3078,7 +3115,14 @@ def grpo_train(
                         repeated_batch["message_log"],
                         pad_value_dict={"token_ids": tokenizer.pad_token_id},
                     )
-                    input_ids = batched_flat["token_ids"]
+                    prompt_grouping_ids = _prompt_grouping_ids(
+                        batched_flat["token_ids"],
+                        num_prompts=batch.size,
+                        num_generations_per_prompt=master_config.grpo.num_generations_per_prompt,
+                        use_nemo_gym=should_use_nemo_gym(master_config),
+                        group_offset=prompt_grouping_id_offset,
+                    )
+                    prompt_grouping_id_offset += batch.size
 
                 # Generate responses - this updates the LLMMessageLogType in repeated_batch
                 memory_tracker.snapshot_start_of_stage("Generation", dir())
@@ -3194,7 +3238,6 @@ def grpo_train(
                                 master_config.grpo.debug_payload_metrics
                             ),
                         )
-                        input_ids = nemo_gym_rollout_result.input_ids
                         repeated_batch = nemo_gym_rollout_result.final_batch
                         rollout_metrics = nemo_gym_rollout_result.rollout_metrics
                         del nemo_gym_rollout_result
@@ -3284,7 +3327,7 @@ def grpo_train(
                         # Just fix the device id for now
                         device_id = 0
                         baseline, std = calculate_baseline_and_std_per_prompt(
-                            input_ids.cuda(device_id),
+                            prompt_grouping_ids.cuda(device_id),
                             rewards.cuda(device_id),
                             torch.ones_like(rewards).cuda(device_id),
                             leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
@@ -3298,13 +3341,16 @@ def grpo_train(
                         std = std.cpu()
                     else:
                         baseline, std = calculate_baseline_and_std_per_prompt(
-                            input_ids,
+                            prompt_grouping_ids,
                             rewards,
                             torch.ones_like(rewards),
                             leave_one_out_baseline=master_config.grpo.use_leave_one_out_baseline,
                             std_rewards=std_rewards,
                         )
 
+                    # Keep source prompt identity aligned with dynamic-sampling rows.
+                    repeated_batch["prompt_grouping_ids"] = prompt_grouping_ids
+
                     # Apply dynamic sampling to filter prompts with non-zero std (DAPO algorithm)
                     repeated_batch, is_batch_complete, batch_cache, ds_metrics = (
                         dynamic_sampling(
@@ -3341,24 +3387,12 @@ def grpo_train(
                     # Save baseline for logging (before deletion)
                     baseline_for_log = baseline.clone()
 
-                    # Must precede prompt extraction: it reuses the same message
-                    # dicts, so this also protects the prompt flatten below.
+                    # Must precede the training flatten because it reuses the same message dicts.
                     backfill_missing_routed_experts(repeated_batch["message_log"])
 
-                    # Extract original prompt messages using the length field
-                    # This correctly handles multi-turn prompts that contain assistant messages
-                    initial_prompt_message_logs = extract_initial_prompt_messages(
-                        repeated_batch["message_log"],
-                        repeated_batch["length"],
-                    )
-                    prompt_batched_flat, _ = batched_message_log_to_flat_message(
-                        initial_prompt_message_logs,
-                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
-                    )
-                    prompt_ids_for_adv = prompt_batched_flat["token_ids"]
-                    del initial_prompt_message_logs
-                    del prompt_batched_flat
-                    del input_ids
+                    prompt_ids_for_adv = repeated_batch["prompt_grouping_ids"]
+                    del repeated_batch["prompt_grouping_ids"]
+                    del prompt_grouping_ids
                     del baseline
                     del std
 
@@ -4026,6 +4060,7 @@ def grpo_train(
             # Reset the batch and set dynamic_sampling_num_gen_batches to 0
             batch_cache = None
             dynamic_sampling_num_gen_batches = 0
+            prompt_grouping_id_offset = 0
 
             # Clear mem
             memory_tracker.snapshot_start_of_stage("After CPU memory clear", dir())
@@ -5068,8 +5103,13 @@ def _flush_collector_telemetry() -> None:
                         f"✅ Sampled {len(trajectories)} trajectory groups from buffer (avg age: {avg_trajectory_age:.2f} steps)"
                     )
 
-                    # Concatenate per-prompt groups into a single training batch
-                    per_prompt_batches = [t["batch"] for t in trajectories]
+                    # Concatenate per-prompt groups into a single training batch.
+                    # Reconstruct grouping from the replay buffer's prompt groups.
+                    # This also repairs groups loaded from older checkpoints.
+                    per_prompt_batches = _replay_prompt_batches(
+                        trajectories,
+                        use_explicit_grouping=should_use_nemo_gym(master_config),
+                    )
                     repeated_batch = BatchedDataDict.from_batches(
                         per_prompt_batches,
                         allow_missing_packed_tensors=(
@@ -5127,24 +5167,24 @@ def _flush_collector_telemetry() -> None:
                         RLSpanGroup.REWARD, "rl.grpo.reward_calculation", tracer=_tracer
                     ),
                 ):
-                    # Must precede prompt extraction: it reuses the same message
-                    # dicts, so this also protects the prompt flatten below.
                     backfill_missing_routed_experts(repeated_batch["message_log"])
 
-                    # Extract original prompt messages using the length field
-                    # This correctly handles multi-turn prompts that contain assistant messages
-                    initial_prompt_message_logs = extract_initial_prompt_messages(
-                        repeated_batch["message_log"],
-                        repeated_batch["length"],
-                    )
-
-                    prompt_batched_flat, _ = batched_message_log_to_flat_message(
-                        initial_prompt_message_logs,
-                        pad_value_dict={"token_ids": tokenizer.pad_token_id},
-                    )
-                    prompt_ids_for_adv = prompt_batched_flat["token_ids"]
-                    del initial_prompt_message_logs
-                    del prompt_batched_flat
+                    if "prompt_grouping_ids" in repeated_batch:
+                        prompt_ids_for_adv = repeated_batch["prompt_grouping_ids"]
+                        del repeated_batch["prompt_grouping_ids"]
+                    else:
+                        # Native rollouts group by the tokenized source prompt.
+                        initial_prompt_message_logs = extract_initial_prompt_messages(
+                            repeated_batch["message_log"],
+                            repeated_batch["length"],
+                        )
+                        prompt_batched_flat, _ = batched_message_log_to_flat_message(
+                            initial_prompt_message_logs,
+                            pad_value_dict={"token_ids": tokenizer.pad_token_id},
+                        )
+                        prompt_ids_for_adv = prompt_batched_flat["token_ids"]
+                        del initial_prompt_message_logs
+                        del prompt_batched_flat
 
                     rewards = repeated_batch["total_reward"]
 

nemo_rl/algorithms/single_controller.py

@@ -103,6 +103,7 @@
     squeeze_trailing_unit_dim,
     tensor_field,
 )
+from nemo_rl.algorithms.utils import grouping_ids_from_identifiers
 from nemo_rl.data.interfaces import DatumSpec
 from nemo_rl.data.multimodal_utils import present_multimodal_fields
 from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta
@@ -173,6 +174,25 @@ def _train_fields_for_step(
     )
 
 
+def _prompt_ids_from_group_tags(
+    meta: KVBatchMeta, fallback: torch.Tensor
+) -> torch.Tensor:
+    """Use explicit prompt-group metadata when the producer provides it."""
+    if not meta.tags:
+        return fallback
+
+    group_ids = [tag.get("group_id") for tag in meta.tags]
+    if not any(group_id is not None for group_id in group_ids):
+        return fallback
+    if len(group_ids) != len(meta.sample_ids) or not all(
+        isinstance(group_id, str) and group_id for group_id in group_ids
+    ):
+        raise ValueError(
+            "Prompt group metadata must provide one non-empty group_id for every sample"
+        )
+    return grouping_ids_from_identifiers(cast(list[str], group_ids), 1)
+
+
 @ray.remote(num_cpus=1, num_gpus=0)  # pragma: no cover
 class SingleControllerActor:
     """CPU-only Ray actor that orchestrates the RL training loop.
@@ -3246,6 +3266,8 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
         )
 
         prompt_ids = tensor_field(data, adv_cfg.prompt_ids_field)
+        if not self._is_ppo:
+            prompt_ids = _prompt_ids_from_group_tags(meta, prompt_ids)
         rewards = squeeze_trailing_unit_dim(
             tensor_field(data, adv_cfg.reward_field)
         ).float()

nemo_rl/algorithms/utils.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import hashlib
 import math
 import random
 import warnings
+from collections.abc import Sequence
 from functools import partial, wraps
 from typing import Any, Optional
 
@@ -32,6 +34,29 @@
 from nemo_rl.utils.logger import Logger
 
 
+def grouping_ids_from_identifiers(
+    identifiers: Sequence[str], repeats_per_group: int
+) -> torch.Tensor:
+    """Return stable tensor identities for explicitly formed prompt groups."""
+    if repeats_per_group <= 0:
+        raise ValueError("repeats_per_group must be positive")
+
+    rows = []
+    for identifier in identifiers:
+        digest = hashlib.sha256(identifier.encode()).digest()[:16]
+        rows.append(
+            [
+                int.from_bytes(digest[:8], byteorder="big", signed=True),
+                int.from_bytes(digest[8:], byteorder="big", signed=True),
+            ]
+        )
+    if not rows:
+        return torch.empty((0, 2), dtype=torch.int64)
+    return torch.tensor(rows, dtype=torch.int64).repeat_interleave(
+        repeats_per_group, dim=0
+    )
+
+
 def get_gdpo_reward_component_keys(batch) -> list[str]:
     """Return batch keys that are named reward components (e.g. reward/correctness) in sorted order."""
     return sorted(

nemo_rl/experience/payload.py

@@ -193,7 +193,7 @@ def pack_payload(
 
     Returns:
         sample_ids of the form {group_id}_g{i}, a jagged-packed TensorDict, and per-row
-        tags carrying weight_version plus any per-row violation counts.
+        tags carrying group identity, weight_version, and per-row violation counts.
     """
     lengths = train_batch["input_lengths"]
     n = int(lengths.shape[0])
@@ -212,6 +212,7 @@ def pack_payload(
         {
             "weight_version": weight_version,
             "prompt_idx": prompt_idx,
+            "group_id": group_id,
             **violations[i],
         }
         for i in range(n)

nemo_rl/experience/sync_rollout_actor.py

@@ -217,7 +217,10 @@ def rollout_to_tq(
         # Lazy imports keep rollout-specific dependencies off the actor startup path.
         # ``_policy_dtype`` sizes the VLM pixel tensors below.
         from nemo_rl.algorithms.grpo import _policy_dtype
-        from nemo_rl.algorithms.utils import get_gdpo_reward_component_keys
+        from nemo_rl.algorithms.utils import (
+            get_gdpo_reward_component_keys,
+            grouping_ids_from_identifiers,
+        )
         from nemo_rl.data.llm_message_utils import (
             MESSAGE_LOG_BULK_FIELDS,
             decompose_message_log,
@@ -249,7 +252,8 @@ def rollout_to_tq(
         )
 
         # Rollout dispatch (mirrors grpo_sync.py:294-349).
-        if should_use_nemo_gym(cfg):
+        use_nemo_gym = should_use_nemo_gym(cfg)
+        if use_nemo_gym:
             r = run_nemo_gym_rollout_sync(
                 **common,
                 max_seq_len=None,
@@ -360,6 +364,23 @@ def rollout_to_tq(
                 else np.asarray(v, dtype=object)
             )
 
+        n_samples = int(bulk_batch["sample_mask"].shape[0])
+        input_size = int(input_batch.size)
+        if group_size <= 0 or input_size % group_size != 0:
+            raise ValueError(
+                f"input_batch.size={input_size} is not divisible by group_size={group_size}"
+            )
+        n_prompts = input_size // group_size
+        if n_prompts == 0 or n_samples % n_prompts != 0:
+            raise ValueError(
+                f"bulk_batch has {n_samples} samples; not divisible by n_prompts={n_prompts}"
+            )
+        n_per_prompt = n_samples // n_prompts
+        uids = [str(uuid.uuid4()) for _ in range(n_prompts)]
+        prompt_ids_for_adv = prompt_flat["token_ids"]
+        if use_nemo_gym:
+            prompt_ids_for_adv = grouping_ids_from_identifiers(uids, n_per_prompt)
+
         # Slice — only what the driver can't derive from a TQ slice fetch
         # (anything containing `message_log` or per-token data would
         # force a fetch). Driver does scale_rewards / reward_shaping /
@@ -376,7 +397,7 @@ def rollout_to_tq(
             "truncated": truncated,
             "length": length,
             "input_lengths": input_lengths,
-            "prompt_ids_for_adv": prompt_flat["token_ids"],
+            "prompt_ids_for_adv": prompt_ids_for_adv,
             # Computed by decompose_message_log above; feeds
             # apply_reward_shaping on the driver without a TQ fetch.
             "response_token_lengths": decomposed["response_token_lengths"],
@@ -399,20 +420,16 @@ def rollout_to_tq(
                 )
             driver_carry = {k: driver_carry[k] for k in carry_keys}
 
-        n_samples = int(bulk_batch["sample_mask"].shape[0])
-        input_size = int(input_batch.size)
-        if group_size <= 0 or input_size % group_size != 0:
-            raise ValueError(
-                f"input_batch.size={input_size} is not divisible by group_size={group_size}"
-            )
-        n_prompts = input_size // group_size
-        if n_prompts == 0 or n_samples % n_prompts != 0:
-            raise ValueError(
-                f"bulk_batch has {n_samples} samples; not divisible by n_prompts={n_prompts}"
-            )
-        n_per_prompt = n_samples // n_prompts
-        uids = [str(uuid.uuid4()) for _ in range(n_prompts)]
-        sample_ids = [f"{uid}_g{i}" for uid in uids for i in range(n_per_prompt)]
+        tags = multimodal_row_tags(multimodal, n_samples) or [
+            {} for _ in range(n_samples)
+        ]
+        sample_ids = []
+        row_index = 0
+        for group_id in uids:
+            for generation_index in range(n_per_prompt):
+                sample_ids.append(f"{group_id}_g{generation_index}")
+                tags[row_index]["group_id"] = group_id
+                row_index += 1
         trace_rollout_payload(keys=sample_ids, data=bulk_batch)
         meta = kv_first_write(
             bulk_batch,
@@ -423,7 +440,7 @@ def rollout_to_tq(
             # Per-row shapes the flattening removes from the payload. ``tags``
             # is the transport's per-sample channel and is projected with the
             # rows, so no consumer re-keys them.
-            tags=multimodal_row_tags(multimodal, len(sample_ids)),
+            tags=tags,
             task_name=partition_id,
             pad_to_multiple=int(
                 cfg.policy.get("make_sequence_length_divisible_by") or 1