How a context-compaction step fits a variable number of rows into one batch

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

Before
train_global_batch_size = prompts × generations, and each rollout is exactly one row.
This PR
Under context compaction (CC), the batch size still counts rollouts. The number of rows per step varies, and padding keeps it even across GPUs.
Status
Explains the PR's code — no change proposed. The example numbers are worked by hand from the code's formulas, not taken from a run.

Nine pages on this PR, one question each. This one: what "batch size" means when one rollout can be several rows.

The config check is unchanged: prompts × generations must equal train_global_batch_size (config.py L1355-1363). Under CC, that number counts rollouts, not rows. A step closes when it holds exactly that many rollouts, however many rows they produced, and raises otherwise (single_controller.py L3206-3213). Two things make a variable row count work. Each group is padded to an even multiple, so every GPU runs the same number of micro batches. And gradients are added up across all micro batches and divided once, by the total token count, at the end of the step.

Ordinary GRPO: one row per rollout g0·r0 g0·r1 g1·r0 g1·r1 g2·r0 g2·r1 3 prompts × 2 generations = 6 rollouts = 6 rows = train_global_batch_size Context compaction: each rollout gives one row per segment, so the count varies g0 r0·s1 r1·s1 r1·s2 r1·s3 1 + 3 = 4 rows, +0 pad → 4 g1 r0·s1 r0·s2 r1·s1 r1·s2 r1·s3 r1·s4 r1·s5 pad 2 + 5 = 7 rows, +1 pad → 8 g2 r0·s1 r1·s1 pad pad 1 + 1 = 2 rows, +2 pad → 4 Each group is padded up to a multiple of Q = 4 GPUs × lcm(micro batch 1, logprob batch 1) = 4. A padding row copies a real row's input, has reward 0, and is masked. rollout_reassembler.py L853-882 config.py L882-908 (Q) 13 real rows + 3 pad = 16 rows for 6 rollouts Each group is split evenly across the 4 GPUs, so every GPU runs the same number of micro batches GPU 0 r0·s1 g0 · micro batch 1 r0·s1 g1 · micro batch 2 r0·s2 g1 · micro batch 3 r0·s1 g2 · micro batch 4 GPU 1 r1·s1 g0 · micro batch 1 r1·s1 g1 · micro batch 2 r1·s2 g1 · micro batch 3 r1·s1 g2 · micro batch 4 GPU 2 r1·s2 g0 · micro batch 1 r1·s3 g1 · micro batch 2 r1·s4 g1 · micro batch 3 pad g2 · micro batch 4 GPU 3 r1·s3 g0 · micro batch 1 r1·s5 g1 · micro batch 2 pad g1 · micro batch 3 pad g2 · micro batch 4 Which row goes to which GPU is drawn for illustration. The code only guarantees that every GPU gets the same number. Gradients from all 16 micro batches add up, then are divided once by N, the valid tokens summed over every GPU (padding counts 0). Then one optimizer step, however many rows there were. megatron_policy_worker.py L2018-2036 single_controller.py L3206-3213 (6 rollouts or raise)

Why the row count does not change the math

  1. Loss: each micro batch adds its summed token loss to the gradients. finish_train_step sums the valid tokens over every GPU and scales the gradients by 1/N (L2018-2036). Each generated token is trained once, whether its rollout sits in 1 row or 3. Padding rows have no valid tokens, so they add nothing.
  2. Advantage: computed once per rollout and copied to all its rows (page 5). Padding rows get 0.
  3. Optimizer: one step per training step, as before. More rows means more micro batches, not more optimizer steps.

What it costs

CostWhy
Megatron only, fixed batch: no sequence packing, no dynamic batching, no interleaved pipelinePadding to a fixed multiple only works when every micro batch has a fixed number of rows. The config check rejects the other modes (config.py L891-900).
Up to Q − 1 padding rows per groupHere 3 of 16 rows (19%) are padding. More GPUs or a larger micro batch raise Q, and with it the waste.
Step time grows with compaction6 rollouts cost 16 forward/backward rows here, against 6 in ordinary GRPO. Memory per micro batch stays the same.
Row-count metrics change meaningglobal_valid_seqs counts rows, so under CC it can exceed the number of rollouts. The proposed L1 test uses this to check that rows were split.
So what. You get no more than train_global_batch_size rollouts per step, and no fewer. Extra compaction rows become extra micro batches inside the same optimizer step, normalized by tokens. This is also why a dropped group cannot simply be skipped (page 6): the step insists on the full rollout count.