Current limitation
RolloutManager._post_process_rewards can infer GRPO reward groups from the total sample count. When a custom rollout returns uneven numbers of samples per prompt, the fallback reshapes the rewards into one row and normalizes across the entire batch instead of within each prompt group.
For example, with group sizes 4, 3, and 4, a prompt whose three rewards are all 5.0 can receive advantages of -1.090909 instead of zero.
Expected behavior
Reward normalization should use each sample's explicit group_index and restore the normalized values in input order. If group identity is unavailable, uneven or mixed unidentified groups should not silently fall back to a global baseline. Singleton groups should produce zero when standard-deviation normalization is enabled.
Impact
Uneven custom rollout groups can receive incorrect advantages, changing the training signal.
Related pull request
Current limitation
RolloutManager._post_process_rewards can infer GRPO reward groups from the total sample count. When a custom rollout returns uneven numbers of samples per prompt, the fallback reshapes the rewards into one row and normalizes across the entire batch instead of within each prompt group.
For example, with group sizes 4, 3, and 4, a prompt whose three rewards are all 5.0 can receive advantages of -1.090909 instead of zero.
Expected behavior
Reward normalization should use each sample's explicit group_index and restore the normalized values in input order. If group identity is unavailable, uneven or mixed unidentified groups should not silently fall back to a global baseline. Singleton groups should produce zero when standard-deviation normalization is enabled.
Impact
Uneven custom rollout groups can receive incorrect advantages, changing the training signal.
Related pull request