fix: normalize rewards by explicit sample groups

.github/workflows/pr-test.yml

@@ -621,6 +621,10 @@ jobs:
               "num_gpus": 0,
               "test_file": "test_rm_deepscaler.py"
             },
+            {
+              "num_gpus": 0,
+              "test_file": "test_reward_utils.py"
+            },
             {
               "num_gpus": 0,
               "test_file": "test_sample.py"

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

@@ -80,6 +80,7 @@
         {'test_file': 'test_rm_math.py', 'num_gpus': 0},
         {'test_file': 'test_rm_math_dapo.py', 'num_gpus': 0},
         {'test_file': 'test_rm_deepscaler.py', 'num_gpus': 0},
+        {'test_file': 'test_reward_utils.py', 'num_gpus': 0},
         {'test_file': 'test_sample.py', 'num_gpus': 0},
         {'test_file': 'test_rollout_validation.py', 'num_gpus': 0},
         {'test_file': 'test_placement_group.py', 'num_gpus': 0},

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
@@ -670,21 +671,16 @@ 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()
+            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],
+                normalize_std=normalize_std,
+                fallback_group_size=self.args.n_samples_per_prompt,
+            )
+            return raw_rewards, rewards
 
         return raw_rewards, raw_rewards
 

slime/rollout/_fanout_test_helpers.py

@@ -3,32 +3,18 @@
 The underscore prefix marks this as test infrastructure — it is not part
 of the user-facing slime API and is not re-exported anywhere. It lives
 in ``slime/`` only so the test can reference it by a dotted module path
-(``--custom-generate-function-path`` / ``--custom-reward-post-process-path``
-resolve a string via ``importlib.import_module``, which can't handle the
-dots in the e2e test's filename).
-
-Two helpers:
-
-  - ``compact_generate``: fans one input sample out to N siblings
-    sharing the same ``rollout_id``. That's the contract the rest of the
-    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.
+(``--custom-generate-function-path`` resolves a string via
+``importlib.import_module``, which can't handle the dots in the e2e
+test's filename).
+
+``compact_generate`` fans one input sample out to N siblings sharing the
+same ``rollout_id``. That's the contract the rest of the framework
+(per-rollout step splitter, per-rollout-mean reducer,
+``_validate_rollout_id_annotated`` validator) is built around.
 """
 
 import copy
 import os
-from collections import defaultdict
 
 
 MAX_FANOUT = 3
@@ -66,50 +52,8 @@ 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.
-
-    Returns ``(raw_rewards, normalized_rewards)`` matching the input
-    ``samples`` order — same shape as the default's return contract.
-    """
-    import torch
-
-    raw_rewards = [s.get_reward_value(args) for s in samples]
-
-    # group_index → list of (original_position, raw_reward)
-    groups: dict[int, list[tuple[int, float]]] = defaultdict(list)
-    for i, s in enumerate(samples):
-        groups[s.group_index].append((i, raw_rewards[i]))
-
-    out = [0.0] * len(samples)
-    use_std = getattr(args, "grpo_std_normalization", True)
-    for indexed in groups.values():
-        positions = [p for p, _ in indexed]
-        rewards = torch.tensor([r for _, r in indexed], dtype=torch.float)
-        rewards = rewards - rewards.mean()
-        if use_std:
-            rewards = rewards / (rewards.std() + 1e-6)
-        for pos, r in zip(positions, rewards.tolist(), strict=True):
-            out[pos] = r
-
-    return raw_rewards, out

slime/rollout/reward_utils.py

@@ -0,0 +1,45 @@
+from collections import defaultdict
+
+import torch
+
+
+def normalize_rewards_by_group(
+    rewards: list[float],
+    group_indices: list[int | None],
+    *,
+    normalize_std: bool,
+    fallback_group_size: int | None = None,
+) -> list[float]:
+    """Normalize rewards within the sample group that produced each response.
+
+    ``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 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():
+        group_rewards = reward_tensor[positions]
+        group_rewards = group_rewards - group_rewards.mean()
+        if normalize_std and len(positions) > 1:
+            group_rewards = 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,16 +100,6 @@ 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.
-        "--custom-reward-post-process-path slime.rollout._fanout_test_helpers.grpo_normalize_by_group_index "
     )
 
     perf_args = (

tests/test_reward_utils.py

@@ -0,0 +1,58 @@
+import pytest
+
+from slime.rollout.reward_utils import normalize_rewards_by_group
+
+NUM_GPUS = 0
+
+
+@pytest.mark.unit
+def test_normalize_rewards_uses_explicit_uneven_groups():
+    rewards = [0.0, 1.0, 2.0, 3.0, 5.0, 5.0, 5.0, 10.0, 11.0, 12.0, 13.0]
+    group_indices = [0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2]
+
+    normalized = normalize_rewards_by_group(rewards, group_indices, normalize_std=False)
+
+    assert normalized == pytest.approx([-1.5, -0.5, 0.5, 1.5, 0.0, 0.0, 0.0, -1.5, -0.5, 0.5, 1.5])
+
+
+@pytest.mark.unit
+def test_normalize_rewards_preserves_order_and_zeroes_singletons():
+    rewards = [-1.0, 4.0, 0.0, 7.0, 1.0, 4.0]
+    group_indices = [10, 20, 10, 30, 10, 20]
+
+    normalized = normalize_rewards_by_group(rewards, group_indices, normalize_std=True)
+
+    assert normalized == pytest.approx([-1.0, 0.0, 0.0, 0.0, 1.0, 0.0], abs=1e-5)
+
+
+@pytest.mark.unit
+def test_normalize_rewards_requires_group_identity():
+    with pytest.raises(ValueError, match="group_index is required.*position 1"):
+        normalize_rewards_by_group([1.0, 2.0], [0, None], normalize_std=False)
+
+
+@pytest.mark.unit
+def test_normalize_rewards_supports_legacy_fixed_size_groups():
+    normalized = normalize_rewards_by_group(
+        [1.0, 3.0, 10.0, 10.0],
+        [None, None, None, None],
+        normalize_std=False,
+        fallback_group_size=2,
+    )
+
+    assert normalized == pytest.approx([-1.0, 1.0, 0.0, 0.0])
+
+
+@pytest.mark.unit
+def test_normalize_rewards_rejects_unidentified_uneven_groups():
+    with pytest.raises(ValueError, match="group_index is required when reward groups are not uniformly sized"):
+        normalize_rewards_by_group(
+            [1.0, 2.0, 3.0],
+            [None, None, None],
+            normalize_std=False,
+            fallback_group_size=2,
+        )
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__]))