Exclude padding from RL trainer diagnostics, and fix three related metrics

[behroozazarkhalili]

Fixes the four issues reported in #6809. This PR primarily corrects reported metrics. It also changes PPO optimization for one edge case: masked_mean now returns 0 on an empty mask, so a micro-batch whose policy mask is empty produces a zero policy loss and zero gradients instead of NaN. Losses and gradients for every nonempty mask are unchanged.

Credit to @mmjerge, who found all four during a code audit.

The four

Two are padding problems and two are not, so they are worth separating.

Trainer	Metric	State on main
Online DPO	objective/kl, objective/non_score_reward, objective/rlhf_reward, objective/entropy	summed over the padded completion width
PPO	policy/approxkl_avg, policy/entropy_avg, policy/ratio_avg	bare .mean() while pg_loss beside them is masked
XPO	logps/chosen, logps/rejected	policy logps plus reference logps, a quantity with no interpretation
GRPO	policy_loss	computed on every loss branch, appended only under the entropy bonus

Every fix uses a mask or a value already in scope at that site, so nothing new is plumbed through. For XPO, Online DPO logs the policy term alone and Nash-MD logs the two separately; the sum matches neither, so this now logs the policy term.

What the padding actually contains

_forward returns selective_log_softmax over the full padded width and nothing zeroes it before the stat sites, so the padded entries hold log p(pad | context). That value rises toward zero as the policy learns to emit pad after EOS, which means the contamination shrinks as training proceeds rather than sitting at a fixed offset.

Holding the true masked objective/entropy fixed at 37.4808 and varying only p(pad|ctx):

 p(pad|ctx)     log p    reported (stale)    apparent trend
       0.01    -4.605             48.9938
       0.10    -2.303             43.2373           -5.7565
       0.50    -0.693             39.2137           -4.0236
       0.90    -0.105             37.7442           -1.4695
       0.99    -0.010             37.5060           -0.2383

The reported value falls by 11.5 while the quantity it measures never moves, so the metric shows a downward trend that nothing real is producing. The effect is largest early in training, when these curves get used to decide whether a run is worth continuing.

Those p(pad|ctx) values are plausible rather than measured from a live run; pinning the real distribution needs a GPU run this change does not otherwise require. The identity below does not depend on them.

Verification

• For any sum-based metric, reported minus correct equals the summand summed over the padded positions. Checked across 200 random fillings of the padded entries, agreeing to 1.8e-14, which is float64 summation rounding. It is zero exactly when every padded entry is zero, which no code path arranges.
• ruff@0.13.3 check and format clean on all four files, and the pinned doc-builder gate clean.
• Loss paths are untouched, traced rather than assumed. In Online DPO, logprobs reaches the loss only through cr_logprobs at line 1213, and line 1220 already multiplies that by ~cr_padding_mask, so the new mask is idempotent there. In PPO the changes are confined to the three *_stats buffers. XPO and GRPO change only what is appended to self.stats and self._metrics.
• Cross-model adversarial review caught a real defect in an earlier revision: a comment claiming policy_loss is logged before any regularizer. It is captured after per_token_loss = per_token_loss + self.beta * per_token_kl, so with beta != 0 it includes KL. The comment now states that. The metric's value is deliberately unchanged, since policy_loss = loss.detach() is pre-existing and this PR moves only where it is appended.

Note for the GRPO item

trl/experimental/gspo_token/grpo_trainer.py copies _compute_loss, and #6856 re-syncs that copy against main including this same gated append. Whichever lands second should carry the policy_loss change across, so the copies do not drift apart again.

Resolves #6809

---

Note

Medium Risk
Changes affect widely watched training metrics and PPO statistic reduction; only behavioral training change is zero policy loss on all-padding micro-batches, while actual optimization paths are otherwise unchanged per the PR.

Overview
Corrects logged metrics across Online DPO, PPO, XPO, and GRPO so training curves match what the loss paths actually optimize—without changing those loss computations in the common case.

Online DPO masks padded completion positions when aggregating objective/kl, objective/entropy, and related objective stats (aligned with the existing loss mask).

XPO reports logps/chosen and logps/rejected as policy log-probability sums only, instead of policy plus reference.

GRPO always logs policy_loss (including default entropy_coef=0 and Liger runs), with capture timing documented and matched per loss family (dapo/cispo/vespo vs others). grpo_trainer.md is updated accordingly.

PPO fixes several aggregation bugs: statistic buffers are sized to real micro-batches per minibatch (fixing ~1/num_mini_batches dilution), reinitialized each update with NaN slots skipped via nanmean, policy vs value masks split for zero-length responses, and approxkl/entropy/ratio use masked means. masked_mean on an empty policy mask now returns 0 (zero policy loss on that micro-batch instead of NaN).

New unit tests pin the expected metric values for these cases.

^{Reviewed by Cursor Bugbot for commit 4d77d03. Bugbot is set up for automated code reviews on this repo. Configure here.}

Update: two further metric defects

PPO statistics were scaled by 1 / num_mini_batches. The buffers in ppo_trainer.py took their last axis from gradient_accumulation_steps, while 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 were all under-reported by that factor. They were right only at the default num_mini_batches=1. The axis now comes from the same loop bounds: across the 86 configurations the exact-division checks admit, the old size is wrong in 50 and the new size in none.

The regression test pins this through ratio, which is exactly 1 while the policy has not moved. With num_mini_batches=2 it fails on the old code with assert 0.4999999701976776 == 1.0 ± 1.0e-04, and passes on the new code.

policy_loss was never logged for Liger runs. compute_loss sends them to compute_liger_loss and returns, so they never reached the append in _compute_loss, even though the comment there says the metric is logged for every run. The Liger path now records it at the point that matches the non-Liger capture: DAPO, CISPO and VESPO divide the normalizer in while building the loss and capture afterwards, and the other loss types capture before the accumulation rescale. The new test covers one loss type from each group.

Update: value diagnostics for zero-length responses, and metric-value oracles

Value statistics were gated by the policy mask. A zero-length response has no policy timestep, but its first value timestep is valid through padding_mask_p1, so loss/value_avg and val/clipfrac_avg were dropped for such a micro-batch. The two value slots now use the value mask; the policy slots keep the policy mask.

The tests now pin values rather than completion. The PPO all-padding test was not an oracle against the previous reduction: padded old and new log-probabilities share a sentinel, so the unmasked ratio is exactly 1 there too. It now controls the training logits, checks policy/entropy_avg against the valid-token entropy, and requires the empty row's value loss to reach loss/value_avg. Online DPO gains an exact oracle for objective/kl, objective/entropy, objective/non_score_reward and objective/rlhf_reward on hand-built completion rows whose padded values would dominate an unmasked sum. XPO gains an oracle showing logps/chosen and logps/rejected hold policy log-probabilities only.

On a compute node (job 58035371): the three oracles pass; restoring the shared policy gate for the value slots fails the PPO test; removing the KL mask in Online DPO fails its oracle; the full PPO, Online DPO and XPO files give 88 passed, 5 skipped; ruff 0.13.3 and the pinned doc-builder gate are clean.

Two related defects live elsewhere on purpose: #6856 re-syncs the gspo_token copy of _compute_loss (which carries the policy_loss append), #7033 adds policy_loss to GMPO, and #7010 covers gathering XPO per-example metrics before reduction.

[bot-ci-comment]

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.

[cursor]

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.}