One rollout, several rows: the reward is copied onto every row, but counted once

Explainer for NVIDIA-NeMo/RL #4124 (paired with NVIDIA-NeMo/Gym #3368). RL links pin c617a2c.

Before
Without compaction, one rollout is one row, so its reward, its advantage and its row are the same thing.
This PR
With context compaction (CC), one rollout can become several rows. Every row carries the same reward, GRPO uses it once per rollout, and its one advantage is copied to every row.
Status
Proposed — open PR, not merged. The numbers below come from a real run of the PR's code, with a stand-in model.

Nine pages on this PR, one question each. This one: how one reward maps onto several rows.

Background. In GRPO each prompt is run several times; each run is a rollout, and the runs of one prompt form a group. A rollout's advantage is its reward minus the group's average reward (the baseline). Every token the model generated in that rollout is then pushed up or down by that advantage. Dotted words have a short explanation: hover or tap them.

Here the two rollouts of one group split differently: g0 was compacted 3 times and became 4 rows, while g1 was not compacted and stayed 1 row. The reassembler writes the rollout's single reward onto each of its rows (rollout_reassembler.py L1350-1376). Before GRPO runs, RL checks that all rows of a rollout agree, and keeps just one of them (logical_advantage.py L133-148). GRPO therefore sees 2 rewards, not 5. The resulting advantage is copied back to every row (single_controller.py L5378-5393).

Two rollouts, five rows, two rewards

① rows from the reassembler ② GRPO, once per rollout ③ advantage copied to every row rollout g0 · reward 1.0 · 7 generated tokens g0_s0 6 tok · reward 1.0 ← kept g0_s1 6 tok · reward 1.0 g0_s2 6 tok · reward 1.0 g0_s3 3 tok · reward 1.0 rollout g1 · reward 0.0 · 7 generated tokens g1_s0 21 tok · reward 0.0 ← kept same reward on each row: rollout_reassembler.py L1350-1376 rewards per rollout: [1.0, 0.0] baseline = mean = 0.5 g0: 1.0 − 0.5 = +0.5 g1: 0.0 − 0.5 = −0.5 one row per rollout: L133-148 g0_s0 [+0.5, +0.5] g0_s1 [+0.5, +0.5] g0_s2 [+0.5, +0.5] g0_s3 [+0.5] g1_s0 [−0.5 × 7] values on each row's generated tokens, as the learner received them copy to rows: single_controller.py L5378-5393 ④ loss weight: every generated token counts once, whatever row it is in The loss divides by the step's total number of generated tokens (here 7 + 7 = 14), not by the number of rows (5). megatron_policy_worker.py L2025-2030 Not what the code does — if GRPO counted rows instead of rollouts: rewards [1.0, 1.0, 1.0, 1.0, 0.0] → baseline 0.8 → g0 +0.2, g1 −0.8 (g0 counted 4 times because it was compacted 3 times) Token values come from the PR test's stand-in for vLLM; the rewards, row counts and advantages are what the real code produced.The baseline in this test is the plain group mean.

Why it is built this way

QuestionAnswer
Is the reward stored once or many times?Many: each row carries a copy, because rows are trained one at a time. RL refuses to train if the copies disagree.
Does a rollout with more rows get more say?No. GRPO sees one reward per rollout, so compaction cannot change any rollout's advantage. Above, g0 and g1 get +0.5 and −0.5 even though g0 has 4 rows and g1 has 1.
Does a rollout with more rows get more gradient?No more than it would without compaction. Each generated token gets its rollout's advantage, and the loss divides by the total token count. This is why CC requires a token-level loss and rejects sequence-level ones, which would divide by rows (config.py L858-879).
What if one row is bad?All rows of that rollout are masked. A rollout's validity is the lowest validity of its rows (logical_advantage.py L148).
What changes, then?Only the context each token is scored against. A token in row g0_s2 is scored against the compacted prompt that row starts with, which is the prompt the model actually saw when it generated that token.
So what. Compaction changes how many rows a rollout has, not how much it counts. The reward is one number per rollout. It is copied onto each row for storage, used once for the group baseline, and its advantage is spread over the rollout's generated tokens, wherever they landed. One metric does count rows: num_mask_sample_filtered counts each row and padding row (review comment on single_controller.py).