Repository navigation
fix(reward): keep the grouped shape in reward_norm when only std is per group - #511
Open
flyingdoubleG wants to merge 1 commit into
Open
flyingdoubleG wants to merge 1 commit into
flyingdoubleG wants to merge 1 commit into
Conversation
…er group With norm_std_type="group" and norm_mean_type "batch", "running" or None, reward_norm computes the per-group std on the [n_groups, n_sample] view but subtracts the mean from the flat [N] rewards, so the division broadcasts to [n_groups, N] and the RLVR pipeline fails when it writes the result back to the batch. The agentic agentic_reward_norm supports this combination. Compute on the grouped view whenever a group statistic is used and restore the original shape afterwards; the other combinations are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
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
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.
What
reward_norm(roll/utils/functionals.py, used by the RLVRreward_postprocess) computes group statistics on the[n_groups, n_sample]view, but only uses that view whennorm_mean_type == "group":With
norm_std_type="group"andnorm_mean_type"batch","running"orNone, the flat[N]rewards are divided by a[n_groups, 1]std, which broadcasts to[n_groups, N]; writing that back to the batch then fails withRuntimeError: batch dimension mismatch, got self.batch_size=torch.Size([N]) and value.shape=torch.Size([n_groups, N]). Both fields accept these values inPPOConfig, and the agenticagentic_reward_normsupports the same combination (an agentic example config has it commented out:norm_mean_type: batch/norm_std_type: group).norm_mean_type,norm_std_typeChange
Use the grouped view whenever a group statistic is involved (
group_mode and n_sample > 0) and restore the original shape afterwards. The combinations that already worked give the same values.Tests
tests/utils/test_functionals.py::test_reward_norm_combinationschecks shape and values for six combinations against a direct computation; the two group-std cases fail before the change, all pass after (CPU).🤖 Generated with Claude Code