NVIDIA-NeMo/RL #3380 — direct path for Megatron-LM prepacked SFT inputs.
Code links pinned to ed4d534, the commit this was found on.
if len(pack_tokens) == pack_length: break. That catches the exact-fill case before it
can reach the code that wrote the duplicate boundary. Everything below describes the behaviour
before that commit.
Two pages on this PR — this one: a row that crashed training, and how it was fixed.
This PR reads SFT data that was already packed offline: one row holds several conversations
laid end to end, and a list of boundaries says where each one starts and stops. That list is
cu_seqlens — "cumulative sequence lengths" — and the attention kernel uses it to keep
conversations from reading each other.
The packing loop adds one boundary per conversation, then asks "am I full?". The check is
len(pack_tokens) >= pack_length + 1.
A pack sitting at exactly the length limit does not pass that check, so the loop takes one more
conversation — and the cleanup that follows writes a boundary that is already there. Two identical
boundaries in a row mean a conversation of length zero.
Two earlier checks wave the row through. The collate step only asks that boundaries never go
down, so 6 followed by 6 is acceptable to it. The divisibility check asks
that each conversation length divides evenly by twice the context-parallel size — and zero divides
by anything. Only the last check, which runs on a GPU worker after the model is in memory, asks for
lengths above zero. That is where the run dies:
The loop is a faithful copy of Megatron-LM's, duplicate boundary and all. What is not copied is what upstream does next. Upstream pads its boundary list out to a fixed width using copies of the pack length, then, when reading it back, keeps everything up to the first copy of that value and throws away the rest. A real duplicate looks exactly like that padding, so it gets thrown away too, and the attention kernel never sees the empty conversation.
This PR uses a variable-length list with an explicit count instead — cleaner, and it does not need the padding trick. But dropping the padding also dropped the accidental cleanup that came with it. So the producer now emits something its own checker refuses.
Same numbers as above. Two ways to close it: reject the duplicate after it is written (what this page originally suggested), or never reach that code at all — which is what the author did.
the branch to change: megatron_sft_packed.py:291-298
Running the loop before and after, on the same row:
The other branch that closes out a pack — the one that pads a short row up to length — cannot produce a repeat, because reaching it means the previous boundary was below the pack length. So this one branch is the only place that needs the check.
max_total_sequence_length, with at least one conversation still to come, kills the job.
It fails loudly rather than training on bad data, but it fails late — on a GPU worker, after
startup, on whichever step first draws that row. Turning on context parallelism makes it more
likely, not less: conversations get padded up to a multiple of twice the context-parallel size,
so there are fewer distinct lengths to land on and exact landings come up more often. The fix is one line and rejects nothing real, since a conversation of length zero has no meaning.