Why one conversation that exactly fills the pack crashed SFT training

NVIDIA-NeMo/RL #3380 — direct path for Megatron-LM prepacked SFT inputs. Code links pinned to ed4d534, the commit this was found on.

Fixed — this page is kept as a record of the bug, not an open issue. The author closed it in 789a2cf9c ("fix(sft): validate direct packed inputs") by adding one line above the truncation branch: 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.

Pack length limit = 6 tokens. Conversation A is 6 tokens; conversation B is 3. 1. Add A (6 tokens) cu_seqlens = [0, 6] full? 6 >= 7 is false → keep going 2. Add B (3 tokens) cu_seqlens = [0, 6, 9] full? 9 >= 7 is true → cut it down 3. Cut to 6, add 1 pad pad 0 6 6 cu_seqlens = [0, 6, 6 ] the last write set 6 — but 6 was already the previous boundary Conversations this describes A: 6 tokens fine B: 0 tokens — a conversation with nothing in it Where the row goes next collate step allows equal neighbours collate_fn.py:115 pass length-divisible check 0 divides by anything megatron_sft_packed.py:327 pass worker check rejects length ≤ 0 megatron/data.py:178 raise happens on a GPU worker, after the model has loaded

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:

ValueError: packed_cu_seqlens must contain increasing boundaries from 0 through the packed row length

Why Megatron-LM itself never hits this

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.

How it was fixed

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.

# megatron_sft_packed.py, inside the "pack is full" branch (line 291) pack_tokens = pack_tokens[:pack_length] # 9 tokens -> 6 pack_tokens.append(pad) # 6 -> 7 (one pad on the end) cu_seqlens[-1] = len(pack_tokens) - 1 # writes 6 ... but [0, 6] already ended in 6 # add this: if len(cu_seqlens) > 1 and cu_seqlens[-1] == cu_seqlens[-2]: cu_seqlens.pop() # [0, 6, 6] -> [0, 6]

the branch to change: megatron_sft_packed.py:291-298

Running the loop before and after, on the same row:

cu_seqlens conversation lengths before A=6 (fills exactly), B=3 [0, 6, 6] [6, 0] crashes after A=6 (fills exactly), B=3 [0, 6] [6] ok rows that were already fine are untouched: before A=4, B=3 [0, 4, 6] [4, 2] ok after A=4, B=3 [0, 4, 6] [4, 2] ok

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.

What this cost before the fix. Any row where some conversation lands the pack on exactly 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.