Why the first pack in a batch decides whether a column stays a tensor

PR #4105 · all links pinned to f31f793

Four pages on this PR, one question each — this one: an order-dependent bug found while reviewing it.

When Energon builds a batch of packs, each token gets some extra per-token columns alongside its id — for images, mm_token_type_ids marks which tokens are picture and which are text. A conversation with no picture has no such column at all. The packing code notices this and fills in zeros for the sources that lack it, so every row lines up. It does that fill once per pack. A batch holds several packs, and if the first one has no picture anywhere in it, the fill never happens for the batch — and the column stops being a tensor.

Order A — the text-only pack comes first pack 1: text only no picture column pack 2: has a picture column present looks at the first row only isinstance(values[0], Tensor) column is a Python list [None, tensor([...])] — nothing raises Order B — same two packs, other way round pack 1: has a picture column present pack 2: text only no picture column same line, first row is a tensor so the zero-fill below it runs column is a proper tensor shape [2, 8] — correct Measured, in two separate environments text-first type=list   value=[None, tensor([1,1,1,1,0,0,0,0])] picture-first type=Tensor  value=[[1,1,1,1,0,0,0,0], [0,0,0,0,0,0,0,0]] Identical packs. Only the order changed. Red = wrong, green = right.
The batch is built from the same two packs both times.

How it happens

Three things line up. Each is reasonable alone. 

# 1. Packs are grouped by a key that ignores pictures, so a text-only pack
#    and a picture pack land in the same batch all the time.
group_key = (self.fingerprint,)
#    generic_sft.py:249

# 2. The zero-fill template is built inside the per-pack loop, so it only
#    ever sees one pack. An all-text pack keeps no picture column.
for pack in packs:
    templates = {...}          # packing.py:107-121

# 3. The batch builder decides the column's type from row zero alone.
if not values or not isinstance(values[0], Tensor):
    result[key] = values       # stores the raw Python list
    continue
#    llm_message_utils.py:402 — the code just below this correctly
#    fills in zeros for missing rows. It never gets there.

The PR's author already reasoned about step 1 and wrote it down at generic_sft.py:239-242: grouping more finely "buys nothing", because the plain batch builder keeps these columns in per-message dictionaries and never stacks them. That is true of the older path. The packed path added by this PR does stack them — so the reason the coarse grouping was safe stops holding exactly where the new code runs.

The fix

Build the zero-fill template once for the whole batch instead of once per pack. Then a text-only pack gets its zeros no matter where it sits. 

# before — inside the loop, sees one pack
for pack in packs:
    templates = {...}

# after — above the loop, sees every pack
templates = {
    key: value
    for pack in packs            # ← the whole batch
    for sample in pack.samples
    for message in sample.message_log
    for key, value in message.items()
    if key not in {"token_ids", "token_loss_mask"}
    and isinstance(value, torch.Tensor)
}
for pack in packs:
    ...

Replaces packing.py:107-121. The committed test at test_energon_packing.py:117 puts the text-only source first inside one pack, which the per-pack fill already handles; the missing case is a text-only pack first in the batch.

Why this is worth catching before it lands. Nothing raises. The batch conversion accepts the Python list and stores it as NonTensorData, so training carries on with a per-token column that is no longer a tensor. There is no error message to search for and no failed assertion — only whichever batches happened to start with a text-only pack behave differently. What the model then does with that column was not measured: there was no GPU available for this review, so the claim stops at "the column stops being a tensor and nothing rejects it".