[flyingdoubleG]

What

agg_loss(..., loss_agg_mode="token-mean") sums loss_mat over every position and divides by the masked token count:

loss = (loss_mat * weights.unsqueeze(-1)).sum() / batch_num_tokens

loss_mask is not applied to the numerator, so tokens the caller masked out still add to the loss and receive gradient. The other modes use masked_sum / masked_mean and are not affected. This came in with v0.2.0 (3077bef), where masked_mean(loss_mat * weights, loss_mask) was replaced to support the global batch_num_tokens; the mask was dropped in that change.

Where it matters

• AgenticActorPGWorker.loss_func and the RLVR ActorPGWorker.loss_func apply the train-infer correction filter as response_mask * filter_mask, then aggregate the policy-gradient loss with agg_loss(pg_loss, response_mask, loss_agg_mode, ...). The IS weight is not zero on filtered tokens and their advantages are not masked, so with loss_agg_mode: token-mean the filtered tokens are still trained on. Five example configs use token-mean (e.g. examples/qwen3-8B-gsm8k-math-rlvr_megatron/rlvr_config.yaml).
• Entropy loss, and the clip-fraction / filter pass-rate metrics computed with token-mean, also count masked positions.

Minimal check with the function from roll/utils/functionals.py (one token filtered, one padding token, loss 1 everywhere):

mode	loss	grad on masked tokens
token-mean (before)	1.3333	0.1667, 0.1667
token-mean (after)	1.0000	0, 0
seq-mean-token-mean	1.0000	0, 0

Change

loss = masked_sum(loss_mat * weights.unsqueeze(-1), loss_mask) / batch_num_tokens

batch_num_tokens keeps its meaning (global valid-token count), so nothing changes when loss_mat is already zero outside the mask.

Tests

tests/utils/test_functionals.py: masked tokens change neither the value nor the gradient for token-mean, seq-mean-token-sum and seq-mean-token-mean; token-mean with weights divides the local masked sum by the global token count. Both token-mean cases fail before the change and pass after; the file passes (8 tests, CPU).

A related question (not changed here)

seq-mean-token-sum-norm returns sum(seq_losses) / loss_mask.shape[-1] without dividing by the number of sequences, so its scale grows with the batch size, unlike the other seq-mean-* modes. This matches the verl version cited in the docstring; verl later added the sequence mean (verl-project/verl#5417). If you'd like the same here I can send a separate PR, since it changes the loss scale for existing configs.

🤖 Generated with Claude Code

[CLAassistant]

All committers have signed the CLA.