Skip to content

fix: skip zero-variance GRPO updates - #1118

Merged
bradhilton merged 3 commits into
mainfrom
tarkin/grpo-zero-variance-advantages
Oct 10, 2026
Merged

bradhilton merged 3 commits into
mainfrom
tarkin/grpo-zero-variance-advantages

Conversation

@bradhilton

@bradhilton bradhilton commented Oct 9, 2026 •

Copy link
Copy Markdown
Collaborator

Agent: Tarkin

Fixes #1110.

Problem: GRPO advantages are (r − mean) / (std + 1e-6), and a trajectory is dropped only when its advantage is exactly 0. In a group with equal rewards, floating-point residue in the mean can leave a tiny nonzero advantage: [0.1, 0.1, 0.1] gives −1.4e-11. Those trajectories were kept. pack.py then divided by the batch-mean |advantage| with no epsilon. So a batch made only of such groups trained on |A| = 1 for every token, all with the same sign, and a batch of exact zeros produced NaN.

Behavior after this change:

  • Groups whose rewards agree within 1e-12 get an advantage of exactly 0. That's the same tolerance the pipeline trainer already used. The check now lives in one helper, art.utils.grpo.rewards_have_zero_variance, shared by tokenization, the pipeline trainer and the data/step_num_groups_trainable metric.
  • pack.py only normalizes weights and advantages when their means are positive, so it can no longer produce NaN or Inf.
  • If nothing trainable remains after packing, the step takes ART's existing skipped-step path: the checkpoint and step advance with no optimizer update. This applies on the Local/Unsloth, Tinker and Megatron backends, and consistently across ranks.
  • Groups with real reward variance get bit-identical advantages.

Validation:

  • New tests: tests/unit/test_grpo_zero_variance.py covers residue, exact-zero and near-equal groups, with scaling on and off and dropping on and off. It checks that normal and mixed batches are unchanged, that zero-weight denominators stay finite, and that the Local, Tinker and Megatron backends skip cleanly with unchanged optimizer state, including a four-rank no-op. 31 of its 50 cases fail on main.
  • Related suites: 180 passed and 1 skipped across the tokenize, packing, metrics, pipeline trainer and MoE routing suites.

Limitations:

@bradhilton
bradhilton force-pushed the tarkin/grpo-zero-variance-advantages branch 2 times, most recently from a234491 to ad8f435 Compare October 10, 2026 17:36
@bradhilton
bradhilton force-pushed the tarkin/grpo-zero-variance-advantages branch from ad8f435 to b66e9ff Compare October 10, 2026 18:58
@bradhilton
bradhilton marked this pull request as ready for review October 10, 2026 19:26
@bradhilton
bradhilton merged commit 9ff5329 into main Oct 10, 2026
12 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

GRPO: equal-reward groups can become a full-scale advantage signal or NaN after pack normalization

1 participant