docs: state the trained-batch divisor in per-rollout mean loss - #602
hiro-nikaitou wants to merge 2 commits into
Conversation
`per_rollout_mean` divides every training row by `rollout_token_count * num_trained_rows`. Issue microsoft#593 read the second factor as an implementation slip; the thread confirms the row-count divisor is intentional, because it is the inverse of the extra PPO mini-batch update units that the same rows create, so the two cancel at the update level. State that where users configure the loss mode, and spell out both operands of the divisor in the normalize_advantages_by_rollout docstring. No behaviour changes. Signed-off-by: hiro-nikaitou <vieteviete@proton.me>
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Replace the unsupported cancellation claim with the directly implemented normalized-mass invariant.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 1
What changed in this PR
Documents the trained-row divisor used by per-rollout mean loss without changing behavior.
Changes:
- Clarifies loss normalization in trainer documentation.
- Expands the normalization function docstring.
| File | Description |
|---|---|
docs/20-trainer-configuration.md |
Documents divisors, but its update-level cancellation claim is not guaranteed by the training loop. |
agentlightning/verl/per_rollout_loss.py |
Clarifies both normalization divisor operands. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
|
||
| `actor_rollout_ref.actor.policy_loss.loss_mode: per_rollout_mean` normalizes the policy loss at the rollout level. It prevents a rollout from receiving more optimization weight only because it produced more training rows. | ||
|
|
||
| Rows that belong to the same rollout share one token denominator, and each row is additionally divided by the number of rows in the trained batch. That second divisor is the inverse of the extra PPO mini-batch update units those rows produce, so the two cancel at the update level: splitting one rollout across more rows, for example by switching `agentlightning.trace_aggregator.level`, adds proportional update units without changing the rollout's effective optimization weight. |
Each retained row is consumed once per PPO epoch, and rows are floored to complete mini-batches and can be dropped by `max_ppo_update_times`, so extra rows do not create proportional update units. State the invariant the code implements instead of the cancellation it does not guarantee. Signed-off-by: hiro-nikaitou <vieteviete@proton.me>
|
Rewritten in 875f838 so the section states the invariant the trainer implements instead of the cancellation it does not guarantee. Each retained rollout now reads as carrying a normalized token mass of |
Summary
per_rollout_meandivides every training row byrollout_token_count * num_trained_rows. #593 read the second factor as an implementation slip and asked whether the divisor should be the number of distinct rollouts instead. The thread confirms the row-count divisor is intentional: it is the inverse of the extra PPO mini-batch update units that the same rows create, and the two cancel at the update level.This change writes that down where users configure the loss mode. No behaviour changes.
docs/20-trainer-configuration.md: state both divisors and thetrace_aggregator.levelconsequence in Per-rollout mean loss.agentlightning/verl/per_rollout_loss.py: spell out both operands of the divisor in thenormalize_advantages_by_rolloutdocstring.Two different reporters have already read the row-based divisor as a bug, so the intended semantics are worth stating explicitly. This only documents semantics that are already pinned by
tests/verl/test_per_rollout_loss.py.Closes #593.
Validation
uv run --locked --no-sync pytest -q tests/verl/test_per_rollout_loss.py->4 passeduv run --locked --no-sync ruff check .->All checks passed!;ruff format --check agentlightning/verl/per_rollout_loss.py->1 file already formatteduv run --locked --no-sync python scripts/check_headers.py-> passedgit diff --check-> no whitespace errorsuv sync --frozen --no-default-groups --group docs+uv run --locked --no-sync mkdocs build --strict->Documentation built in 17.44 seconds, exit 0num_trained_rows=len(batch)intrainer.py, therollout_token_counts[rollout_id] * num_trained_rowsdivisor inper_rollout_loss.py, and the presence of both doc statements. 3/5 checks passed before this change, 5/5 after.normalize_advantages_by_rolloutand the mini-batch computation copied fromtrainer.py:518-523: two rollouts over two rows -> per-rollout mass{0.5, 0.5}, batch mass1.0, update units2; the same two rollouts over three rows ->{0.3333, 0.3333}, batch mass0.6667, update units3.mass * unitsstays2.0in both cases.Not run: GPU rollouts / real training (no GPU here); the numbers above come from the adapter's own normalization function, not from an end-to-end loss.
This patch was prepared with AI assistance. I have reviewed every line and can explain it.