Who packs the sequences: the loader, or the trainer?

NVIDIA-NeMo/RL #4105 · predecessor #4070 — all links pinned to f31f793.

Four pages on this PR — this one: how the new packing path compares with the packing NeMo-RL already had.

"Packing" means putting several short conversations into one long row so the GPU is not padding its way through a batch. NeMo-RL already did this inside the trainer. This PR adds a second way: Energon, the data loader, packs the rows and hands the trainer one that is already packed.

Both ways end up handing the model the same kind of object. What changes is who decides the boundaries, when they are decided, and how wide the row ends up. Below: where each one packs, what the data looks like at each hop, and the exact line where the two paths become one.

Where each path packs

Both run under the same SFTv2 driver — the difference is which box builds the boundary list.

The one-sentence difference. The trainer packs a row as wide as the batch happens to need. The loader packs a row that is always the full max_input_seq_length, because it does not know what else is coming.

What the data looks like at each hop

Same two conversations — A is 5 tokens, B is 10 — carried through both paths, with each source padded to a multiple of 4 and a pack capacity of 36.

Trainer packs: a padded rectangle squeezed inside the worker

# the batch that reaches the GPU worker: nothing is packed yet input_ids [2, 10] # 2 rows of 10; 5 of the 20 slots are padding input_lengths tensor([5, 10]) # get_packer picks which rows share a bin, by token budget bins [[0, 1]] # both fit in one # _pack_sequences_for_megatron squeezes them end to end packed_input_ids [1, 20] # A, its pad, B, its pad cu_seqlens [0, 5, 15] # where the real tokens end cu_seqlens_padded [0, 8, 20] # where each padded source ends

the squeeze: megatron/data.py:1284 — built fresh every microbatch, thrown away after the step

Loader packs: the squeeze already happened, one row per pack

# what prepare_packed_sft_batch hands back, already model-ready input_ids [1, 36] # A, pad, B, pad, then 16 tokens of tail pad input_lengths tensor([36])# the whole row counts as one token_mask [1, 36] # 1 where the token counts toward the loss sample_mask tensor([1.])# one entry per pack, not per conversation cu_seqlens PackedTensor([0, 5, 15]) cu_seqlens_padded PackedTensor([0, 8, 36]) source_ids [["convA", "convB"]] # the worker adds the shapes so nothing downstream re-packs MICRO_BATCH_INDICES [[[0, 1]]] MICRO_BATCH_LENGTHS [[36]]

the pack: packing.py:89 · the shapes: sft_worker.py:177

Three fields exist only on the new path. source_ids names which conversations went into the row, and the two MICRO_BATCH_* fields tell the worker the shapes are already decided — without them the worker would pack the row again locally and the data-parallel ranks would disagree about what they were training on. They also explain sample_mask: on the trainer path it has one entry per conversation, here one per pack, so anything counting samples now counts packs.

Note that cu_seqlens_padded ends at 36, not 20. The last padded boundary is set to the pack capacity, which folds the 16 tokens of tail padding into conversation B's padded extent. That is not cosmetic — it is the only shape the consumer accepts.

Side by side

	Trainer packs	Loader packs (this PR)
When boundaries are chosen	every step, from the batch at hand	as the loader fills its buffer
Who chooses them	`get_packer`, fitting rows to a token budget	the same packers, called from the loader
Row width	as wide as the batch needs — 20 here	always `max_input_seq_length` — 36 here
Rows per microbatch	several	exactly one pack
Where the work happens	on the GPU worker, per step	on the loader worker, ahead of the step
`cu_seqlens_q` given to the kernel	the real boundaries, separate from the padded ones	the padded boundaries, the same array twice
`pad_between_seqs`	`True` — the two arrays differ	`False` — they are equal
Works for	text and images	multimodal SFT through Energon

Where the two become one

They meet in get_microbatch_iterator, which picks a branch by looking for a key in the batch:

prepacked = "cu_seqlens" in data or "cu_seqlens_padded" in data ... if prepacked: # loader: boundaries came with the row ... one pack per microbatch, reuse the stored boundaries elif cfg["sequence_packing"]["enabled"]: # trainer: pack now ... bin-pack, concatenate, build boundaries # both end up holding the same two things: # a CP-sharded token tensor, and a PackedSeqParams(qkv_format="thd")

the branch: megatron/data.py:300

From there the attention kernel cannot tell them apart. qkv_format="thd" means one flat run of tokens plus a boundary list, instead of a rectangle of rows. The kernel reads cu_seqlens to know where one conversation stops and the next starts, so it never lets them attend to each other. Both paths produce exactly that.

They do match — but one field differs, on purpose. The trainer path gives the kernel the real boundaries in cu_seqlens_q and the padded ones in cu_seqlens_q_padded. The loader path passes the padded array for both, and a comment says why: some consumers read cu_seqlens_q as the wrap-around point, so real boundaries would let a value roll into the padding. Because the two arrays are then equal, pad_between_seqs=False is the consistent answer — causal masking keeps real tokens from seeing the pad behind them, and the loss mask zeroes them.

What the fixed width costs. Every pack ships at the full max_input_seq_length — in the example above, 36 tokens carrying 15 real ones — and each pack is its own microbatch. So step cost tracks the number of packs, which makes the choice of packing algorithm matter more here than it did on the trainer path. Both new recipes pick balanced_greedy_knapsack while the design doc calls Modified First Fit Decreasing the default recommendation. Worth a line in the PR on which one to reach for, and a tokens-per-second number next to the convergence curves.