[Bug] KL-Cov token selection changes with the PPO microbatch size

Summary

KL-Cov selects the penalized tokens separately within each microbatch. Changing the microbatch size therefore changes which tokens are selected and, in turn, the parameter update, even when the PPO minibatch, rollout data, and initial model state are identical.

A batch setting that should only affect memory ends up changing the effective loss.

System Info

Environment: veRL commit 24f25b03, native FSDP engine, one visible GPU per run.

The results below are from the pinned version. Source inspection on 2026-09-10 found the same calculation in 1252cc71; the full experiment has not been rerun on that commit.

The pinned experiment uses the native FSDP actor engine with one visible GPU per run. Its relevant configuration is included below. A full python scripts/diagnose.py output from that GPU run is not attached; the later source inspection is not a new GPU execution.

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 test freezes 64 response sequences of length 128 (8,192 valid response tokens) together with their advantages, old log probabilities, the initial model state, and the optimizer settings. It compares one microbatch of 8,192 tokens against eight dynamic microbatches of 1,024 tokens. Both perform one SGD update with learning rate 0.02, clipping norm 0.05, and kl_cov_ratio=0.0002.

The current path already produces different selected-token masks before backward().

The discrepancy repeats exactly across three blocks. The repair brings the difference down to the level of the penalty-disabled control. Reproduced on a second host.

The quota difference can be checked directly from the formula in compute_policy_loss_kl_cov:

ratio = 0.0002
full_minibatch = max(1, int(8192 * ratio))        # 1 selected token
eight_microbatches = 8 * max(1, int(1024 * ratio))  # 8 selected tokens
assert (full_minibatch, eight_microbatches) == (1, 8)

This snippet isolates the quota calculation; it does not reproduce the full FSDP update or the covariance ordering measured above.

Expected behavior

For the fixed PPO minibatch and frozen inputs above, changing only the microbatch size should preserve the selected-token mask and the resulting objective. The selection quota and covariance statistics should be evaluated over that minibatch.

Root cause

compute_policy_loss_kl_cov computes the advantage and log-probability means and the top-k covariance selection on whatever tensors each loss call receives, and those calls receive individual microbatches.

The quota is max(1, int(valid_token_count * kl_cov_ratio)). For these sizes it selects one token for the full minibatch, but one token per microbatch, eight in total, when split. The means and the candidate set change as well, so adjusting the quota alone does not recover the full-minibatch selection.

Proposed fix

Compute the covariance statistics and the selected-token mask once over the optimizer minibatch, then pass each microbatch its slice of the mask.

The tested repair runs a forward-only prepass to obtain the current log probabilities, builds the mask, and keeps the original microbatch loop for backward. This costs one extra forward pass. The reference implementation was tested with entropy_coeff=0 and use_kl_loss=False; an upstream version must keep the other supported loss terms intact.

Related pull request

#7830.