Repository navigation
Exclude padding from RL trainer diagnostics, and fix three related metrics - #6861
Open
behroozazarkhalili wants to merge 17 commits into
Open
behroozazarkhalili wants to merge 17 commits into
behroozazarkhalili wants to merge 17 commits into
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Two of the four are padding problems and two are not. No loss or gradient path changes in any of them. Online DPO summed raw `(B, T)` logprobs for `objective/kl`, `objective/entropy` and, through `non_score_reward`, `objective/rlhf_reward`, two lines below a loss that already multiplies by `~cr_padding_mask`. PPO masked `pg_loss` and `pg_clipfrac` but left `policy/approxkl_avg`, `policy/entropy_avg` and `policy/ratio_avg` on a bare `.mean()`. Both now use the mask already in scope. The padded positions hold `log p(pad | context)`: `_forward` returns `selective_log_softmax` over the full padded width and nothing zeroes it before the stat sites. That value rises toward zero as the policy learns to emit pad after EOS, so the contamination shrinks during training instead of sitting at a fixed offset. Holding the true masked entropy fixed and moving only `p(pad|ctx)` from 0.01 to 0.99 slides the reported value from 48.99 to 37.51, a downward trend that nothing real is producing. XPO logged `logps/chosen` and `logps/rejected` as policy plus reference logps. Online DPO logs the policy term alone and Nash-MD logs the two separately, so the sum matches neither and has no interpretation. Now the policy term. GRPO computed `policy_loss` on every loss branch but appended it only inside the entropy-bonus block, so the metric vanished for standard runs. The append moves out. Its value is unchanged: it is captured after the KL term is folded in, so with `beta != 0` it carries that too, and excludes only the entropy bonus and the MoE auxiliary loss. Resolves #6809
behroozazarkhalili
force-pushed
the
fix/6809-rl-metrics-masking
branch
from
August 21, 2026 17:04
654abcd to
c31418c
Compare
The PPO statistics buffers took their last axis from gradient_accumulation_steps, but the micro-batch loop writes local_mini_batch_size // per_device_train_batch_size slots per minibatch and resets its index for each one. The unwritten slots stayed at zero, and every statistic is reduced with .mean(), so approxkl, pg_clipfrac, pg_loss, vf_loss, vf_clipfrac, entropy and ratio all came out scaled by 1 / num_mini_batches. They were right only at the default num_mini_batches=1. Taking the axis from the same loop bounds leaves no unwritten slot: over the 86 configurations the exact-division checks admit, the old size is wrong in 50 and the new size in none. The new test pins the scaling through ratio, which is exactly 1 while the policy has not moved yet. With two minibatches it reported 0.5 before this change. GRPO recorded policy_loss only in _compute_loss. compute_loss sends Liger runs to compute_liger_loss and returns, so a Liger run never reached that line, though the comment above it says the metric is logged for every run. The Liger path now records it where the non-Liger path does: DAPO, CISPO and VESPO divide the normalizer in as they build the loss and capture afterwards, and the remaining loss types capture before the accumulation rescale.
The doc-builder style hook reflows docstrings to a wider column than these were wrapped at, so it rewrote both files and failed the quality check. This is its own output.
This was referenced Sep 2, 2026
`masked_mean` divided by `mask.sum()` with no floor. A response whose first token is the pad token yields `sequence_length = -1`, so `~padding_mask` is empty, the divisor is zero, and the result is NaN where the bare `.mean()` it replaced returned a finite number. The statistics buffers are reduced with `.mean()`, so one such micro-batch turned `policy/approxkl_avg`, `policy/entropy_avg`, `val/ratio` and `val/ratio_var` into NaN for the entire update, which is the opposite of what moving those metrics to a masked mean was meant to achieve. Clamping the denominator reports 0 for the empty case and leaves every finite case bit-identical, matching the `.clamp(min=1.0)` normalization already used throughout the library. Rebuild the Liger `policy_loss` parity test so it can actually fail. The old configuration let the DAPO normalizer collapse to 1, where both branches of the capture-point ternary agree, and the loss it compared was structurally zero because GRPO advantages are centred within a group and the policy ratio is exactly 1 without `old_per_token_logps`. The test now offsets the fused loss by a constant, reads the accumulation count where `compute_liger_loss` reads it rather than after `train()` returns, and requires that at least one call had a normalizer other than 1. Both parametrizations pass clean and both fail against an inverted ternary. Correct three statements that were false as written: the PPO comment claimed every statistic is mean-reduced, but `val/ratio_var` uses `.var()`, where the padding zeros manufacture spread instead of scaling it; the Liger test comment said the accumulation count always falls to the epoch remainder on the last optimizer step, which does not hold when the dataloader length divides evenly; and the `policy_loss` doc entry did not say that `cispo`, `dapo` and `vespo` report a per-micro-batch value while the other loss types report the window value.
The previous commit clamped the `masked_mean` denominator so an all-padding micro-batch yields a loss of 0 instead of NaN. That kept the weights finite but had two side effects. The 0 it returns was then written into the statistics buffers, where the `.mean()` reductions counted it as an observation: with one empty and one real micro-batch, `val/ratio` read 0.5 for a real ratio of 1. Each buffer now starts as NaN, a slot is written only when the micro-batch holds at least one valid token, and the reductions use `nanmean`; `val/ratio_var` is the variance of the written slots. The loss path is unchanged, so the optimizer still sees 0 for an empty micro-batch and the accumulation cadence is untouched. `clamp(min=1.0)` also rescaled every mask whose sum was below one: a fractional mask of 0.5 on a value of 2.0 returned 1.0 where the plain division returned 2.0. Only an exactly zero denominator is replaced now, and the comment no longer claims the clamp matched the rest of the library, which normalizes both ways. The `policy_loss` note in grpo_trainer.md said the per-micro-batch value for cispo, dapo and vespo is smaller by `current_gradient_accumulation_steps / steps_per_generation`. The code divides by that ratio, and `gradient_accumulation_steps=1` with `steps_per_generation=2` is an accepted configuration where it is 0.5 and the value is larger. The note now says what the division is and when it shrinks the value. A regression test makes generation emit the pad token first for one row, which gives that row a sequence length of -1 and an all-padding micro-batch, and asserts `val/ratio` stays within 1e-4 of 1; on the previous commit it reads 0.5. The test pads with the tokenizer's own pad token: padding with EOS instead moved the real micro-batch's ratio by up to 2e-3 even without the forced row, while the distinct pad token keeps it within 1.2e-7. A second regression test covers the other half of the `policy_loss` change in this PR. On main the append sits inside `if self._entropy_bonus_enabled:`, so a run with the default `entropy_coef=0.0` logs no `policy_loss` at all; measured on main at 8397289, `entropy_coef=0.0` logs no key and `0.01` does. The two existing entropy tests both enable the bonus, so neither noticed. The new test trains one step with the default and asserts the key is logged; gating the append again makes it fail.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit d6da644. Configure here.
…r run The NaN-initialised buffers from the previous commit were allocated once before the update loop. A slot is written only for a micro-batch with a valid token, so an all-padding micro-batch in a later update left the value the previous update had written there, and `nanmean` folded that stale number into the current update's averages. Before the NaN change every slot was overwritten every update, so the single allocation was safe; it stopped being safe the moment slots could be skipped. The buffers are now created at the top of each update. A regression test makes every row of the second update all padding: that update has nothing to report, so its `val/ratio` must be NaN, while the stale buffers reported the first update's value.
# Conflicts: # tests/test_grpo_trainer.py
…e metric values The value loss and value clip fraction were gated by the policy mask, so a micro-batch whose only response is empty dropped its value statistics although its first value timestep is valid through padding_mask_p1. The two value slots now use the value mask. The PPO all-padding test was not an oracle: padded old and new log-probabilities share a sentinel, so the unmasked ratio is exactly one on the previous code too. It now checks the valid-token entropy against logits it controls and requires the empty row's value loss to be logged. Online DPO gains an exact oracle for the masked KL, entropy and reward metrics; XPO gains one showing logps/chosen and logps/rejected hold policy log-probabilities only.
PPOTrainer left main in #7020, so the PPO stats-buffer and masking changes and their tests have no target. The GRPO, GSPO-token, Online DPO and XPO changes are unaffected.
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
2 participants
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.
Online DPO included padding in metrics, XPO combined policy and reference log probabilities, and GRPO omitted
policy_losson some paths. I corrected those metrics and tested their values against independent calculations, with the selected trainer tests passing except for Liger tests blocked by missing GPU support. This overlaps #6856, so whichever merges second must carry the metric change into the GSPO-token copy.Resolves #6809.
Note
Low Risk
Changes affect diagnostics and documentation only; training losses are not modified, with tests pinning the corrected metric values.
Overview
Fixes misleading logged metrics in several RL trainers; optimization loss paths are unchanged.
Online DPO now masks padded completion tokens when computing
objective/kl,objective/entropy, and derived objectives (non_score_reward,rlhf_reward), matching the masked loss path instead of summing post-EOS padding log-probs.XPO
logps/chosenandlogps/rejectednow record policy log-probability sums only, not policy plus reference.GRPO (and the experimental GSPO-token copy) always appends
policy_lossto metrics—not only when an entropy bonus is enabled—and the Liger branch logs the same quantity at the same normalization point as the non-Liger path (with loss-type-specific scaling documented ingrpo_trainer.md).Tests add numeric oracles for Online DPO padding and XPO logps, plus GRPO/Liger
policy_losscoverage.Reviewed by Cursor Bugbot for commit a77d8c3. Bugbot is set up for automated code reviews on this repo. Configure here.