From 01921e4400242e201e4ffb1e90cfec0f240df28d Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Tue, 14 Jul 2026 04:47:21 +0000 Subject: [PATCH 1/2] Fix reward normalization for uneven sample groups --- .github/workflows/pr-test.yml | 4 ++ .github/workflows/pr-test.yml.j2 | 1 + slime/ray/rollout.py | 27 ++++----- slime/rollout/_fanout_test_helpers.py | 76 ++++--------------------- slime/rollout/reward_utils.py | 41 +++++++++++++ tests/test_qwen2.5_0.5B_fanout_short.py | 10 ---- tests/test_reward_utils.py | 58 +++++++++++++++++++ 7 files changed, 126 insertions(+), 91 deletions(-) create mode 100644 slime/rollout/reward_utils.py create mode 100644 tests/test_reward_utils.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 761496b949..9bd1af8e86 100644 --- a/.github/workflows/pr-test.yml +++ b/.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" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 415586cb84..4742f33e82 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.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}, diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index d0653522f7..dbd390c8d1 100644 --- a/slime/ray/rollout.py +++ b/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,17 @@ 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 diff --git a/slime/rollout/_fanout_test_helpers.py b/slime/rollout/_fanout_test_helpers.py index e065dd13dd..0e178af20e 100644 --- a/slime/rollout/_fanout_test_helpers.py +++ b/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 diff --git a/slime/rollout/reward_utils.py b/slime/rollout/reward_utils.py new file mode 100644 index 0000000000..4f60c13d48 --- /dev/null +++ b/slime/rollout/reward_utils.py @@ -0,0 +1,41 @@ +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() diff --git a/tests/test_qwen2.5_0.5B_fanout_short.py b/tests/test_qwen2.5_0.5B_fanout_short.py index 292f18913c..eb6c62b7b7 100644 --- a/tests/test_qwen2.5_0.5B_fanout_short.py +++ b/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 = ( diff --git a/tests/test_reward_utils.py b/tests/test_reward_utils.py new file mode 100644 index 0000000000..9f251b88c8 --- /dev/null +++ b/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__])) From 55d49e4c53fda48056a2ed8a36e2f0066a87c35d Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Tue, 14 Jul 2026 04:55:14 +0000 Subject: [PATCH 2/2] Apply repository formatting --- slime/ray/rollout.py | 3 +-- slime/rollout/reward_utils.py | 8 ++++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index dbd390c8d1..944a36ba95 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -672,8 +672,7 @@ def _post_process_rewards(self, samples: list[Sample] | list[list[Sample]]): and self.args.rewards_normalization ): normalize_std = ( - self.args.advantage_estimator in ["grpo", "gspo", "cispo"] - and self.args.grpo_std_normalization + self.args.advantage_estimator in ["grpo", "gspo", "cispo"] and self.args.grpo_std_normalization ) rewards = normalize_rewards_by_group( raw_rewards, diff --git a/slime/rollout/reward_utils.py b/slime/rollout/reward_utils.py index 4f60c13d48..20c1d9f939 100644 --- a/slime/rollout/reward_utils.py +++ b/slime/rollout/reward_utils.py @@ -16,7 +16,9 @@ def normalize_rewards_by_group( ``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)}") + 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: @@ -26,7 +28,9 @@ def normalize_rewards_by_group( 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") + 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)