Skip to content

fix(genrm): propagate prompt cohort retry identity - #3861

Open
macandro96 wants to merge 2 commits into
mainfrom
amahishi/genrm-fixes
Open

fix(genrm): propagate prompt cohort retry identity#3861
macandro96 wants to merge 2 commits into
mainfrom
amahishi/genrm-fixes

Conversation

@macandro96

@macandro96 macandro96 commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Summary

Propagate stable logical prompt-group identity from NeMo-RL to Gym and GenRM so retries of the same
prompt group can be distinguished from new logical groups.

This is the NeMo-RL side of the corresponding Gym GenRM cohort fix:

Identity model

Each Gym rollout now carries:

  • _ng_group_id: stable across retries of one logical prompt group.
  • _ng_group_attempt: starts at 0 and increments when the entire group is retried.
  • _ng_rollout_index: stable logical sibling index within the group.

The Transfer Queue identity remains separate:

  • A new TQ group ID is allocated for every physical dispatch attempt.
  • The logical Gym group_id remains stable across those attempts.
  • This prevents stale data from an old TQ transaction from colliding with the replacement transaction.

Conceptually:

Logical prompt group
    group_id = G
        |
        +-- attempt 0
        |      group_attempt = 0
        |      fresh TQ group ID = TQ-A
        |      rollout_index = 0..N-1
        |
        +-- retry
               group_attempt = 1
               fresh TQ group ID = TQ-B
               rollout_index = 0..N-1

Changes

  • Define constants for the three logical cohort coordinates.
  • Stamp identity in both NeMo-Gym row-construction paths.
  • Assign one group_id to all generations belonging to the same prompt.
  • Assign stable sibling indices from 0 through num_generations - 1.
  • Preserve the logical group_id across infrastructure and data retries.
  • Increment group_attempt before each replacement group dispatch.
  • Continue allocating a fresh TQ group ID for every physical attempt.
  • Preserve explicitly supplied group identity when valid.
  • Reject invalid, negative, boolean, or inconsistent group-attempt values.
  • Reject batches that do not contain complete prompt groups.
  • Pin the Gym submodule to the matching GenRM cohort implementation.

Why this is needed

Without stable logical identity, GenRM cannot distinguish:

  1. A new occurrence of the same prompt.
  2. A retry of an existing prompt group.
  3. A duplicate physical request for an existing sibling.

The combination of group_id, group_attempt, and rollout_index gives GenRM enough information to isolate
replacement cohorts and deduplicate logical sibling slots safely.

Compatibility

A new group_id is generated when callers do not provide one, and the initial group_attempt defaults to 0.

The change does not reuse TQ transaction identities. Existing cleanup and rollback behavior continues to
operate on fresh physical TQ group IDs.

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 requested review from a team as code owners August 27, 2026 00:06
@copy-pr-bot

copy-pr-bot Bot commented Aug 27, 2026

Copy link
Copy Markdown

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.

@github-actions

Copy link
Copy Markdown

✅ Submodule Fast-Forward Check Results

Check based on commit: 0c22e3e (PR #3861 from amahishi/genrm-fixes)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

@ananthsub

ananthsub commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

the synchronous Gym path uses the total batch size as the size of one prompt cohort.

run_nemo_gym_rollout_sync() then calls run_async_nemo_gym_rollout() with num_generations=input_batch.size

_prepare_nemo_gym_rows() therefore gives all rows one _ng_group_id and assigns _ng_rollout_index from 0 through N-1.

Those rows do not belong to one GenRM comparison. Gym can reject prompt B because its text differs from the prompt already stored for that group. Larger batches can also produce rollout indices outside GenRM's configured cohort size.

This helper is used by synchronous GRPO training, GRPO validation, PPO, distillation, and SyncRolloutActor. Any of those paths can hit this when a batch contains more than one prompt.

The fix needs to keep two values separate:

  • num_generations_per_prompt defines which rows share _ng_group_id; _ng_rollout_index must restart at zero for each prompt.
  • input_batch.size defines how many total rows the synchronous call waits for and returns.

this fix should pass the per-prompt generation count into run_nemo_gym_rollout_sync() and use it only for identity grouping. it can continue collecting every row in the full batch, then return all prompt groups in their original input order.

replacing the current num_generations value is not enough because that value also controls whole-batch result collection today

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96

Copy link
Copy Markdown
Contributor Author

/ok to test 492afc5

@macandro96 macandro96 added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants