fix: normalize rewards by explicit sample groups - #2204
Open
morluto wants to merge 2 commits into
Open
Conversation
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
morluto
changed the title
FunJim
added a commit
to FunJim/slime
that referenced
this pull request
Aug 19, 2026
_post_process_rewards built its groups by reshaping the flat reward vector to (-1, n_samples_per_prompt), falling back to view(-1, total) -- one group spanning the whole batch -- whenever the total did not equal n_samples_per_prompt * rollout_batch_size. A reshape cannot express uneven groups, and agent fan-out makes them uneven on every batch: one trajectory becomes several training samples through sub-agent dispatch, auto-compaction and token-drift forks, so a run with n_samples_per_prompt=8 produced group sizes like [8, 9, 10, 12, 28]. The fallback fired every time and silently replaced per-prompt centering with batch-wide centering, which is not GRPO -- the group mean is the prompt's difficulty baseline, and subtracting the batch mean instead leaves task difficulty in the advantage. Measured against per-prompt normalization on dumped rollouts: 80.0% of advantages had the wrong sign on one run (18.7% on a second, larger one), and non-zero advantages inflated from 454 to 1898. The inflation is the symptom -- prompts that solved nothing picked up advantages purely because other prompts in the batch solved something. Normalize by Sample.group_index, the data source's per-prompt counter, which survives the deepcopy that builds a prompt's samples and so identifies the prompt however many samples a trajectory emitted. Batches whose group_index is missing everywhere still fall back to fixed-size slices for custom rollouts that predate the field; a partially-missing group_index now raises rather than guessing a baseline. On the uniform layout the reshape handled, the result is elementwise identical. Grouping by prompt can produce a size-1 group, which reshaping never could. torch.std of one element is the sample std, i.e. NaN, so std normalization is skipped there; centering alone already sends a lone sample to zero. raw_reward_group_indices deliberately keeps its own permissive grouping: it feeds pass-rate logging, where a batch mixing samples with and without group_index should degrade to positional grouping rather than kill the run over a metric. The fan-out e2e test drops its --custom-reward-post-process-path workaround so it exercises the default path, which is where the collapse happened. grpo_normalize_by_group_index stays as an independent oracle for the new tests and because users may still have it configured as a hook. The normalization itself -- slime/rollout/reward_utils.py and the call site that replaces the reshape -- is taken from THUDM#2204 by morluto, which has been open since 2026-07-14 without review. Cherry-picked here rather than waited on because our agent-fanout training is affected on every batch. Added on top: the measurements above, the size-1 NaN rationale, manager-level wiring tests, four more unit tests, and keeping grpo_normalize_by_group_index so a configured --custom-reward-post-process-path does not break. Upstream is the same fix, so this should drop out cleanly when THUDM#2204 lands. Prior reports of the same defect: THUDM#904, THUDM#1414 (closed), THUDM#2230, and PRs THUDM#487, THUDM#1415, THUDM#1918.
FunJim
added a commit
to FunJim/slime
that referenced
this pull request
Aug 19, 2026
_post_process_rewards built its groups by reshaping the flat reward vector to (-1, n_samples_per_prompt), falling back to view(-1, total) -- one group spanning the whole batch -- whenever the total did not equal n_samples_per_prompt * rollout_batch_size. A reshape cannot express uneven groups, and agent fan-out makes them uneven on every batch: one trajectory becomes several training samples through sub-agent dispatch, auto-compaction and token-drift forks, so a run with n_samples_per_prompt=8 produced group sizes like [8, 9, 10, 12, 28]. The fallback fired every time and silently replaced per-prompt centering with batch-wide centering, which is not GRPO -- the group mean is the prompt's difficulty baseline, and subtracting the batch mean instead leaves task difficulty in the advantage. Measured against per-prompt normalization on dumped rollouts: 80.0% of advantages had the wrong sign on one run (18.7% on a second, larger one), and non-zero advantages inflated from 454 to 1898. The inflation is the symptom -- prompts that solved nothing picked up advantages purely because other prompts in the batch solved something. Normalize by Sample.group_index, the data source's per-prompt counter, which survives the deepcopy that builds a prompt's samples and so identifies the prompt however many samples a trajectory emitted. Batches whose group_index is missing everywhere still fall back to fixed-size slices for custom rollouts that predate the field; a partially-missing group_index now raises rather than guessing a baseline. On the uniform layout the reshape handled, the result is elementwise identical. Grouping by prompt can produce a size-1 group, which reshaping never could. torch.std of one element is the sample std, i.e. NaN, so std normalization is skipped there; centering alone already sends a lone sample to zero. raw_reward_group_indices deliberately keeps its own permissive grouping: it feeds pass-rate logging, where a batch mixing samples with and without group_index should degrade to positional grouping rather than kill the run over a metric. The fan-out e2e test drops its --custom-reward-post-process-path workaround so it exercises the default path, which is where the collapse happened. grpo_normalize_by_group_index stays as an independent oracle for the new tests and because users may still have it configured as a hook. The normalization itself -- slime/rollout/reward_utils.py and the call site that replaces the reshape -- is taken from THUDM#2204 by morluto, which has been open since 2026-07-14 without review. Cherry-picked here rather than waited on because our agent-fanout training is affected on every batch. Added on top: the measurements above, the size-1 NaN rationale, manager-level wiring tests, four more unit tests, and keeping grpo_normalize_by_group_index so a configured --custom-reward-post-process-path does not break. Upstream is the same fix, so this should drop out cleanly when THUDM#2204 lands. Prior reports of the same defect: THUDM#904, THUDM#1414 (closed), THUDM#2230, and PRs THUDM#487, THUDM#1415, THUDM#1918.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
RolloutManager._post_process_rewardsinferred GRPO reward groups from the total sample count. When a custom rollout returns uneven numbers of samples per prompt, the fallback reshaped the reward vector to one row and centered rewards across the entire batch.For groups of sizes 4, 3, and 4, a prompt whose three rewards are all
5.0received advantages of-1.090909instead of zero.Changes
group_indexand restore results in input order.group_indexis missing; reject mixed or uneven unidentified groups instead of silently using a global baseline.Validation
\n-
uvx --from pre-commit pre-commit run --all-files --show-diff-on-failure --color=alwaysPYTHONPATH=. python3 tests/test_reward_utils.py(5 passed)ruff check slime/ray/rollout.py slime/rollout/reward_utils.py slime/rollout/_fanout_test_helpers.py tests/test_reward_utils.py tests/test_qwen2.5_0.5B_fanout_short.pypython3 -m py_compile slime/ray/rollout.py slime/rollout/reward_utils.py slime/rollout/_fanout_test_helpers.py tests/test_reward_utils.py tests/test_qwen2.5_0.5B_fanout_short.pypython3 .github/workflows/generate_github_workflows.pytests/test_qwen2.5_0.5B_fanout_short.pywas not run locally because it requires four GPUs; it remains registered in the GPU CI matrix.Tracking issue
Fixes #2230