per_rollout_mean divides by training rows instead of rollouts, so the loss scale changes when a rollout spans several rows

[KunyangZhang]

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.5

Question

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_ids
  makes a rollout keep 1 / n_rollouts regardless of how it was split. I have that patch ready
  (per_rollout_loss.py, trainer.py, and the test expectation 1/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 switching trace_aggregator.level then silently
  rescales the loss.

Happy to open the PR for the first option, or rework it, once you confirm the intended semantics.

[KunyangZhang]

Found the authoritative definition in the v1.0 technical report (arXiv:2608.17528), §2.3
Loss Normalization, which docs/20-trainer-configuration.md points to. Eq. 16 there is
the rollout-level token-mean loss:

L_rollout-mean = (1/R) * sum_rho [ (sum_j sum_t l_rho,j,t) / (sum_j L_rho,j) ]

so each rollout contributes 1/R of the loss regardless of how many samples it produced,
and the report states the principle explicitly: "sample count should not be allowed to
affect gradient normalization, because it is often driven by incidental factors such as
retokenization". It also reports that in their own SWE-bench runs only 36% of rollouts
stay a single merged row, with 2.41 samples per rollout on average (Figure 10) — so this is
the common case, not an edge case.

Checked the current implementation against Eq. 16 using the report's own Figure 5 batch
(A: 50 + 100 tokens, B: 30 x 3, C: 40) with l = 1 per token and no clipping, so the batch
total is proportional to the post-normalization advantage mass:

	per-rollout mass	batch total	vs Eq. 16
current main	1/N = 1/6 each	0.5	R/N = 0.5
dividing by distinct rollout_ids	1/R = 1/3 each	1.0	matches Eq. 16

With the 2.41 samples/rollout reported in the paper, main would carry 1/2.41 ≈ 0.41 of
the intended gradient, and the factor moves with samples-per-rollout, which Eqs. 14-16 are
specifically written to avoid.

So I read this as the row-count divisor being an implementation slip rather than the
intended semantics — thanks for confirming, and I will open the PR for the
rollout_ids-based divisor if you agree. The patch is 3 files (per_rollout_loss.py,
trainer.py, tests/verl/test_per_rollout_loss.py), keeps the existing per-rollout
weighting, and the full suite passes (101 passed).

[hzy46]

Siwei Zhang (@SiweiPro) Would you please check this?

[SiweiPro]

Hi, Kunyang KunyangZhang,

Thanks for the detailed investigation. Eq. 16 in the technical report is indeed one source of the normalization, but the current implementation combines two normalization dimensions. Looking only at the sum of the normalized advantages captures the first one, but not the mini-batch update normalization.

Let

• $R$ = train_batch_size * $n$: the nominal number of rollouts in one training batch;
• $N$ = train_collected: the actual number of training samples/rows entering the actor update;
• $M$ = ppo_mini_batch_size: the actor mini-batch size.

The first normalization is the rollout-level token mean from the technical report:

$$ \frac{1}{R} \sum_{\rho} \left[ \frac{\sum_j \sum_t \ell_{\rho,j,t}} {\sum_j L_{\rho,j}} \right]. $$

Here, all rows belonging to the same rollout share the denominator $\sum_j L_{\rho,j}$. This ensures that splitting one rollout into multiple rows does not change its within-rollout token weighting.

There is also a second, implementation-level normalization arising from mini-batched actor training. The nominal number of actor mini-batches is

$$ K_{\mathrm{expected}} = \frac{R}{M}, $$

while the actual number is

$$ K_{\mathrm{actual}} = \frac{N}{M}. $$

Therefore, the actual-to-nominal amount of mini-batch gradient updates is

$$ \frac{K_{\mathrm{actual}}}{K_{\mathrm{expected}}}=\frac{N/M}{R/M}=\frac{N}{R}. $$

We compensate for this with its inverse:

$$ \frac{1}{N/R} = \frac{R}{N}. $$

Combining the report-level normalization and the mini-batch normalization gives

$$ \frac{1}{R}\cdot\frac{1}{N/R}=\frac{1}{N}. $$

This is why the implementation divides by train_collected $N$, rather than only by the number of distinct rollout_ids. For the Figure 5 example, $R=3$ and $N=6$. The observed normalized advantage mass of $R/N=0.5$ is the second correction factor; at the same time, the actor processes $N/R=2$ times as many mini-batch update units. These factors cancel at the training-update level.

Thanks again for the careful investigation and the detailed reproduction. Please feel free to follow up if you have any other questions or if any part of the normalization remains unclear.

[KunyangZhang]

Thanks Siwei — that explains it, and it checks out against the code. I verified the second factor
on main before replying:

• agentlightning/verl/trainer.py:502 — mini_bs = ppo_mini_batch_size * rollout.n
• agentlightning/verl/trainer.py:505 — n_remained_transition = n_transition // mini_bs * mini_bs

so the number of mini-batch updates in a step is driven by the collected rows N, not by the
nominal rollout count R, and max_ppo_update_times in
docs/20-trainer-configuration.md:207 ("limits the maximum number of PPO mini-batch updates
performed for one batch") exists precisely because more rows mean more update units. Measuring the
advantage mass alone shows the R/N correction factor but not the N/R extra update units it
cancels against — that was the gap in my analysis, and I was wrong to read the divisor as a slip.

I am closing #594 since it removes that compensation.

Thanks again for taking the time to work through it with the numbers — closing #594 now.

[KashmirAwana]

It looks like the issue is about how the loss scale changes when a rollout spans multiple rows. I'd check if the divisor should be based on the number of rollouts instead of rows.

[DeepanshuPal]

I reproduced this against current main at ff945758.

The behavior still comes from normalize_advantages_by_rollout() multiplying each rollout's token count by num_trained_rows, with trainer.py passing len(batch). That means a two-rollout batch keeps total mass 1.0 when each rollout is one row, but only 2/3 when the same two rollouts occupy three rows. The extra row changes the batch scale even though it does not add another rollout.

My reading is that the denominator should be the number of unique rollout IDs, not the number of training rows. That matches the documented purpose: preventing a rollout from getting more weight merely because it produced more rows.

Before changing the pinned test and training semantics: is unique rollout count the intended denominator here? If so, the focused fix is to derive it from rollout_ids, update the 1/3 expectation to 1/2, and add coverage showing the same per-rollout mass across 1/2/3-row splits.