Group GRPO reward normalization by prompt instead of reshaping the batch

.github/workflows/pr-test.yml

@@ -597,6 +597,10 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_cispo_loss.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_reward_utils.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_ppo_logprob_entropy.py"

.github/workflows/pr-test.yml.j2

@@ -74,6 +74,7 @@
         {'test_file': 'test_logprob_response_spans.py', 'num_gpus': 0},
         {'test_file': 'test_value_temperature.py', 'num_gpus': 0},
         {'test_file': 'test_cispo_loss.py', 'num_gpus': 0},
+        {'test_file': 'test_reward_utils.py', 'num_gpus': 0},
         {'test_file': 'test_ppo_logprob_entropy.py', 'num_gpus': 0},
         {'test_file': 'test_rm_f1.py', 'num_gpus': 0},
         {'test_file': 'test_rm_gpqa.py', 'num_gpus': 0},

examples/coding_agent_rl/README.md

@@ -184,6 +184,7 @@ prompt-base restarts.
 - `generate()` returns `list[Sample]` — one Sample per root-to-leaf chain in the per-session message tree.
 - Per-trajectory reward is split as `reward / K` across chains; `rollout_id` is shared so the per-rollout-mean loss reducer still counts the trajectory once.
 - Sub-agent dispatch and auto-compaction increase `K` (each prompt-prefix divergence forks a new branch), so the effective batch after flatten can be much larger than `rollout_batch_size * n_samples_per_prompt`.
+- The per-prompt Sample count is therefore uneven. GRPO's reward normalization handles that: `_post_process_rewards` groups by `Sample.group_index` (the data source's per-prompt counter), not by reshaping the flat batch to `n_samples_per_prompt`, so each prompt is still centered against itself.
 
 ## Porting to a New Sandbox Backend
 

slime/ray/rollout.py

@@ -18,6 +18,7 @@
 from slime.backends.sglang_utils.sglang_config import ModelConfig, ServerGroupConfig, SglangConfig
 from slime.backends.sglang_utils.sglang_engine import SGLangEngine
 from slime.rollout.base_types import call_rollout_fn
+from slime.rollout.reward_utils import normalize_rewards_by_group
 from slime.utils import logging_utils
 from slime.utils.data import get_source
 from slime.utils.dp_schedule import build_dp_schedule
@@ -694,21 +695,29 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]):
             self.args.advantage_estimator in ["grpo", "gspo", "cispo", "reinforce_plus_plus_baseline"]
             and self.args.rewards_normalization
         ):
-            # group norm
-            rewards = torch.tensor(raw_rewards, dtype=torch.float)
-            if rewards.shape[-1] == self.args.n_samples_per_prompt * self.args.rollout_batch_size:
-                rewards = rewards.reshape(-1, self.args.n_samples_per_prompt)
-            else:
-                # when samples count are not equal in each group
-                rewards = rewards.view(-1, rewards.shape[-1])
-            mean = rewards.mean(dim=-1, keepdim=True)
-            rewards = rewards - mean
-
-            if self.args.advantage_estimator in ["grpo", "gspo", "cispo"] and self.args.grpo_std_normalization:
-                std = rewards.std(dim=-1, keepdim=True)
-                rewards = rewards / (std + 1e-6)
-
-            return raw_rewards, rewards.flatten().tolist()
+            # Group norm, one group per prompt. Grouping is by group_index rather
+            # than by reshaping to (-1, n_samples_per_prompt): a rollout function
+            # that fans one trajectory out to several training samples (sub-agent
+            # dispatch, auto-compaction, token-drift forks) makes the per-prompt
+            # count uneven, so the total no longer equals
+            # n_samples_per_prompt * rollout_batch_size and a reshape cannot
+            # express the grouping at all. The previous fallback for that case
+            # normalized over the whole batch as ONE group, which silently
+            # replaced per-prompt centering with batch-wide centering: measured on
+            # agent-fanout dumps, 80% of advantages came out with the wrong sign
+            # and prompts that solved nothing picked up non-zero advantages from
+            # other prompts in the batch.
+            normalize_std = (
+                self.args.advantage_estimator in ["grpo", "gspo", "cispo"] and self.args.grpo_std_normalization
+            )
+            rewards = normalize_rewards_by_group(
+                raw_rewards,
+                [sample.group_index for sample in samples],
+                [sample.rollout_id for sample in samples],
+                normalize_std=normalize_std,
+                fallback_group_size=self.args.n_samples_per_prompt,
+            )
+            return raw_rewards, rewards
 
         return raw_rewards, raw_rewards
 
@@ -744,6 +753,12 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
             "truncated": [1 if sample.status == Sample.Status.TRUNCATED else 0 for sample in samples],
             "sample_indices": [sample.index for sample in samples],
             "rollout_ids": rollout_ids,
+            # Deliberately NOT routed through normalize_rewards_by_group's grouping:
+            # that one rejects a partially-missing group_index, which is right for
+            # rewards (a wrong baseline corrupts training) but wrong here. This
+            # feeds pass-rate logging only, so a batch mixing samples with and
+            # without group_index should degrade to positional grouping rather
+            # than crash the run over a metric.
             "raw_reward_group_indices": [
                 sample.group_index if sample.group_index is not None else i // self.args.n_samples_per_prompt
                 for i, sample in enumerate(samples)

slime/rollout/_fanout_test_helpers.py

@@ -14,16 +14,13 @@
     framework (per-rollout step splitter, per-rollout-mean reducer,
     ``_validate_rollout_id_annotated`` validator) is built around.
 
-  - ``grpo_normalize_by_group_index``: replaces the default
-    ``_post_process_rewards`` reshape-by-shape logic with a proper
-    ``group_index``-keyed grouping. The default at
-    ``slime/ray/rollout.py:618`` assumes every prompt produced exactly
-    ``n_samples_per_prompt`` samples and reshapes by that constant; when
-    compact/fanout makes the per-prompt count uneven, the reshape fails
-    and the fallback ``view(-1, total)`` collapses everything into ONE
-    group, destroying per-prompt centering. ``group_index`` (set by the
-    data source per-prompt, preserved through ``deepcopy``) is the right
-    key here.
+  - ``grpo_normalize_by_group_index``: per-prompt GRPO reward
+    normalization, still wired into the e2e test above via
+    ``--custom-reward-post-process-path``. It is now redundant with the
+    default (see ``slime/rollout/reward_utils.py``), but is kept both
+    because users may have the flag pointed at it and because it gives
+    the normalization tests an oracle that is not the implementation
+    under test.
 """
 
 import copy
@@ -66,28 +63,29 @@ async def compact_generate(args, sample, sampling_params):
         # Critical invariant: all siblings share ``rollout_id`` so the
         # per-rollout reducer aggregates them as ONE rollout (not N) and
         # the rollout-aware step splitter keeps them in the same step.
-        # ``group_index`` is inherited via ``deepcopy`` and is what the
-        # post-process reward hook below groups on for GRPO normalize.
+        # ``group_index`` is inherited via ``deepcopy`` so production reward
+        # normalization keeps the siblings in their prompt group.
         s.rollout_id = sample.index
         siblings.append(s)
     return siblings
 
 
 def grpo_normalize_by_group_index(args, samples):
-    """Drop-in ``--custom-reward-post-process-path`` for compact/fanout.
-
-    The default ``_post_process_rewards`` (``slime/ray/rollout.py:618``)
-    reshapes the flat reward tensor as ``(-1, n_samples_per_prompt)``
-    when ``total == n_samples_per_prompt * rollout_batch_size``, falling
-    back to ``view(-1, total)`` (= one giant group) otherwise. With
-    fanout the count per prompt is uneven, so the fallback fires and
-    centering is computed across ALL samples in the batch instead of
-    per-prompt — that's silently wrong for GRPO.
-
-    This helper groups by ``Sample.group_index`` (the data-source-set
-    per-prompt counter, preserved through deepcopy in
-    ``compact_generate``) and applies the same mean-center + optional
-    std-normalize the default does, just with the correct grouping.
+    """Reference implementation of per-prompt GRPO reward normalization.
+
+    Equivalent to what the default ``_post_process_rewards`` now does via
+    ``slime.rollout.reward_utils.normalize_rewards_by_group``: group by
+    ``Sample.group_index`` -- the data-source-set per-prompt counter, preserved
+    through the deepcopy in ``compact_generate`` -- then mean-center and
+    optionally std-normalize within each group.
+
+    Written as a ``--custom-reward-post-process-path`` workaround back when the
+    default built its groups by reshaping to ``(-1, n_samples_per_prompt)`` and
+    fell back to one batch-wide group whenever fan-out made the per-prompt count
+    uneven. That fallback is gone, so this is no longer load-bearing, but it is
+    kept: the e2e fan-out test still passes the flag, users may have it
+    configured too, and duplicating the logic gives the normalization tests an
+    oracle that is not the implementation under test.
 
     Returns ``(raw_rewards, normalized_rewards)`` matching the input
     ``samples`` order — same shape as the default's return contract.

slime/rollout/reward_utils.py

@@ -0,0 +1,98 @@
+"""Reward normalization for group-relative advantage estimators.
+
+Grouping matches how the loss aggregates. An agentic rollout can split one
+attempt into several training samples (sub-agent dispatch, auto-compaction,
+token-drift forks), and those siblings share ``rollout_id``. The loss already
+collapses them into one per-attempt mean (``rollout_mask_sums`` ->
+``cp_utils.get_sum_of_sample_mean``, which gives each attempt total weight 1.0
+however many segments it produced), ``dp_schedule`` sizes a training step in
+rollouts rather than samples, and ``compute_grouped_pass_rate`` deduplicates by
+``(group_index, rollout_id)``. So the baseline has to be per-attempt too --
+otherwise a heavily-forked attempt votes several times in its own group's mean.
+"""
+
+from collections import defaultdict
+
+import torch
+
+
+def normalize_rewards_by_group(
+    rewards: list[float],
+    group_indices: list[int | None],
+    rollout_ids: list[int | str | None] | None = None,
+    *,
+    normalize_std: bool,
+    fallback_group_size: int | None = None,
+) -> list[float]:
+    """Normalize rewards within the sample group that produced each response.
+
+    Grouping is two-level: ``group_indices`` selects the prompt group, then
+    ``rollout_ids`` reduces that group to one reward per attempt, so mean and std
+    are attempt-level and match the loss reducer. Each attempt's normalized
+    reward is broadcast back to all of its samples.
+
+    ``rollout_ids=None``, or a ``None`` entry, treats the sample as its own
+    attempt -- correct for the default path, where one execution is one training
+    sample and ``Sample.rollout_id`` is left unset. That case reduces to plain
+    per-group normalization, bit for bit.
+
+    ``fallback_group_size`` preserves fixed-size custom rollouts that predate
+    ``Sample.group_index``. Uneven groups must provide explicit identities.
+    """
+    if len(rewards) != len(group_indices):
+        raise ValueError(
+            f"rewards and group_indices must have the same length, got {len(rewards)} and {len(group_indices)}"
+        )
+    if rollout_ids is not None and len(rollout_ids) != len(rewards):
+        raise ValueError(
+            f"rewards and rollout_ids must have the same length, got {len(rewards)} and {len(rollout_ids)}"
+        )
+
+    if group_indices and all(group_index is None for group_index in group_indices):
+        if fallback_group_size is None or fallback_group_size <= 0 or len(group_indices) % fallback_group_size != 0:
+            raise ValueError("group_index is required when reward groups are not uniformly sized")
+        group_indices = [position // fallback_group_size for position in range(len(group_indices))]
+
+    positions_by_group: dict[int, list[int]] = defaultdict(list)
+    for position, group_index in enumerate(group_indices):
+        if group_index is None:
+            raise ValueError(
+                f"group_index is required for reward normalization, but sample at position {position} has none"
+            )
+        positions_by_group[group_index].append(position)
+
+    reward_tensor = torch.tensor(rewards, dtype=torch.float)
+    normalized_rewards = torch.empty_like(reward_tensor)
+    for positions in positions_by_group.values():
+        # One entry per attempt. A missing rollout id makes the sample its own
+        # attempt; the tuple key keeps that apart from a real id of the same
+        # value, since rollout_id may be a string on custom rollout paths.
+        # Segments of one attempt carry the same reward, so the reduction is
+        # lossless -- but take the max rather than first-write-wins so that a
+        # custom path emitting disagreeing segments degrades to "solved if any
+        # segment solved" instead of depending on sample order.
+        attempt_rewards: dict[tuple[str, object], float] = {}
+        for position in positions:
+            rollout_id = None if rollout_ids is None else rollout_ids[position]
+            key = ("position", position) if rollout_id is None else ("rollout", rollout_id)
+            reward = float(rewards[position])
+            attempt_rewards[key] = max(attempt_rewards.get(key, reward), reward)
+
+        attempt_tensor = torch.tensor(list(attempt_rewards.values()), dtype=torch.float)
+        mean = attempt_tensor.mean()
+        # Take the std of the centered attempts, not of the raw ones. The two are
+        # equal in exact arithmetic but not in float32, and centering first is
+        # what per-sample normalization did -- so the one-attempt-per-sample case
+        # stays bit-identical to it rather than drifting by ~1e-7.
+        centered_attempts = attempt_tensor - mean
+        # A lone attempt is already 0 after centering, and torch.std of one
+        # element is NaN (it is the sample std), which would poison the gradient.
+        # Reshaping could never produce a size-1 group; grouping by prompt can.
+        std = centered_attempts.std() if normalize_std and len(attempt_rewards) > 1 else None
+
+        group_rewards = reward_tensor[positions] - mean
+        if std is not None:
+            group_rewards = group_rewards / (std + 1e-6)
+        normalized_rewards[positions] = group_rewards
+
+    return normalized_rewards.tolist()

tests/test_qwen2.5_0.5B_fanout_short.py

@@ -100,15 +100,12 @@ def execute():
         "--global-batch-size 4 "
         "--balance-data "
         "--custom-generate-function-path slime.rollout._fanout_test_helpers.compact_generate "
-        # GRPO normalization needs per-prompt grouping. The default
-        # ``_post_process_rewards`` (slime/ray/rollout.py:618) reshapes
-        # by ``n_samples_per_prompt`` and falls back to "one big group"
-        # when the per-prompt count is uneven — fan-out trips exactly
-        # that fallback. The helper here groups by ``Sample.group_index``
-        # (the per-prompt counter the data source stamps; deepcopy in
-        # compact_generate preserves it across siblings) so each prompt's
-        # siblings normalize against each other, matching the GRPO
-        # semantics the default targets in the uniform case.
+        # Redundant since the default ``_post_process_rewards`` started
+        # grouping by ``Sample.group_index`` itself, which handles the
+        # uneven per-prompt counts fan-out produces. Kept because dropping
+        # it changes what this test exercises, and that swap should land
+        # with an actual e2e run behind it rather than on the argument
+        # that it ought to be equivalent.
         "--custom-reward-post-process-path slime.rollout._fanout_test_helpers.grpo_normalize_by_group_index "
     )