[AmineDiro]

AsyncGRPO: pack as arrays

What this changes

TreeRow.pack used to build the sampling mask as a dense (N, width) tensor: one row per packed position, padded to the widest candidate set in the row, built from a nested Python list with torch.tensor. With top-p 0.95 and top-k 512 a position holds about 3 candidates on average, but any row of a few thousand positions contains one near the top-k cap, so the table was more than 98% padding and torch.tensor walked every padded cell on rank 0, before the broadcast and before the GPUs could start.

Now pack returns the candidate ids laid end to end in position order (sampling_mask, length T) and the size of each position's set (sampling_mask_counts, length N). compute_loss scatters them into the padded (N, width) table on the device with two index ops. The fused head and the kernel are unchanged: they already turned that table into a tile indicator with a scatter.

The sampler top-K log-probs (score centering) followed the same path: a dict per token from the worker, two nested comprehensions in pack. The worker now hands them over as two (len(input_ids), K) arrays padded with -1 and -inf, the layout the async distillation worker already uses for its teacher candidates, and pack fills the row's table with one indexed copy per sample. Nothing read those dicts by key; they were transport only.

Why it matters

On the sanity bed (R1-Distill-Qwen-1.5B, 128 completions of up to 8k tokens per step, two trainer GPUs, two vLLM GPUs), a py-spy profile of the previous version put 499 s of the 25-step run in rank 0's collate, 386 s of it in pack, against 12.8 s of forward and backward per step. Rank 1 spent the same time blocked in broadcast_object_list waiting for the batch. Both GPUs idled for about 20 s of every 50 s step, and the sampler meanwhile produced 2.5x more completions than the trainer consumed, dropped as stale.

Measurements

Packing one row, CPU

Synthetic samples in the trainer's queue format, two harness shapes, 8 micro-batches of 2 rows x 8 samples each. "main" is main (ragged mask only applies there), "before" is the score-centering branch this PR sits on. Generated by bench-packing/gen_samples.py, timed by bench-packing/charts.py on a login node.

shape	packed tokens per row	main	before	after	batch bytes per row, before -> after
agent: 6k-token shared prompt, 400-token response	9.2k	256 ms	282 ms	15 ms	38 MB -> 3 MB
multi-turn: 6 turns, thinking dropped after 2	28k	881 ms	1040 ms	66 ms	122 MB -> 9 MB

Of the 1040 ms before on the multi-turn row, 784 ms was the one sampling_mask line and 149 ms the two sampler tables. After, the remaining 66 ms is the DFS layout and the per-token label loop.

The per-sample pickle through the rollout queue goes from 1.2 ms dumps and 2.5 ms loads to 0.5 ms each for a 4.2k-token multi-turn sample; the sample is 1.1 MB instead of 0.5 MB because context positions now carry padding rows, which the trainer skips.

Training step, sanity bed

Two 15-step runs on the same bed and base commit, each under py-spy: prof4-before-sc-mis (branch head) and prof4-after-sc-mis (this change), both on the trackio Space aminediroHF/trl-sanity. Seconds of active Python stacks on rank 0 over the run, from the dumps (bench-packing/profile_charts.py):

rank 0, whole run	before	after
TreeRow.pack	250 s	31 s
trie add	38 s	36 s
collator pad and tensors	53 s	21 s
forward	52 s	53 s
compute_loss	56 s	52 s

Per optimizer step that is 22.7 s of collate on rank 0 before and 5.9 s after, while rank 1 waits in broadcast_object_list. The rollout worker's queue pickling went from 23 s to 6 s; its JSON decode and per-token conversion scale with the number of responses and are unchanged per response.

On these hosts the vLLM pause at weight sync stalled many steps for one to eight minutes, in both runs alike and unrelated to this change (the earlier profiled run on October 16 paused in 0.1 to 17 s). The clean comparison is the steps where the pause took under 2 s:

per step, quiet steps	before (2 steps)	after (5 steps)
step time	37.0 s	22.1 s
forward + backward	13.2 s	13.1 s
rollout queue wait	1.7 s	1.5 s

So the step goes from 37 s to 22 s with the same compute, and the trainer-side overhead outside forward and backward from about 22 s to 7 s. Over all steps, stalls included, trained tokens per second of wall clock went from 9.2k to 21.1k and samples dropped as stale from 400 to 178 per step, but those two numbers carry the pause noise.

The py-spy flamegraphs of the trainer ranks (width is the share of active samples; blue is pack, light gray the wait for the batch, green the weight sync wait on rank 1):

---

Note

Medium Risk
Changes how truncated-sampling batches are tensorized and how compute_loss builds support_ids, though parity tests cover loss/gradients; main risk is subtle shape/index bugs in the device-side scatter.

Overview
AsyncGRPO no longer builds top-k/top-p sampling masks as a dense (positions × max_candidates) table during tree packing on CPU. TreeRow.pack now emits candidate token ids concatenated in position order (sampling_mask) plus per-position counts (sampling_mask_counts), and the collator pads those ragged 1D buffers across ranks instead of a wide 2D matrix.

In compute_loss, the trainer reconstructs the padded mask on the device (scatter from counts + flat candidates) before the forward, so the fused head/kernel still see the same layout. Trainer signature columns include sampling_mask_counts.

Tests assert the new collator layout and extend tree-packing vs sequence-packing parity with a truncated-sampling case that fills realistic candidate sets.

^{Reviewed by Cursor Bugbot for commit c1648a9. Bugbot is set up for automated code reviews on this repo. Configure here.}

[lewtun]

LGTM with a question on whether we should have a dedicated test for the new width > 0 path

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.