diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 7df8f4db53a..1d1ae8d9d80 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/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"] diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 525205c4cf9..e307bfd8484 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/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() diff --git a/nemo_rl/algorithms/utils.py b/nemo_rl/algorithms/utils.py index f76133443ea..819f0a2872a 100644 --- a/nemo_rl/algorithms/utils.py +++ b/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( diff --git a/nemo_rl/experience/payload.py b/nemo_rl/experience/payload.py index b0441937942..88f40944182 100644 --- a/nemo_rl/experience/payload.py +++ b/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) diff --git a/nemo_rl/experience/sync_rollout_actor.py b/nemo_rl/experience/sync_rollout_actor.py index 48bf096d530..b25c4f7364b 100644 --- a/nemo_rl/experience/sync_rollout_actor.py +++ b/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 diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 901f986b65f..21a4b237203 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -46,7 +46,9 @@ _initial_policy_generation_stale, _maybe_restore_async_replay_buffer_checkpoint, _needs_hf_refit_handshake, + _prompt_grouping_ids, _raise_if_reward_penalties_enabled_without_nemo_gym, + _replay_prompt_batches, _resolve_logprob_skip_flags, _resolve_message_level_advantage_penalties, _save_async_replay_buffer_checkpoint, @@ -263,6 +265,68 @@ def test_missing_mask_sample_is_noop(self): ) +def test_nemo_gym_prompt_grouping_uses_source_group_positions() -> None: + grouping_ids = _prompt_grouping_ids( + torch.empty((4, 0)), + num_prompts=2, + num_generations_per_prompt=2, + use_nemo_gym=True, + ) + + assert torch.equal(grouping_ids, torch.tensor([[0], [0], [1], [1]])) + + +def test_nemo_gym_prompt_grouping_offsets_dynamic_sampling_batches() -> None: + grouping_ids = _prompt_grouping_ids( + torch.empty((4, 0)), + num_prompts=2, + num_generations_per_prompt=2, + use_nemo_gym=True, + group_offset=2, + ) + + assert torch.equal(grouping_ids, torch.tensor([[2], [2], [3], [3]])) + + +def test_native_prompt_grouping_preserves_token_identity() -> None: + prompt_token_ids = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4]]) + + grouping_ids = _prompt_grouping_ids( + prompt_token_ids, + num_prompts=2, + num_generations_per_prompt=2, + use_nemo_gym=False, + ) + + assert grouping_ids is prompt_token_ids + + +def test_replay_prompt_batches_restore_grouping_without_task_indices() -> None: + def make_batch(prompt: str) -> BatchedDataDict: + return create_mock_batch( + 2, + ["math", "math"], + [ + [ + {"role": "user", "content": f"{prompt}_{index}"}, + {"role": "assistant", "content": "response"}, + ] + for index in range(2) + ], + ) + + batches = _replay_prompt_batches( + [ + {"batch": make_batch("first")}, + {"batch": make_batch("second")}, + ], + use_explicit_grouping=True, + ) + + assert torch.equal(batches[0]["prompt_grouping_ids"], torch.tensor([[0], [0]])) + assert torch.equal(batches[1]["prompt_grouping_ids"], torch.tensor([[1], [1]])) + + def test_initial_policy_generation_stale() -> None: generation = MagicMock() generation.weight_synchronizer.is_stale = False @@ -2537,6 +2601,58 @@ def test_dapo_dynamic_sampling_batch_caching(mock_grpo_components): assert batch_cache is not None +def test_dapo_dynamic_sampling_carries_prompt_grouping_ids(mock_grpo_components): + def make_batch(grouping_id: int) -> BatchedDataDict: + batch = create_mock_batch( + 3, + ["math"] * 3, + [ + [ + {"role": "user", "content": f"prompt_{grouping_id}"}, + {"role": "assistant", "content": f"response_{i}"}, + ] + for i in range(3) + ], + ) + batch["total_reward"] = torch.tensor([1.0, 0.0, 0.5]) + batch["prompt_grouping_ids"] = torch.full((3, 1), grouping_id) + return batch + + master_config = mock_grpo_components["master_config"] + master_config.grpo.use_dynamic_sampling = True + master_config.grpo.num_prompts_per_step = 2 + master_config.grpo.num_generations_per_prompt = 3 + master_config.grpo.dynamic_sampling_max_gen_batches = 5 + std = torch.tensor([0.4, 0.4, 0.4]) + baseline = torch.tensor([0.5, 0.5, 0.5]) + + _, complete, cache, _ = dynamic_sampling( + make_batch(10), + std, + baseline, + dynamic_sampling_num_gen_batches=1, + master_config=master_config, + timer=Timer(), + ) + assert not complete + + result, complete, _, _ = dynamic_sampling( + make_batch(11), + std, + baseline, + dynamic_sampling_num_gen_batches=2, + master_config=master_config, + timer=Timer(), + batch_cache=cache, + ) + + assert complete + assert torch.equal( + result["prompt_grouping_ids"], + torch.tensor([[10], [10], [10], [11], [11], [11]]), + ) + + def test_dapo_cache_aligns_deduplicated_media_with_text_only_batch( mock_grpo_components, ): @@ -4059,6 +4175,31 @@ def _run_single_grpo_train_step(mock_grpo_components, train_func, monkeypatch): ) +def test_grpo_train_passes_source_prompt_grouping_ids_to_estimator( + mock_grpo_components, monkeypatch +): + source_prompt_grouping_ids = torch.tensor([[1234]]) + mock_adv_estimator = MagicMock() + mock_adv_estimator.compute_advantage.side_effect = ( + lambda **kwargs: torch.zeros_like(kwargs["mask"]) + ) + monkeypatch.setattr( + "nemo_rl.algorithms.grpo._create_advantage_estimator", + lambda _cfg: mock_adv_estimator, + ) + monkeypatch.setattr( + "nemo_rl.algorithms.grpo._prompt_grouping_ids", + lambda *args, **kwargs: source_prompt_grouping_ids, + ) + + _run_single_grpo_train_step(mock_grpo_components, grpo_train, monkeypatch) + + assert torch.equal( + mock_adv_estimator.compute_advantage.call_args.kwargs["prompt_ids"], + source_prompt_grouping_ids, + ) + + @pytest.mark.parametrize("train_func", [grpo_train, async_grpo_train]) def test_grpo_train_clips_advantages_when_configured( mock_grpo_components, train_func, monkeypatch diff --git a/tests/unit/algorithms/test_utils.py b/tests/unit/algorithms/test_utils.py index 559d2f5c13a..f9d607e864b 100755 --- a/tests/unit/algorithms/test_utils.py +++ b/tests/unit/algorithms/test_utils.py @@ -28,6 +28,7 @@ WALL_CLOCK_EFFICIENCY_CATEGORIES, calculate_baseline_and_std_per_prompt, get_tokenizer, + grouping_ids_from_identifiers, maybe_pad_last_batch, print_efficiency_summary, print_performance_metrics, @@ -36,6 +37,19 @@ from nemo_rl.distributed.batched_data_dict import BatchedDataDict +def test_grouping_ids_from_identifiers_are_stable_and_distinct() -> None: + grouping_ids = grouping_ids_from_identifiers(["group-a", "group-b"], 2) + + assert grouping_ids.shape == (4, 2) + assert torch.equal(grouping_ids[0], grouping_ids[1]) + assert torch.equal(grouping_ids[2], grouping_ids[3]) + assert not torch.equal(grouping_ids[0], grouping_ids[2]) + assert torch.equal( + grouping_ids, + grouping_ids_from_identifiers(["group-a", "group-b"], 2), + ) + + @pytest.fixture def conversation_messages(): """Fixture providing a multi-turn conversation for testing chat templates""" diff --git a/tests/unit/data_plane/test_smoke.py b/tests/unit/data_plane/test_smoke.py index e373adf076f..3574df6b7fc 100644 --- a/tests/unit/data_plane/test_smoke.py +++ b/tests/unit/data_plane/test_smoke.py @@ -164,3 +164,111 @@ def test_sync_rollout_actor_prompt_extraction_and_masks_match_grpo() -> None: flat["generation_logprobs"], torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.1, 0.2]]), ) + + +def test_sync_rollout_actor_uses_explicit_nemo_gym_prompt_groups() -> None: + """Harness prompt rewrites must not change synchronous TQ grouping.""" + from types import SimpleNamespace + from unittest.mock import patch + + import torch + + from nemo_rl.data_plane.interfaces import KVBatchMeta + from nemo_rl.distributed.batched_data_dict import BatchedDataDict + from nemo_rl.experience.sync_rollout_actor import SyncRolloutActor + + message_logs = [ + [ + { + "role": "user", + "content": f"workspace-{index}", + "token_ids": torch.tensor([100 + index]), + }, + { + "role": "assistant", + "content": "answer", + "token_ids": torch.tensor([200 + index]), + "generation_logprobs": torch.tensor([-0.1]), + }, + ] + for index in range(4) + ] + final_batch = BatchedDataDict( + { + "message_log": message_logs, + "length": torch.ones(4, dtype=torch.long), + "total_reward": torch.tensor([1.0, 0.0, 0.5, 0.25]), + "loss_multiplier": torch.ones(4), + "truncated": torch.zeros(4, dtype=torch.bool), + } + ) + rollout_result = SimpleNamespace( + final_batch=final_batch, + rollout_metrics={"mean_gen_tokens_per_sample": 1.0}, + ) + + controller_cls = SyncRolloutActor.__ray_metadata__.modified_class + actor = object.__new__(controller_cls) + actor.policy_generation = None + actor.tokenizer = SimpleNamespace(pad_token_id=0) + actor.task_to_env = {} + actor._dp_client = object() + actor.master_config = SimpleNamespace( + policy={ + "generation": {}, + "make_sequence_length_divisible_by": 1, + "precision": "float32", + }, + logger={"wandb_enabled": False, "wandb": {}}, + env={"nemo_gym": {}}, + grpo=SimpleNamespace( + deduplicate_multimodal_data=False, + debug_payload_metrics=False, + ), + reward_penalties=SimpleNamespace(), + ) + + def fake_write(batch, *, sample_ids, partition_id, tags, **kwargs): + del batch, kwargs + return KVBatchMeta( + partition_id=partition_id, + task_name=partition_id, + sample_ids=sample_ids, + tags=tags, + ) + + with ( + patch( + "nemo_rl.experience.sync_rollout_actor.run_nemo_gym_rollout_sync", + return_value=rollout_result, + ), + patch( + "nemo_rl.experience.sync_rollout_actor.get_nemo_gym_thinking_tags", + return_value=None, + ), + patch( + "nemo_rl.environments.nemo_gym.should_use_nemo_gym", + return_value=True, + ), + patch( + "nemo_rl.experience.sync_rollout_actor.kv_first_write", + side_effect=fake_write, + ), + ): + meta, carry, _, _ = actor.rollout_to_tq( + BatchedDataDict({"row": torch.arange(4)}), + partition_id="train", + group_size=2, + ) + + grouping_ids = carry["prompt_ids_for_adv"] + assert torch.equal(grouping_ids[0], grouping_ids[1]) + assert torch.equal(grouping_ids[2], grouping_ids[3]) + assert not torch.equal(grouping_ids[0], grouping_ids[2]) + assert meta.tags is not None + assert [tag["group_id"] for tag in meta.tags] == [ + meta.sample_ids[0].rsplit("_g", 1)[0], + meta.sample_ids[1].rsplit("_g", 1)[0], + meta.sample_ids[2].rsplit("_g", 1)[0], + meta.sample_ids[3].rsplit("_g", 1)[0], + ] diff --git a/tests/unit/experience/test_payload.py b/tests/unit/experience/test_payload.py index a30e8f86c9e..b46a03c0f5b 100644 --- a/tests/unit/experience/test_payload.py +++ b/tests/unit/experience/test_payload.py @@ -151,8 +151,8 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None: "num_assistant_messages": 1, } assert tags == [ - {"weight_version": 3, "prompt_idx": 17, **no_violations}, - {"weight_version": 3, "prompt_idx": 17, **no_violations}, + {"weight_version": 3, "prompt_idx": 17, "group_id": "group", **no_violations}, + {"weight_version": 3, "prompt_idx": 17, "group_id": "group", **no_violations}, ] @@ -359,6 +359,7 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None: { "weight_version": 7, "prompt_idx": 17, + "group_id": "g", "num_invalid_tool_calls": 1, "num_malformed_thinking": 0, "num_assistant_messages": 1, @@ -366,6 +367,7 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None: { "weight_version": 7, "prompt_idx": 17, + "group_id": "g", "num_invalid_tool_calls": 0, "num_malformed_thinking": 1, "num_assistant_messages": 1, @@ -373,6 +375,7 @@ def test_pack_payload_stamps_violation_counts_on_tags() -> None: { "weight_version": 7, "prompt_idx": 17, + "group_id": "g", "num_invalid_tool_calls": 0, "num_malformed_thinking": 0, "num_assistant_messages": 0, diff --git a/tests/unit/single_controller/test_single_controller_actor.py b/tests/unit/single_controller/test_single_controller_actor.py index 50e361f5a1d..6b69381c0e2 100644 --- a/tests/unit/single_controller/test_single_controller_actor.py +++ b/tests/unit/single_controller/test_single_controller_actor.py @@ -34,6 +34,7 @@ from nemo_rl.algorithms.single_controller import ( SingleControllerActor, _pooled_opd_metrics, + _prompt_ids_from_group_tags, ) from nemo_rl.algorithms.single_controller_utils.config import ( AdvantageConfig, @@ -46,6 +47,37 @@ from nemo_rl.utils.timer import TimeoutChecker, Timer +def test_prompt_group_tags_are_explicit_and_complete() -> None: + fallback = torch.tensor([[10], [20]]) + grouped_meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["opaque-a", "opaque-b"], + tags=[{"group_id": "same"}, {"group_id": "same"}], + ) + + grouped = _prompt_ids_from_group_tags(grouped_meta, fallback) + + assert torch.equal(grouped[0], grouped[1]) + custom_meta = KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["custom_g0", "custom_g1"], + tags=[{}, {}], + ) + assert _prompt_ids_from_group_tags(custom_meta, fallback) is fallback + with pytest.raises(ValueError, match="one non-empty group_id"): + _prompt_ids_from_group_tags( + KVBatchMeta( + partition_id="train", + task_name="train", + sample_ids=["opaque-a", "opaque-b"], + tags=[{"group_id": "same"}, {}], + ), + fallback, + ) + + class FakeWeightSynchronizer: pass @@ -542,10 +574,12 @@ def put_samples(self, *, fields, **kwargs) -> None: class _MaskRecordingAdvantageEstimator: def __init__(self) -> None: self.mask: torch.Tensor | None = None + self.prompt_ids: torch.Tensor | None = None - def compute_advantage(self, *, rewards, mask, **kwargs) -> torch.Tensor: + def compute_advantage(self, *, prompt_ids, rewards, mask, **kwargs) -> torch.Tensor: del kwargs self.mask = mask.clone() + self.prompt_ids = prompt_ids.clone() return rewards.unsqueeze(-1).expand_as(mask).clone() @@ -610,8 +644,14 @@ def test_advantage_stage_composes_all_filters_before_computing_advantages( meta = KVBatchMeta( partition_id="rollout_data", task_name="train", - sample_ids=[f"sample-{i}" for i in range(batch_size)], + sample_ids=["first_g0", "first_g1", "second_g0", "second_g1"], fields=list(data.keys()), + tags=[ + {"group_id": "first"}, + {"group_id": "first"}, + {"group_id": "second"}, + {"group_id": "second"}, + ], ) result_meta, has_valid_training_tokens = asyncio.run(ctrl._advantage_stage(meta)) @@ -638,6 +678,10 @@ def test_advantage_stage_composes_all_filters_before_computing_advantages( assert estimator.mask is not None assert estimator.mask[0].all() assert estimator.mask[1:].count_nonzero() == 0 + assert estimator.prompt_ids is not None + assert torch.equal(estimator.prompt_ids[0], estimator.prompt_ids[1]) + assert torch.equal(estimator.prompt_ids[2], estimator.prompt_ids[3]) + assert not torch.equal(estimator.prompt_ids[0], estimator.prompt_ids[2]) assert ctrl._step_log_dict["num_mask_sample_filtered"] == [1] metrics = ctrl._step_log_dict["seq_logprob_error_metrics"] assert len(metrics) == 1 @@ -2202,11 +2246,10 @@ def test_advantage_stage_writes_gae_returns_alongside_advantages() -> None: """The critic's regression target has to reach TQ, or the value train step fetches a column nobody wrote.""" batch_size, sequence_length = 2, 4 + prompt_ids_for_adv = torch.tensor([[1, 2, 0, 0], [3, 4, 0, 0]], dtype=torch.long) data = TensorDict( { - "prompt_ids_for_adv": torch.zeros( - batch_size, sequence_length, dtype=torch.long - ), + "prompt_ids_for_adv": prompt_ids_for_adv, "total_reward": torch.tensor([1.0, 0.0]), "token_mask": torch.ones(batch_size, sequence_length), "sample_mask": torch.ones(batch_size), @@ -2252,8 +2295,9 @@ def compute_advantage(self, *, rewards, mask, **kwargs): meta = KVBatchMeta( partition_id="rollout_data", task_name="train", - sample_ids=[f"sample-{i}" for i in range(batch_size)], + sample_ids=["same_g0", "same_g1"], fields=list(data.keys()), + tags=[{"group_id": "same"}, {"group_id": "same"}], ) result_meta, has_valid_training_tokens = asyncio.run(ctrl._advantage_stage(meta)) @@ -2261,6 +2305,7 @@ def compute_advantage(self, *, rewards, mask, **kwargs): assert has_valid_training_tokens assert "values" in (data_plane.selected_fields or []) assert estimator.kwargs is not None + assert torch.equal(estimator.kwargs["prompt_ids"], prompt_ids_for_adv) assert torch.equal(estimator.kwargs["values"], torch.zeros(2, 4)) assert data_plane.written_fields is not None assert torch.equal( diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 4aa323b9d45..f426d081b28 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -504,7 +504,11 @@ def test_commit_writes_tq_then_fills_meta(self, monkeypatch): assert buf.ready_list == [True] assert buf.meta_list[0].sample_ids == meta.sample_ids # TQ tags preserve both dispatch-time weight and dataset identity. - assert meta.tags == [{"weight_version": 3, "prompt_idx": 418}] * _N_GENS + assert ( + meta.tags + == [{"weight_version": 3, "prompt_idx": 418, "group_id": group_id}] + * _N_GENS + ) assert len(dp.put_calls) == 1 assert len(trace_calls) == 1 assert trace_calls[0]["keys"] == meta.sample_ids