[Bug] GMPO geo_mean loss changes when a minibatch is split into more microbatches

System Info

verl: 24f25b03aa4b54249a273655ebbcce06f484192b
PyTorch: 2.11.0+cu130
Transformers: 5.5.3
Python: 3.12
GPU: NVIDIA GeForce RTX 5090, one visible GPU per process

Source: core_algos.py at 24f25b03.

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

The GMPO update changes when the same optimizer minibatch is split into more microbatches. The reproduction below goes through the native FSDP engine with one data-parallel rank.

Related: PR #5614, which discusses the same aggregation problem.

compute_policy_loss_geo_mean ends with pg_loss = torch.mean(pg_losses), a mean over the sequences in the current microbatch, and does not go through agg_loss or the global batch metadata. The FSDP loop then calls backward() on each microbatch loss with no rescaling, so a minibatch split into eight pieces contributes eight local means instead of one minibatch mean. Steps:

1. Go through the FSDP engine's train_batch → forward_backward_batch → ppo_loss → compute_policy_loss_geo_mean with eight fixed sequences (response length 16, 128 valid response tokens).
2. One SGD update, learning rate 0.02, gradient clipping 1.0, KL coefficient 0.001, same initial model and optimizer state.
3. Compare one microbatch of eight sequences with eight microbatches of one sequence.

Three repeated runs give the same result, and a second host reproduces the vectors exactly.

Expected behavior

Each microbatch should contribute its share of the minibatch mean. With one rank:

contribution = local_mean * local_sequence_count / minibatch_sequence_count

The denominator has to be the sequence count of the optimizer minibatch. Dividing by the number of microbatches only works when all microbatches have the same number of sequences.

A local weighting like this, applied without changing the eight-microbatch execution, brings the gradient and update differences down to numerical noise and agrees with an independent computation of the objective. The implementation also compensates for the engine's DP gradient averaging, but the numbers above are single-rank. This does not measure downstream reward and does not assume #5614 is part of the tested commit.