fix(grpo): group advantages by rollout identity

[aroshanghias-nvd]

What does this PR do ?

Fix GRPO credit assignment by grouping completions by their rollout sampling identity instead of tokenized prompt equality.

Problem

GRPO baselines are defined over the completions sampled together for one prompt instance. NeMo RL previously used the tokenized initial prompt as that instance's group identifier. Independent rollout groups were therefore merged whenever their initial token rows were identical.

This is especially likely for multimodal and stateful environments. Different images can share the same text and image-placeholder token sequence, while text-only environment instances can share initial instructions but differ in reward-relevant state such as their seed, tool state, or later observations.

For example, consider two independent 16-generation groups with identical initial token rows:

Group A: image/state A, rewards [0, 0, ..., 0]
Group B: image/state B, rewards [1, 1, ..., 1]

The correct within-group RLOO advantages are zero for both groups. Token-based grouping instead forms one 32-sample group. A zero-reward sample receives a baseline of 16/31, and a one-reward sample receives a baseline of 15/31, creating false negative and positive advantages.

The issue is often not visible in normal training because most text datasets have unique prompt tokens. Even when collisions occur, pooling has no numerical effect if the colliding groups have the same reward distribution. A limited amount of corrupted credit assignment can also look like ordinary optimization noise rather than an obvious failure.

Fix

• Assign explicit sampling-group IDs in legacy synchronous and asynchronous GRPO.
• Recover the existing group UUID from {group_id}_g{generation_index} sample IDs in TransferQueue and Single Controller paths.
• Preserve group identity through filtering, caching, concatenation, and slicing.
• Rebuild TransferQueue group IDs after dynamic-sampling caches from multiple generation batches are combined.
• Validate malformed, duplicate, and incomplete unfiltered groups while permitting legitimate partial groups after dynamic sampling.

This does not change the sample-ID or checkpoint schema; it consumes group identity already present in rollout metadata.

Issues

None.

Usage

No user-facing configuration or API changes are required.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added regression tests for the affected execution paths.
• Ran 17 targeted unit tests in the project container.
• No documentation update is required because this is an internal correctness fix with no user-facing behavior or configuration change.

Additional Information

The regression coverage includes:

• identical token prompts belonging to different reward groups;
• rollout-group ID construction and validation;
• Single Controller advantage computation;
• legacy dynamic-sampling cache concatenation;
• TransferQueue caching across two generation batches whose temporary dense IDs both begin at zero;
• expected RLOO baselines after cached groups are separated correctly.

This fixes incorrect GRPO credit assignment. It does not claim to resolve entropy collapse or other policy-degeneration behavior generally.

Full repository CI remains the merge gate.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[aroshanghias-nvd]

Self-review of my own PR, run with this repo's /review-pr-team skill (rl-expert, bug-finder, test-agent, design-reviewer, comment-reviewer, devil-advocate). Posting the surviving findings for transparency. The adversarial pass cut or downgraded 12 of 15 clusters and disputed 1 outright, including two of my own — what's left is what survived it.

2 action items, 3 follow-ups.

AI-1 — ruff format is not clean on 4 of the 8 changed files

Every offending hunk is a line this PR adds:

File	Lines
nemo_rl/algorithms/utils.py	126-128
nemo_rl/algorithms/grpo_sync.py	237-239
tests/unit/algorithms/test_grpo.py	2234-2236, 2279-2281
tests/unit/algorithms/test_utils.py	46-48, 79-81, 92-94

Reproduced three times independently, twice at the v0.9.9 pinned in .pre-commit-config.yaml. ruff check passes; only format fails. Action: uv run ruff format .. Inline suggestions cover the first hunk in each file.

This is invisible right now because Request NVSkills CI is skipped, not passing — the copy-pr-bot gate means the lint job has never run on this branch.

AI-2 — this PR orphans a full-batch prompt flatten in grpo_train

input_ids was the GRPO group key. Both consumers are now reward_group_ids, so grpo.py:2995-3000 builds an unused [B, S_prompt] tensor every generation batch. Details and the required deletion set inline — note it is not a single-block delete.

Follow-up 1 — attach CI:L2

labels is currently empty. The cicd skill maps "changes that could affect convergence" to CI:L2, which this is by definition — it changes the GRPO baseline denominator.

Follow-up 2 — the "usually not visible" framing in the description is too narrow

The description argues collisions mainly matter for multimodal and stateful environments. That understates it: any text-only dataset containing duplicate prompt strings previously pooled those rows into one baseline group and will now get separate per-instance baselines. Advantages change on those runs too, so this is not a text-recipe no-op.

Worth one cheap number in the description instead of a full A/B: for one step on one text and one VLM recipe, report len(torch.unique(prompt_ids_for_adv, dim=0)) against num_prompts_per_step. Equal for the first, strictly less for the second, is the whole empirical case for the PR.

Follow-up 3 — prompt_ids_for_adv is now a write-only column

Removing the last readers (_advantage_input_fields, and the overwrite at grpo_sync.py:702) leaves payload.py:91, sync_rollout_actor.py:363 and data_plane/schema.py:54 still producing and shipping the column, and AdvantageConfig.prompt_ids_field with zero readers repo-wide. Retaining the write is defensible for a correctness-only PR that promises no schema change — but a follow-up to drop the column, or a comment saying why it stays, would save the next reader from re-deriving that it is dead. Not asking for it here.

---

Context — no action.

The core fix was checked rather than assumed. Invariants confirmed independently by more than one agent:

• Row order is preserved on all three legacy rollout branches, which positional grouping now depends on and token grouping did not: the sync path re-sorts by original index (rollouts.py:836-860), the async path relies on asyncio.gather ordering (:1573), and the NeMo-Gym path returns in input-row order (:2587).
• repeated_batch is batch.repeat_interleave(G), so it is group-contiguous and divisible at every build_rollout_group_ids call site.
• rollout_group_ids survives select_indices / from_batches / slice, and next_rollout_group_id resets in lockstep with batch_cache, so cached generation batches cannot collide.
• Multimodal dedup does not corrupt repeated_batch.size — PackedTensor.__len__ returns logical rows.
• [B,1] is correct at every consumer, including the calculate_advantages_on_gpu branch where group IDs, rewards and mask all move to the same device together.
• SC cannot see a partial group: dynamic sampling and overlong filtering are rejected outright at single_controller_utils/config.py:739.
• Sample-ID producers all mint f"{uuid4()}_g{i}", and uuid4 contains no underscore, so rsplit("_g", 1) is unambiguous.

test_apply_dynamic_sampling_rebuilds_group_ids_across_generation_batches is the one test here that genuinely fails on the old behaviour.

Two design notes worth recording: carrying identity as a batch key in legacy grpo.py rides select_indices/concat for free — the TQ path has to rebuild IDs at grpo_sync.py:237 precisely because it lacks that — and keeping both new helpers as pure functions is what lets the bug be pinned in six lines with no fixtures.

Merge is currently blocked by the copy-pr-bot runner gate, which needs a vetter and is not self-serviceable. All checks that did run (DCO, copyright, secrets-detector, submodule fast-forward, semantic title) pass.

Findings dropped after adversarial review: a sample-ID-grammar refactor (its motivating scenario turned out to be invented — the redispatch path reuses the same group_id), two suggested tests (one cosmetic, one whose pytest.raises(match=...) would have failed on arrival — [0, 1] is a regex character class), a docstring-completeness ask (D417 is in this repo's ruff ignore list), and an exception-type nit (the friendly error already fires, via the unpack ValueError).

[yfw]

single_controller.py:2119

1 action item. (PR-caused, but the edits land in adjacent files.)

TL;DR — dropping prompt_ids_field here removes the last reader of prompt_ids_for_adv, leaving a column that every producer still writes and nobody ever reads.

Three pieces, all downstream of this one deletion:

1. The SC data-plane column. payload.py:91 still jagged-packs prompt_ids_for_adv into every rollout group's TQ put, and it is still registered in SC_ROLLOUT_SCHEMA_FIELDS at schema.py:54. With this line gone, _advantage_input_fields() never fetches it — a fully-written, never-read per-row column.
2. The sync actor's flatten. sync_rollout_actor.py:93 runs extract_initial_prompt_messages + batched_message_log_to_flat_message every generation batch, ships the result across the Ray boundary at sync_rollout_actor.py:363, and grpo_sync.py:702 overwrites it before any read. Nothing touches driver_carry in between, it is not in bulk_batch, and validation (carry_keys=["total_reward", ...]) pays the flatten and discards it too.
3. The knob. AdvantageConfig.prompt_ids_field now has exactly one reference repo-wide — its own definition at config.py:1065. It went 2 readers to 0 in this PR. Setting it now silently does nothing.

Your Follow-up 3 covers (1) and reasonably defers it as a schema change. The part that isn't covered is (2): driver_carry is a Ray return value, not a data-plane column, so it carries no schema commitment — it is the same orphan you already deleted from grpo.py in fc1a3c193, just on the TQ actor. Note it is not a single-block delete: you'd also narrow the return type at sync_rollout_actor.py:76, drop the now-unused prompt_lengths param and the extract_initial_prompt_messages local import, and update the :210 docstring.

Action: delete the prompt_ids_for_adv production in _flatten_rollout_message_log_for_tq / rollout_to_tq, and delete AdvantageConfig.prompt_ids_field (a plain internal @dataclass, constructed no-args, no YAML or checkpoint path — deletion is safe). Leave the schema column for the follow-up, but consider a one-line comment at config.py:1065 so the next reader doesn't re-derive that it's dead.

[yfw]

utils.py:183

1 action item. (PR-introduced — this PR is what makes the line below wrong.)

The docstring of the parameter you're redefining still says the opposite of this PR's thesis. At utils.py:194:

    prompts:    tensor (b, s)     Tensor of prompts the model used. May be on any device

That's the one line in the tree still asserting prompt token rows are the group key — eight lines below the new helper whose whole purpose is to replace them. Highest-probability spot for the next reader to reintroduce the bug.

Action: reword line 194 to describe a group key rather than prompts:

    prompts:    tensor (b, k)     Per-row group key; rows with an identical key row
                                  share a baseline. Production passes a (b, 1)
                                  rollout-group-id column (see build_rollout_group_ids).
                                  May be on any device

No suggestion block because line 194 sits just outside this hunk (it ends at 184).

Deliberately not asking you to rename the prompts parameter — it's a misnomer now, but all six call sites pass it positionally and it reaches untouched PPO files. And the wording above stays generic on purpose: this helper legitimately accepts any (b, k), and 8 existing tests in test_utils.py pass (6, 3).

[yfw]

single_controller.py:1977

1 action item — low severity. (PR-introduced.)

This runs unconditionally, before the _is_ppo branch, so it fires on paths that have no use for the value:

• PPO runs reach it, but their only reachable estimators are gae / raw_reward (ppo.py:1135 restricts the factory; GAEConfig.name is Literal["gae", "raw_reward"]), and neither reads prompt_ids at all — we checked by AST rather than docstring, since **kwargs hides unused params. Zero load references in either body.
• It also runs when has_valid_training_tokens is False and compute_advantage is never called.

So on those paths this is both an unconditional cost and, more importantly, an unconditional new ValueError — the strict expected_group_size check can now fail a run that previously trained fine, for a value that would have been discarded.

To be clear this is a robustness nit, not a live bug: we could not construct a config where a real SC run produces incomplete groups (the sole SC producer mints f"{group_id}_g{i}" at payload.py:126, each PromptGroupRecord holds exactly num_generations_per_prompt completions, and config.py:736 rejects dynamic sampling and overlong filtering outright for both PPO and GRPO).

Action: move the build_rollout_group_ids_from_sample_ids(...) call so it only runs when the estimator actually consumes prompt_ids — or at minimum below the early-return for has_valid_training_tokens.