NLP Foundations
Sequence Packing and Document Masking
Concatenating short documents to fill a fixed context length eliminates padding waste, but only if the attention mask and position IDs are made document-aware; done naively, packing quietly teaches the model to attend across unrelated documents.
intermediate · 7 min read
Real documents almost never land exactly on the model's context length, T. A news article might tokenise to 400 tokens against a context length of 8192; padding it out to T means roughly 95% of every training example is pad tokens, contributing nothing to the loss while still consuming the full compute of a forward and backward pass. Sequence packing fixes the waste by concatenating multiple documents end to end until the buffer is full, but the fix introduces a subtler problem than the one it solves.
The basic idea
Instead of one document per training example, packing fills each fixed-length chunk with as many documents as fit, separated by an end-of-sequence token:
doc A (400 tok) <EOS> doc B (1800 tok) <EOS> doc C (900 tok) <EOS> doc D (...)
[chunk cut at position T]
Every position in every chunk now carries a real token and a real label; there is effectively no padding waste. This is standard practice in large-scale pretraining precisely because at trillion-token scale, even a few percent of wasted compute is enormous in absolute terms.
The leak packing introduces
The causal mask used for ordinary pretraining (see causal-lm-pretraining-task) only prevents attending to future positions; it says nothing about document boundaries. Applied naively to a packed chunk, a token near the start of document C can freely attend to every token of documents A and B that precede it in the buffer, purely because they happen to sit earlier in the same chunk, not because they are related in any way. The model can waste attention capacity, and occasionally learning capacity, modelling spurious correlations between unrelated documents that were only ever adjacent by chance of packing order.
Fixing it: block-diagonal attention and position resets
Correct packing implementations add two document-aware adjustments on top of plain causal masking:
- Block-diagonal attention masking. Each token may attend only within its own document's span, in addition to the ordinary causal (no-future) restriction. Combined, the mask for a packed chunk looks like several small triangular causal blocks stacked along the diagonal, one per document, with everything off those blocks masked out. FlashAttention-style kernels implement this efficiently via variable-length sequence indices (
cu_seqlens) rather than materialising a denseT x Tmask (see flash-attention). - Position ID resets. Positional encodings (see positional-encodings) are normally assigned by absolute position within the chunk,
0, 1, 2, ..., T-1. In a packed chunk that means document C's first token might receive position2201purely because of what happened to precede it, a number with no relationship to its actual position within document C. Resetting position IDs to0at the start of every document keeps positional semantics meaningful and consistent with how the model will see documents at inference, where each one naturally starts at position 0.
Why some pretraining runs skip the fix anyway
Implementing document-aware masking correctly costs engineering complexity: dense attention kernels either need a full mask tensor (memory-expensive at long context) or a variable-length-aware kernel. Some large-scale runs, particularly earlier GPT-3-era recipes, packed documents naively and accepted the cross-document leakage as negligible noise relative to trillions of training tokens, on the reasoning that any single spurious cross-document attention pattern is diluted into irrelevance by the sheer diversity of what else is packed around it. Whether that trade-off is acceptable depends heavily on scale and on how structured the downstream use is; instruction and chat data, packed the same way, has sharper failure modes (see loss-masking-in-fine-tuning) because a leaked fragment of an unrelated conversation is a much more visible artefact than a leaked fragment of unrelated web text.
When it falls down
- Naive packing without masking silently degrades quality in ways that rarely show up as an obvious bug, only as a slightly worse model, which makes it easy to ship unnoticed.
- Document-aware attention costs more to compute than a single dense causal mask, especially at long context lengths where the mask itself becomes memory-significant.
- Position resets interact with long-context extrapolation. Some long-context techniques rely on position IDs growing smoothly to structure how the model generalises beyond its trained length (see context-windows-long-context); frequent resets inside a packed chunk change that structure and need to be accounted for in the positional scheme.
- Document order within a chunk still matters at the margins, even with correct masking, because packing determines what a data-loading pass groups together, and shuffling strategy interacts with how documents get batched.
Further reading
- Krell et al., 2021, Efficient Sequence Packing without Cross-contamination - a systematic treatment of packing and the cross-document attention leak.
- Dao et al., 2022, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness - the variable-length attention kernels that make document-aware masking affordable at scale.