PPOTrainer statistics average per-micro-batch means without weighting by micro-batch size

[behroozazarkhalili]

Summary

PPOTrainer reduces its per-micro-batch statistics with an unweighted mean over slots. Each slot holds a masked_mean of its own micro-batch, so when the last micro-batch of a minibatch is short, it carries the same weight as a full one and the reported metric is not the token or sample mean it is taken to be.

Where

trl/experimental/ppo/ppo_trainer.py. The micro-batch loop stores one masked mean per slot, and the reduction at 897-904 averages those slot means:

metrics["policy/approxkl_avg"]  = self.accelerator.gather_for_metrics(approxkl_stats).mean().item()
metrics["policy/clipfrac_avg"]  = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item()
metrics["loss/policy_avg"]      = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item()
metrics["loss/value_avg"]       = self.accelerator.gather_for_metrics(vf_loss_stats).mean().item()
metrics["val/clipfrac_avg"]     = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item()
metrics["policy/entropy_avg"]   = self.accelerator.gather_for_metrics(entropy_stats).mean().item()
metrics["val/ratio"]            = self.accelerator.gather_for_metrics(ratio_stats).mean().item()

A mean of per-group means equals the pooled mean only when the groups are equal in size.

When the micro-batches are uneven

exact_div requires num_mini_batches to divide per_device_train_batch_size * gradient_accumulation_steps. It does not require per_device_train_batch_size to divide local_mini_batch_size, so a short final micro-batch is reachable from a valid config:

from trl.experimental.ppo import PPOConfig

args = PPOConfig(
    output_dir="/tmp/ppo-uneven", per_device_train_batch_size=4,
    gradient_accumulation_steps=3, num_mini_batches=2, num_ppo_epochs=1,
    use_cpu=True, bf16=False, report_to="none",
)
local_batch = args.per_device_train_batch_size * args.gradient_accumulation_steps
local_mini_batch = local_batch // args.num_mini_batches
sizes = [min(args.per_device_train_batch_size, local_mini_batch - s)
         for s in range(0, local_mini_batch, args.per_device_train_batch_size)]
print(local_batch, local_mini_batch, sizes)

Output: 12 6 [4, 2]. The config is accepted, and the two micro-batches of each minibatch hold 4 and 2 samples.

Size of the error

With slot values 1.0 and 3.0 over micro-batches of 4 and 2 samples, computed with exact rationals:

quantity	value
what the code reports, mean(1, 3)	2
size-weighted mean, (4*1 + 2*3) / 6	5/3
relative error	1/5, that is 20%

The error is exactly zero when every micro-batch is full, which is why the default num_mini_batches=1 never shows it.

Not the same as the zero-slot bug

stats_shape used to size its last axis by gradient_accumulation_steps while the loop writes ceil(local_mini_batch_size / per_device_train_batch_size) slots, leaving unwritten zeros in the average. #6861 corrects that axis. The weighting described here is independent: it is present both before and after that change, since both versions take an unweighted mean over whatever slots exist.

Possible direction

Store a numerator and an active count per slot and reduce as sum / count for the token-level metrics, and weight the sample-level metrics by the real micro-batch size. val/ratio_var needs its own decision, since a variance over slot means is not a variance over samples either.