fix(grpo): group advantages by rollout identity #3882
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -76,6 +76,7 @@ | |
| squeeze_trailing_unit_dim, | ||
| tensor_field, | ||
| ) | ||
| from nemo_rl.algorithms.utils import build_rollout_group_ids_from_sample_ids | ||
| from nemo_rl.data.interfaces import DatumSpec | ||
| from nemo_rl.data_plane import KVBatchMeta | ||
| from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS | ||
|
|
@@ -1971,7 +1972,12 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: | |
| select_fields=self._advantage_input_fields(), | ||
| ) | ||
|
|
||
| # The selected metadata retains each prompt group's UUID in its sample | ||
| # IDs. Text-token equality is insufficient for multimodal prompts. | ||
| prompt_ids = build_rollout_group_ids_from_sample_ids( | ||
| meta.sample_ids, | ||
| expected_group_size=self._algo_cfg.num_generations_per_prompt, | ||
| ) | ||
| rewards = squeeze_trailing_unit_dim( | ||
| tensor_field(data, adv_cfg.reward_field) | ||
| ).float() | ||
|
|
@@ -2110,7 +2116,6 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]: | |
| def _advantage_input_fields(self) -> list[str]: | ||
| adv_cfg = self._advantage_cfg | ||
| fields = [ | ||
| adv_cfg.reward_field, | ||
|
Contributor
There was a problem hiding this comment. 1 action item. (PR-caused, but the edits land in adjacent files.) TL;DR — dropping Three pieces, all downstream of this one deletion:
Your Follow-up 3 covers (1) and reasonably defers it as a schema change. The part that isn't covered is (2): Action: delete the |
||
| adv_cfg.token_mask_field, | ||
| adv_cfg.sample_mask_field, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -101,6 +101,84 @@ def calculate_kl( | |
| return kl | ||
|
|
||
|
|
||
| def build_rollout_group_ids( | ||
| batch_size: int, | ||
| group_size: int, | ||
| *, | ||
| start_group_id: int = 0, | ||
| device: torch.device | str | None = None, | ||
| ) -> torch.Tensor: | ||
| """Return an explicit identifier for every rollout's sampling group. | ||
|
|
||
| GRPO baselines are defined over the completions sampled together for one | ||
| prompt instance. Prompt token IDs are not a valid group identifier for | ||
| multimodal prompts because different media can have identical text and | ||
| therefore identical token rows. | ||
| """ | ||
| if group_size <= 0: | ||
| raise ValueError(f"group_size must be positive, got {group_size}") | ||
| if batch_size < 0 or batch_size % group_size != 0: | ||
| raise ValueError( | ||
| f"batch_size={batch_size} must be non-negative and divisible by " | ||
| f"group_size={group_size}" | ||
| ) | ||
| if start_group_id < 0: | ||
| raise ValueError(f"start_group_id must be non-negative, got {start_group_id}") | ||
| num_groups = batch_size // group_size | ||
| return torch.arange( | ||
| start_group_id, | ||
| start_group_id + num_groups, | ||
| device=device, | ||
| dtype=torch.long, | ||
| ).repeat_interleave(group_size)[:, None] | ||
|
|
||
|
|
||
| def build_rollout_group_ids_from_sample_ids( | ||
| sample_ids: list[str], | ||
| *, | ||
| expected_group_size: int | None = None, | ||
| device: torch.device | str | None = None, | ||
| ) -> torch.Tensor: | ||
| """Derive explicit prompt-group IDs from ``{group_id}_g{index}`` sample IDs.""" | ||
| if expected_group_size is not None and expected_group_size <= 0: | ||
| raise ValueError( | ||
| f"expected_group_size must be positive, got {expected_group_size}" | ||
| ) | ||
| group_to_index: dict[str, int] = {} | ||
| group_to_generation_indices: dict[str, set[int]] = {} | ||
| group_indices: list[int] = [] | ||
| for sample_id in sample_ids: | ||
| try: | ||
| group_id, generation_index = sample_id.rsplit("_g", 1) | ||
| generation_index = int(generation_index) | ||
| except (ValueError, TypeError) as error: | ||
| raise ValueError( | ||
| "Expected sample ID in '{group_id}_g{generation_index}' format, " | ||
| f"got {sample_id!r}" | ||
| ) from error | ||
| if not group_id: | ||
| raise ValueError(f"Sample ID has an empty group ID: {sample_id!r}") | ||
| if generation_index < 0: | ||
| raise ValueError( | ||
| f"Sample ID has a negative generation index: {sample_id!r}" | ||
| ) | ||
| generation_indices = group_to_generation_indices.setdefault(group_id, set()) | ||
| if generation_index in generation_indices: | ||
| raise ValueError(f"Duplicate generation index in sample ID: {sample_id!r}") | ||
| generation_indices.add(generation_index) | ||
| group_indices.append(group_to_index.setdefault(group_id, len(group_to_index))) | ||
| if expected_group_size is not None: | ||
| expected_indices = set(range(expected_group_size)) | ||
| for group_id, generation_indices in group_to_generation_indices.items(): | ||
| if generation_indices != expected_indices: | ||
| raise ValueError( | ||
| f"Rollout group {group_id!r} has generation indices " | ||
| f"{sorted(generation_indices)}, expected " | ||
| f"{list(range(expected_group_size))}" | ||
| ) | ||
| return torch.tensor(group_indices, device=device, dtype=torch.long)[:, None] | ||
|
|
||
|
|
||
| def calculate_baseline_and_std_per_prompt( | ||
| prompts: torch.Tensor, | ||
|
Contributor
There was a problem hiding this comment. 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 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: No Deliberately not asking you to rename the |
||
| rewards: torch.Tensor, | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
single_controller.py:19771 action item — low severity. (PR-introduced.)
This runs unconditionally, before the
_is_ppobranch, so it fires on paths that have no use for the value:gae/raw_reward(ppo.py:1135restricts the factory;GAEConfig.nameisLiteral["gae", "raw_reward"]), and neither readsprompt_idsat all — we checked by AST rather than docstring, since**kwargshides unused params. Zero load references in either body.has_valid_training_tokensis False andcompute_advantageis never called.So on those paths this is both an unconditional cost and, more importantly, an unconditional new
ValueError— the strictexpected_group_sizecheck 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}"atpayload.py:126, eachPromptGroupRecordholds exactlynum_generations_per_promptcompletions, andconfig.py:736rejects 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 consumesprompt_ids— or at minimum below the early-return forhas_valid_training_tokens.