per_rollout_mean divides by training rows instead of rollouts, so the loss scale changes when a rollout spans several rows
per_rollout_mean is the default loss mode (agentlightning/verl/config.yaml:39). In trainer.py:617-626
the advantages are rescaled once for the whole batch before the actor update:
normalize_advantages_by_rollout(
batch.batch["advantages"],
batch.batch["response_mask"],
rollout_ids,
num_trained_rows=len(batch), # rows, not rollouts
)normalize_advantages_by_rollout (per_rollout_loss.py:16-41) then divides each row by
rollout_token_count * num_trained_rows.
That is exactly the per-rollout mean when a rollout occupies a single row. But a rollout can span
several rows: trace_aggregator.level: transition emits one row per turn, and even the default
trajectory level flushes an extra row when a turn's prompt does not extend the running context
(rollout_adapter.py:542-571, reported as training/n_trace_merge_mismatch_rows). In that case the
divisor counts rows, so every rollout keeps 1 / n_rows instead of 1 / n_rollouts and the batch
total becomes n_rollouts / n_rows.
Reproduction
Calling the real adapter and then the same normalization trainer.py performs, on main @ 218f1f7:
[A] 2 rollouts, 1 row each -> per-rollout mass {r1: 0.5, r2: 0.5}, total = 1.0000
[B] 2 rollouts, 3 rows total -> per-rollout mass {r1: 0.3333, r2: 0.3333}, total = 0.6667
(r1 hit a trace-merge mismatch; the adapter reports n_trace_merge_mismatch_rows = 1)Case B differs from A only by a trace-merge mismatch, yet the whole batch carries 2/3 of the signal. The same effect appears in isolation:
normalize_advantages_by_rollout(ones(4, 3), ones(4, 3, dtype=long), ["A", "A", "A", "B"],
num_trained_rows=4)
# A mass 0.25, B mass 0.25, total 0.5Question
Is the divisor the number of rollouts or the number of rows? The current test
tests/verl/test_per_rollout_loss.py::test_normalize_advantages_by_rollout pins the row-based value
(2 rollouts over 3 rows keep 1/3 each), so I did not want to assume it is a bug.
- If it is meant to be the per-rollout mean: deriving the divisor from the distinct
rollout_idsmakes a rollout keep1 / n_rolloutsregardless of how it was split. I have that patch ready (per_rollout_loss.py,trainer.py, and the test expectation1/3→1/2, plus a regression test over 1/2/3 rows per rollout). Full suite passes: 101 passed. - If row-based scaling is intended: the docstring ("
rollout's token count and batch size") and the test would be worth stating explicitly, since switchingtrace_aggregator.levelthen silently rescales the loss.
Happy to open the PR for the first option, or rework it, once you confirm the intended semantics.
Source: microsoft/agent-lightning