diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 26ba9dff02..b07480718a 100644 --- a/.github/workflows/pr-test.yml +++ b/.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" diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index ef7aa001e6..3193918028 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.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}, diff --git a/examples/coding_agent_rl/README.md b/examples/coding_agent_rl/README.md index f16c7c3dfc..c3af274d1f 100644 --- a/examples/coding_agent_rl/README.md +++ b/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 diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 285a1e619e..0564481e7a 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 @@ -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) diff --git a/slime/rollout/_fanout_test_helpers.py b/slime/rollout/_fanout_test_helpers.py index e065dd13dd..b670c21085 100644 --- a/slime/rollout/_fanout_test_helpers.py +++ b/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. diff --git a/slime/rollout/reward_utils.py b/slime/rollout/reward_utils.py new file mode 100644 index 0000000000..43cae57f4b --- /dev/null +++ b/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() diff --git a/tests/test_qwen2.5_0.5B_fanout_short.py b/tests/test_qwen2.5_0.5B_fanout_short.py index 292f18913c..41058171e8 100644 --- a/tests/test_qwen2.5_0.5B_fanout_short.py +++ b/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 " ) diff --git a/tests/test_reward_utils.py b/tests/test_reward_utils.py new file mode 100644 index 0000000000..6ff67c5f94 --- /dev/null +++ b/tests/test_reward_utils.py @@ -0,0 +1,269 @@ +from pathlib import Path + +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, + ) + + +# --------------------------------------------------------------------------- +# Additions beyond the upstream PR: the properties that make this a bug fix +# rather than a refactor. Each of these fails against the reshape-based +# grouping it replaced. +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_a_solved_prompt_does_not_lift_an_unsolved_one(): + # The concrete symptom of the collapsed grouping. Prompt 1 solved nothing, so + # every one of its advantages must be 0 no matter what prompt 0 did; under + # batch-wide centering they came out negative purely because prompt 0 solved. + rewards = [1.0, 0.0] + [0.0] * 6 + group_indices = [0, 0] + [1] * 6 + + normalized = normalize_rewards_by_group(rewards, group_indices, normalize_std=True) + + assert normalized[2:] == pytest.approx([0.0] * 6, abs=1e-6) + assert normalized[0] > 0.0 > normalized[1] + + +@pytest.mark.unit +def test_a_lone_sample_is_finite(): + # Grouping by prompt can produce a size-1 group, which reshaping never could. + # torch.std of one element is NaN, and a NaN advantage poisons the gradient. + normalized = normalize_rewards_by_group([1.0, 1.0, 0.0], [0, 1, 1], normalize_std=True) + + assert all(value == value for value in normalized) # NaN != NaN + assert normalized[0] == pytest.approx(0.0) + + +@pytest.mark.unit +def test_matches_the_reshape_result_when_counts_are_uniform(): + # Backward compatibility: on the uniform layout the old reshape handled, the + # grouping must be elementwise identical to it. + import torch + + rewards = [1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0] + group_size = 4 + + normalized = normalize_rewards_by_group( + rewards, [i // group_size for i in range(len(rewards))], normalize_std=True + ) + + reference = torch.tensor(rewards, dtype=torch.float).reshape(-1, group_size) + reference = reference - reference.mean(dim=-1, keepdim=True) + reference = reference / (reference.std(dim=-1, keepdim=True) + 1e-6) + assert normalized == pytest.approx(reference.flatten().tolist(), abs=1e-5) + + +@pytest.mark.unit +def test_length_mismatch_is_rejected(): + with pytest.raises(ValueError, match="same length"): + normalize_rewards_by_group([1.0, 2.0], [0], normalize_std=False) + + +# --------------------------------------------------------------------------- +# Attempt-level baseline. An agentic rollout splits one attempt into several +# samples sharing a rollout_id, and the loss already collapses those into a +# single per-attempt mean (measured: total weight 1.0 per attempt regardless of +# segment count). The baseline has to use the same unit, or a heavily-forked +# attempt votes several times in its own group's mean. +# --------------------------------------------------------------------------- + + +@pytest.mark.unit +def test_segment_count_does_not_move_the_baseline(): + # Same 8 attempts, 5 solved. Forking one of them must not shift the mean. + unforked = normalize_rewards_by_group( + [1.0] * 5 + [0.0] * 3, + [0] * 8, + list(range(8)), + normalize_std=False, + ) + # Attempt 5 (a failure) now arrives as three segments. + forked = normalize_rewards_by_group( + [1.0] * 5 + [0.0, 0.0, 0.0] + [0.0, 0.0], + [0] * 10, + [0, 1, 2, 3, 4, 5, 5, 5, 6, 7], + normalize_std=False, + ) + + # Solved attempts keep the same advantage either way: 1 - 5/8. + assert forked[0] == pytest.approx(unforked[0]) + assert forked[0] == pytest.approx(0.375) + # Counting samples instead would give 1 - 5/10 = 0.5. + assert forked[0] != pytest.approx(0.5) + + +@pytest.mark.unit +def test_segments_of_one_attempt_share_its_advantage(): + normalized = normalize_rewards_by_group( + [1.0, 1.0, 1.0, 0.0], + [0] * 4, + ["a", "a", "a", "b"], + normalize_std=True, + ) + + assert normalized[0] == pytest.approx(normalized[1]) == pytest.approx(normalized[2]) + assert normalized[3] < 0.0 < normalized[0] + + +@pytest.mark.unit +def test_a_single_attempt_split_into_segments_has_no_signal(): + # One attempt cannot be compared against anything, however many segments it + # produced. Centering alone must send every segment to 0, with no NaN from + # a one-element std. + normalized = normalize_rewards_by_group([1.0] * 4, [0] * 4, ["a"] * 4, normalize_std=True) + + assert normalized == pytest.approx([0.0] * 4) + + +@pytest.mark.unit +def test_disagreeing_segments_reduce_deterministically(): + # Segments of one attempt should carry the same reward; if a custom path + # breaks that, the reduction must not depend on sample order. + forward = normalize_rewards_by_group([1.0, 0.0, 0.0], [0] * 3, ["a", "a", "b"], normalize_std=False) + reverse = normalize_rewards_by_group([0.0, 1.0, 0.0], [0] * 3, ["a", "a", "b"], normalize_std=False) + + assert forward[2] == pytest.approx(reverse[2]) + # "solved if any segment solved": attempt a counts as 1.0, so mean is 0.5. + assert forward[2] == pytest.approx(-0.5) + + +@pytest.mark.unit +def test_rollout_ids_are_scoped_to_their_prompt_group(): + # Two prompts, each with attempts numbered from scratch. Grouping happens + # first, so identical ids in different groups must not merge. + normalized = normalize_rewards_by_group( + [1.0, 0.0, 0.0, 1.0], + [0, 0, 1, 1], + [0, 1, 0, 1], + normalize_std=False, + ) + + assert normalized == pytest.approx([0.5, -0.5, -0.5, 0.5]) + + +@pytest.mark.unit +def test_omitting_rollout_ids_matches_passing_distinct_ones(): + rewards = [1.0, 0.0, 0.5, 1.0] + groups = [0, 0, 0, 0] + + assert normalize_rewards_by_group(rewards, groups, normalize_std=True) == pytest.approx( + normalize_rewards_by_group(rewards, groups, list(range(4)), normalize_std=True) + ) + + +@pytest.mark.unit +def test_a_none_rollout_id_is_its_own_attempt(): + # The default path leaves rollout_id unset; a mixed batch must not merge the + # unset ones into a single bucket. + normalized = normalize_rewards_by_group( + [1.0, 0.0, 0.0, 0.0], + [0] * 4, + [None, None, None, None], + normalize_std=False, + ) + + assert normalized == pytest.approx([0.75, -0.25, -0.25, -0.25]) + + +@pytest.mark.unit +def test_rollout_ids_length_is_validated(): + with pytest.raises(ValueError, match="rollout_ids must have the same length"): + normalize_rewards_by_group([1.0, 2.0], [0, 0], [0], normalize_std=False) + + +# --------------------------------------------------------------------------- +# Call-site guard. The bug lived in RolloutManager._post_process_rewards, not in +# the function above, so the function passing its own tests is not enough -- a +# revert of the call site would leave every test here green. Asserted against +# the source text rather than by importing slime.ray.rollout, which pulls in +# sglang at module scope and cannot be imported in the CPU test job. Same +# read_text approach as tests/plugin_contracts/test_plugin_runtime_hook_contracts.py. +# --------------------------------------------------------------------------- + +ROLLOUT_SOURCE = Path(__file__).resolve().parents[1] / "slime" / "ray" / "rollout.py" + + +@pytest.mark.unit +def test_rollout_manager_normalizes_through_this_function(): + source = ROLLOUT_SOURCE.read_text() + + assert "from slime.rollout.reward_utils import normalize_rewards_by_group" in source + assert "rewards = normalize_rewards_by_group(" in source + + +@pytest.mark.unit +def test_rollout_manager_passes_per_prompt_group_identity(): + source = ROLLOUT_SOURCE.read_text() + + # Grouping must come from the samples, not from the batch shape, and the + # baseline must be reduced to one entry per attempt. + assert "[sample.group_index for sample in samples]" in source + assert "[sample.rollout_id for sample in samples]" in source + assert "fallback_group_size=self.args.n_samples_per_prompt" in source + + +@pytest.mark.unit +def test_the_batch_wide_collapse_is_gone(): + # The exact expression that normalized every prompt against one shared mean + # whenever fan-out made the per-prompt counts uneven. + source = ROLLOUT_SOURCE.read_text() + + assert "rewards.view(-1, rewards.shape[-1])" not in source + assert "rewards.reshape(-1, self.args.n_samples_per_prompt)" not in source + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__]))