Group GRPO reward normalization by prompt instead of reshaping the batch

[FunJim]

Fixes GRPO reward normalization for the uneven per-prompt sample counts that agent fan-out produces on every batch.

The defect

_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 fan-out makes them uneven by construction: one trajectory becomes several training samples through sub-agent dispatch, auto-compaction and token-drift forks. With n_samples_per_prompt=8 we measured group sizes like [8, 9, 10, 12, 28], so the fallback fired and silently replaced per-prompt centering with batch-wide centering.

How often it fires, per run:

• run A (agent training): every dump, 6/6.
• run B (rollout scoring): 11 dumps of 102 — and those 11 are exactly the dumps that carry any reward at all. The other 91 are all-zero, so they train nothing under either grouping. Every batch with something to learn from went through the collapsed path.

That is not GRPO. The group mean is the prompt's difficulty baseline; subtracting the batch mean instead leaves task difficulty in the advantage.

Measured cost

Against per-prompt normalization, on dumped rollouts from two runs:

	run A	run B
samples	435	7722
advantages that agree	14.3%	75.4%
advantages with the wrong sign	80.0%	18.7%
non-zero advantages: correct → collapsed	25 → 373	454 → 1898

The 4–15x inflation of non-zero advantages is the clearest symptom: prompts that solved nothing received gradient purely because other prompts in the batch solved something.

Run B's agreement rate is higher only because 91 of its 102 dumps are all-zero and agree trivially; the divergence is concentrated in the 11 that matter. Both columns replay the pre-fix code faithfully, reshape branch included — not a blanket assumption that the fallback always fires.

After the fix, both runs match an independently computed per-prompt reference on 100% of samples, with no NaN.

Why it went unnoticed

Three layers of cover:

1. No error. view(-1, total) is legal, and the fallback even carried a comment describing when it fires, so it read as handled.
2. Correct on the uniform layout. Every non-fan-out run has exactly n_samples_per_prompt samples per prompt and takes the reshape branch.
3. Training does not crash. Wrong gradient direction, normal magnitude — indistinguishable from "the task is hard".

Worth noting raw_reward_group_indices, ~50 lines below in the same file, already grouped by group_index. Pass-rate logging was correct; only normalization was not.

Changes

• Normalize by Sample.group_index, the data source's per-prompt counter, which survives the deepcopy that builds a prompt's samples.
• Batches with group_index missing everywhere still fall back to fixed-size slices (custom rollouts predating the field); a partially missing group_index now raises rather than guessing a baseline.
• Skip std normalization for size-1 groups: torch.std of one element is the sample std, i.e. NaN, which would poison the gradient. Reshaping could never produce a size-1 group; grouping by prompt can.
• raw_reward_group_indices deliberately keeps its own permissive grouping — it feeds a metric, and a mixed batch should degrade to positional grouping rather than kill the run.
• The fan-out e2e test is unchanged on every executable line. Its --custom-reward-post-process-path flag is now redundant with the fixed default, but dropping it would change what that test exercises and this fork has no self-hosted runner to validate the swap, so it is left for a change that can be backed by a real e2e run. Only its comment is updated, to stop describing the collapse as current behaviour.

On the uniform layout the old reshape handled, the result is elementwise identical (pinned by a test).

Tests

12 CPU tests in tests/test_reward_utils.py, registered in the cpu-unittest matrix. It imports nothing beyond torch, so it runs there cleanly.

Nine cover the grouping function: uneven groups, order preservation, singletons, legacy fallback, strictness, and elementwise equivalence to the old reshape on a uniform layout — plus the properties that make this a bug fix rather than a refactor (an unsolved prompt getting no signal from a solved one, size-1 groups staying finite).

Three guard the call site. The bug lived in _post_process_rewards, not in the function, so the function passing its own tests would not catch a revert. These assert that rollout.py imports and calls normalize_rewards_by_group, passes [sample.group_index for sample in samples] and fallback_group_size, and no longer contains the collapsing view(-1, rewards.shape[-1]). They check the source text rather than importing the module — slime.ray.rollout pulls in sglang at module scope, which the CPU job does not install; this is the same read_text() approach tests/plugin_contracts/test_plugin_runtime_hook_contracts.py already uses for call sites in that file.

Verified all twelve pass against the fix and that the three call-site guards fail against the pre-fix source, in a venv built from the exact CI install commands (python 3.11, torch-cpu, no sglang) — a local dev environment has sglang and cannot reproduce that job.

Attribution

slime/rollout/reward_utils.py and the call site replacing the reshape are taken from THUDM/slime#2204 by @morluto, open since 2026-07-14 without review. Cherry-picked 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, seven more unit tests including the call-site guards, and keeping grpo_normalize_by_group_index so an already-configured --custom-reward-post-process-path does not break (upstream deletes it, which is an undeclared breaking change given THUDM#904 recommended that flag as the official workaround).

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.

[Tsing-git]

• The modified e2e fan-out test has never been run in the CI for this PR
• Wiring test: zero coverage in CI

[FunJim]

Both valid. Fixed in f843508.

1. e2e fan-out test — reverted, now byte-identical to main on every
executable line (comment only). It can't run here at all: this fork has 0
self-hosted runners and every e2e job is runs-on: self-hosted. Dropping the
redundant flag should land with a real e2e run behind it, not with this PR.

2. Wiring test — deleted. It can't be registered: slime.ray.rollout
imports sglang at module scope, and stubbing needs seven modules.

Replaced with three call-site guards in tests/test_reward_utils.py (in the
matrix), asserting rollout.py calls normalize_rewards_by_group, passes
group_index and fallback_group_size, and no longer contains
view(-1, rewards.shape[-1]). Source-text assertions, same read_text()
approach plugin_contracts/test_plugin_runtime_hook_contracts.py already uses
on this file. All three fail against the pre-fix source.

CI coverage 9 → 12. PR body updated.