Sequence packing: concatenating short examples to fill a fixed context window

The problem, from first principles

A model trains on fixed-length windows of context_length tokens, but real training examples come in wildly varying lengths, many of them far SHORTER than context_length. Padding every short example up to the full context length individually (a common, simple approach) wastes an enormous amount of compute: [03-attention-quadratic-complexity] already showed attention's cost scales with sequence length, so running a mostly-padding sequence through the model spends real compute on positions that carry no actual information at all. Sequence packing fixes this directly: concatenate MANY short examples end-to-end into one long stream, then chop that stream into fixed-length context_length chunks, so nearly every position in every training batch is genuine content, only the very LAST chunk of the whole stream needs any padding at all.

This creates a genuinely new problem, though: after packing, two UNRELATED examples can end up sitting right next to each other within the same context_length window, and an ordinary [04-seq-modeling/04-attention/02-causal-mask]'s causal mask alone would happily let a later example's tokens attend BACKWARD into an entirely unrelated earlier example, just because they happen to share a packed window, corrupting what each example's attention actually "sees." The fix: track which original example every packed position came from, and additionally mask out any attention that would cross a document boundary, even when ordinary causality alone would permit it.

From theory to code

Implement pack_sequences(examples, context_length, pad_token), concatenating and chunking, tracking each packed position's original example via a parallel doc_ids array, and build_intra_document_mask(doc_ids), an additive mask combining ordinary causality with a "same document only" restriction.

Constraints

• Examples are concatenated in order, then chopped into context_length-sized chunks; only the FINAL chunk of the whole concatenated stream may need padding (with pad_token).
• doc_ids tracks the ORIGINAL example index for every packed position, -1 for padding positions.
• build_intra_document_mask's attention rule: position i may attend to position j only if j <= i (ordinary causality) AND doc_ids[j] == doc_ids[i] (same original example) AND doc_ids[i] != -1 (not itself padding).
• A packed sequence's chunk BOUNDARIES may (and typically will) fall in the MIDDLE of an original example, when that example is longer than what's left in the current chunk; the example simply continues into the next chunk, still tagged with its own doc_ids.

Hints

all_tokens, all_doc_ids = [], []
for doc_id, example in enumerate(examples):
    all_tokens.extend(example)
    all_doc_ids.extend([doc_id] * len(example))

for start in range(0, len(all_tokens), context_length):
    chunk = all_tokens[start:start + context_length]
    # pad the LAST chunk only, if it's short

doc_ids_arr = np.array(doc_ids)
causal = np.arange(len(doc_ids))[None, :] <= np.arange(len(doc_ids))[:, None]
same_doc = doc_ids_arr[None, :] == doc_ids_arr[:, None]
not_padding = doc_ids_arr[:, None] != -1
allowed = causal & same_doc & not_padding
return np.where(allowed, 0.0, -np.inf)