[behroozazarkhalili]

Fixes #7012.

What

PPOTrainer reduced its per-update statistics with an unweighted mean over micro-batch slots. Each slot holds a masked_mean (or plain mean) of its own micro-batch, so with a short final micro-batch the reported number was a mean of means, not the token or sample mean it is documented as. Unwritten slots, when stats_shape is larger than the loop writes, also contributed zeros.

The seven means now store weighted numerators and reduce as total numerator over total count:

• loss/policy_avg, policy/clipfrac_avg weight by the active response-token count of the micro-batch.
• loss/value_avg, val/clipfrac_avg weight by the active value-token count (padding_mask_p1).
• policy/approxkl_avg, policy/entropy_avg, val/ratio are per-token statistics: each slot now stores a masked_mean over the response tokens times the response-token count, and the reduction divides by the gathered response-token count. Before, the plain .mean() in those three slots included pad positions, and the first version of this PR only reweighted that mean by micro-batch size (corrected in 42c71bf).
• val/ratio_var keeps its variance over slot means, on its own buffer, because the issue leaves that population undefined, but now only over the slots the loop wrote: the response-token count buffer marks them, and the unwritten zeros used to inflate it (with learning_rate=0 under a 3-over-2 split every ratio is 1 and the logged variance was 0.2667 instead of 0).

Test

test_uneven_micro_batch_stats_are_pooled runs one CPU update with per_device_train_batch_size=4, gradient_accumulation_steps=3, num_mini_batches=2, which splits each minibatch into micro-batches of 4 and 2 samples. It spies on the module-level masked_mean and selective_log_softmax to record the live tensors entering each statistic, pools them independently, and compares with the logged metrics.

Verification

On a compute node (job 58032838):

• fixed tree: the test passes, and the full tests/experimental/test_ppo_trainer.py gives 57 passed, 2 skipped;
• base ppo_trainer.py: the test fails at the first pooled comparison;
• mutant dividing loss/policy_avg by the sample count instead of the token count: fails;
• mutant with the entropy weight off by one: fails;
• ruff 0.13.3 check and format: clean.

test_ratio_var_ignores_unwritten_stats trains with learning_rate=0, gradient_accumulation_steps=3 and num_mini_batches=2 and asserts val/ratio_var is 0; on the previous commit it reads 0.2667.

Limits

• Multi-process aggregation was not run here. The count and numerator buffers go through the same gather_for_metrics path as before.
• This overlaps with Exclude padding from RL trainer diagnostics, and fix three related metrics #6861, which changes the same stat buffers (per-update allocation, NaN for unwritten slots). Whichever merges second needs a small conflict resolution; the count-based reduction here makes the unwritten-slot handling unnecessary for the seven means.

---

Note

Medium Risk
Changes how PPO training metrics are reduced across micro-batches, which affects monitoring and any logic keyed on those logs; training loss/backprop paths are unchanged but reported KL/entropy/ratio values will differ when batching is uneven.

Overview
PPO logged averages no longer take an unweighted mean over micro-batch slots (which skewed uneven 4 vs 2 splits and let empty slots pull metrics toward zero). Each step now accumulates weighted numerators—policy-side stats (loss/policy_avg, clip fractions, approx KL, entropy, val/ratio) scale by active response-token count; value-side stats use padding_mask_p1 counts—and divides by the gathered totals after gather_for_metrics.

Per-micro-batch computation also aligns with masking: approximate KL and entropy use masked_mean on the response mask instead of raw .mean(), and ratio logging uses masked means.

val/ratio_var is still variance across per-slot mean ratios, but only slots with padding_count > 0, so unwritten buffer zeros no longer inflate variance (e.g. LR=0 case).

Tests add test_uneven_micro_batch_stats_are_pooled (monkeypatched masked_mean / selective_log_softmax vs logged metrics) and test_ratio_var_ignores_unwritten_stats.

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

[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 c054088. Configure here.}

[behroozazarkhalili]

Closing: PPOTrainer was removed from the tree in #7020 (main 700b845), so the stats-buffer weighting fix here has no target. Nothing else in this PR applied outside trl/experimental/ppo. Issue #7012 is closed for the same reason.