Two pages on this PR — this one: how the new 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 at training time. This PR adds a second way:
read rows that were packed ahead of time and use them as-is.
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 much work happens per step. 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 start from the same driver — run_sft.py
builds the dataset, sft_train
pulls batches and calls policy.train(), and the policy fans the batch out to one Megatron
worker per data-parallel rank. The difference is which box does the packing.
The one-sentence difference. Online packing decides which conversations share a row
while the step is running, from whatever happened to land in that batch. Offline packing
decided it once, before training started, and the row is replayed the same way every epoch.
What the data looks like at each hop
Same two conversations — A is 4 tokens, B is 3 — carried through both paths.
Online: a padded rectangle that gets squeezed at the last moment
# after rl_collate_fn — nothing is packed yet, just a list of message logs
message_log list of 2 conversations
length tensor([4, 3])
# after sft.py flattens it into token tensors: a padded rectangle
input_ids [2, 10] # 2 rows of 10; 13 of the 20 slots are padding
input_lengths tensor([4, 3])
# inside the worker, _pack_sequences_for_megatron squeezes the padding out
packed_input_ids [1, 7] # A and B end to end
cu_seqlens [0, 4, 7] # built right now, from input_lengths
Offline: the squeeze already happened, and the labels came with it
# what the preprocessor wrote when the dataset was built
input_ids [7] # A and B already end to end
target_ids [7] # next-token labels, already shifted by one
token_mask [7] # 1 where the token counts toward the loss
position_ids [7] # restart at 0 for each conversation
packed_cu_seqlens [0, 4, 7]
# after _collate_megatron_sft_packed stacks a batch of 2 such rows
input_ids [2, 7]
target_ids [2, 7]
token_mask [2, 7]
position_ids [2, 7]
packed_cu_seqlens [2, K] # padded to the widest row with -1
packed_cu_seqlens_lengths tensor([3, 3]) # how much of each K row is real
packed_max_seqlen tensor([4, 4])
Three fields exist only on the offline path — target_ids, token_mask and
position_ids — because the preprocessor computed them when it built the pack. On the
online path the labels are derived later, from the message log, and the loss is computed by NeMo-RL
rather than by the model.
Note the shape of packed_cu_seqlens: a batch is a rectangle, but different rows can hold
different numbers of conversations. So the list is padded to the widest row with -1, and
packed_cu_seqlens_lengths says how much of each row to read. Online packing never needs
this, because it re-derives one exact-length list per microbatch.
They meet inside process_microbatch,
which has one branch per path and a shared exit. The check that picks the branch is simply whether
the batch carries a packed_cu_seqlens key:
if "packed_cu_seqlens" in data_dict: # offline: boundaries came with the row
... CP-shard the tensors, reuse the stored boundaries
elif pack_sequences: # online: pack now
... bin-pack, concatenate, build boundaries, CP-shard
# both arrive here holding the same two things:# a CP-sharded token tensor, and a PackedSeqParams(qkv_format="thd")
From there on the attention kernel cannot tell them apart. qkv_format="thd" means
"tokens, height, depth" — 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.
So they do match — but only at the last hop, and not in every field. The tensors and
the boundary list line up. Two things still differ past the merge: total_tokens (the whole
row on the online path, this rank's share on the offline one), and who computes the loss — the offline
path hands labels to the model and lets Megatron do it, the online path brings logprobs
back out and lets NeMo-RL's loss function do it.
Why keep both
They are not the same feature wearing two hats. Online packing solves "my batch has short rows and
I do not want to pay for the padding". Offline packing solves "someone else already decided the
layout, reproduce it exactly" — which is what makes a run comparable against a Megatron-LM baseline
that used the same file.
That is also why the offline path is so restrictive. It forces microbatch size 1, refuses dynamic
batching, and rejects a mixed batch of packed and unpacked rows
(_validate_direct_megatron_sft_setup).
Every one of those would let the framework re-group rows, and re-grouping is the one thing this path
exists to avoid.
One catch worth knowing. Because the boundaries are fixed on disk, a bad row is bad
every epoch — there is no re-packing step to smooth it over. That is what made the empty-conversation
bug on the next page
worth fixing: one unlucky row killed the job on whichever step first drew it, and it would draw it
again on every pass through the data. The author has since fixed it.
The same property is why it matters that the row stores text rather than tensors. The
tokenizer configured for the run re-derives the token counts, so a tokenizer that disagrees with the
one the offline packer used shifts every boundary — and conversations that no longer fit are dropped
off the end of the row without a warning. Nothing in the file records which tokenizer built it.