[flyingdoubleG]

Description

Problem. With verl ≥ 0.8 (is_verl_legacy() is False; PyPI currently ships 0.9.1), the actor is trained through verl's engine API using TrinityPolicyLoss. The engine sums the micro-batch losses:

• FSDP: forward_backward_batch calls loss.backward() for every micro-batch, with no averaging.
• Megatron: the engine multiplies by num_micro_batch to undo Megatron's per-micro-batch division.

Normalization is therefore left to the loss function. The engine attaches batch_num_tokens (global across DP), dp_size and global_batch_size to each micro-batch, and verl's own ppo_loss uses them to return each micro-batch's share of the global mini-batch loss. In verl's words (engine_workers.py): "we should always sum the loss of each micro-batch as we scale by global_bsz/global_token".

TrinityPolicyLoss returns a loss normalized within the micro-batch: aggregate_loss is a local token or sequence mean. As a result the actor gradient is multiplied by the number of micro-batches per rank.

Under the default dynamic batching (use_dynamic_bsz with max_token_len_per_gpu), that number changes from step to step with response length, so steps with longer responses get proportionally larger gradients. Gradient clipping and the logged grad_norm see the inflated values too.

The legacy path (verl_legacy/dp_actor.py) does not have this problem: it scales each micro-batch by 1/gradient_accumulation, or by its token share. The engine path dropped that scaling.

Reproduction on CPU: the real TrinityPolicyLoss + verl's left_right_2_no_padding, summing micro-batch losses as the engine does. Gradient norm relative to a single micro-batch:

loss_agg_mode	2 micro-batches	4 micro-batches
token-mean, main	2.017	4.027
seq-mean-token-mean, main	2.000	4.000
both, this PR	1.000	1.000

Fix. TrinityPolicyLoss now multiplies the combined loss (pg − entropy + kl) by the micro-batch's share of the global mini-batch, times dp_size, following the same contract as verl's ppo_loss:

• token-mean: local response tokens / batch_num_tokens.
• seq-mean-token-mean, seq-mean-token-sum, and losses without loss_agg_mode (DPO): local sequences / global_batch_size. Trinity's _update_actor already sets global_batch_size.
• seq-mean-token-sum-norm: this is already a sum, so only dp_size is applied, to undo the DP averaging.

When the engine does not provide these values (e.g. outside the engine), the weight is 1, so behaviour there is unchanged. Logged loss metrics are unchanged as well, because they use the unscaled value.

How to test. New tests/trainer/verl_losses_test.py builds micro-batches the way the engine does. It checks that the gradient of the summed losses with 2 and 4 micro-batches equals the single-micro-batch gradient, for token-mean and seq-mean-token-mean. On main it fails with 13/40 elements mismatched (the rest are padding with zero gradient); with this change it passes (1 test, 4 subtests, CPU). I have not run a GPU training job.

Checklist

• Code has passed all tests (new test on CPU; GPU tests not run)
• Docstrings have been added/updated in Google Style
• Documentation has been updated (no user-facing change)
• Code is ready for review

🤖 Generated with Claude Code