Sequence Packing Leaks Across Documents Unless You Mask It

You packed four short training examples into one 4096-token sequence to kill padding waste, your throughput doubled, the loss curve looks clean — and your fine-tuned model is quietly worse than the padded baseline. The culprit is almost always sequence packing without a block-diagonal attention mask: token 3000 in document C is attending back to document A, and RoPE thinks document C starts at position 2400 instead of position 0.

This is one of the most common silent bugs in fine-tuning pipelines because nothing crashes and the loss doesn't spike. Let me show you the exact mechanism and the three ways to fix it.

TL;DR

• Sequence packing concatenates multiple training examples into one fixed-length sequence to eliminate padding, often doubling tokens-per-batch throughput.
• A plain causal mask lets every token attend to all prior tokens in the packed sequence, so document B contaminates document A's context — cross-document attention leakage.
• Contiguous position_ids (0…N across the whole pack) also corrupt RoPE: later documents get encoded as if they're thousands of tokens into a conversation.
• The fix is a block-diagonal (document) mask plus per-document position ID reset. Use FlashAttention varlen with cu_seqlens, PyTorch FlexAttention with a mask_mod, or pass position_ids so the boundaries are inferred.
• The failure is silent: loss looks normal, throughput is great, evals drift down. Always verify your packing path masks boundaries before trusting a run.

What is sequence packing and why use it?

Sequence packing is concatenating several variable-length training examples into a single fixed-length sequence so you don't waste compute on padding tokens. If your examples average 800 tokens and your context is 4096, padding each example to 4096 means ~80% of your FLOPs go to <pad>. Packing five examples end-to-end instead fills the sequence, and every token contributes to the loss.

For instruction-tuning datasets with a long tail of short samples, packing routinely gives a 2–5x effective throughput win. That's why TRL, Axolotl, and most serious fine-tuning stacks enable it. The problem is that "concatenate into one sequence" is the easy 90%, and the attention masking is the 10% that actually determines whether the run is correct.

Why does naive packing leak attention across documents?

Because a causal mask permits every position to attend to every earlier position, and packing puts unrelated documents in the same sequence. The causal mask only enforces "no looking at the future." It says nothing about document boundaries.

Concretely, pack documents A (0–999), B (1000–1999), and C (2000–2999). Standard causal attention lets token 2500 in C attend to any token in [0, 2500], which includes all of A and B. During training the model learns to use that cross-document context — context that will never exist at inference, when C is prompted alone. You're training on a distribution you'll never serve.

The second, subtler corruption is positional. With rotary embeddings (RoPE), the relative position between query and key is baked into the dot product. If position_ids run 0…3071 straight through the pack, then the first real token of document C sits at absolute position 2000. RoPE encodes C's internal token-to-token relationships against that offset. The model sees C's opening token as "2000 tokens deep into a document," which is nothing like how C will appear at inference (position 0). Even if you somehow fixed the mask but left positions contiguous, you'd still degrade quality.

Both bugs push in the same direction: the packed training distribution diverges from the unpacked inference distribution, and the divergence grows with how many documents you cram per sequence.

What does the correct mask actually look like?

A block-diagonal mask: within a packed sequence, each document may only attend to itself, with a normal causal mask inside each block. Visually, for a pack of [2, 3] token documents:

        A0 A1 B0 B1 B2
   A0  [ 1  0  0  0  0 ]
   A1  [ 1  1  0  0  0 ]
   B0  [ 0  0  1  0  0 ]
   B1  [ 0  0  1  1  0 ]
   B2  [ 0  0  1  1  1 ]

The off-diagonal blocks (B attending to A) are zeroed. This is sometimes called an "intra-document" or "document causal" mask. Pair it with position_ids = [0, 1, 0, 1, 2] so each document restarts at position 0 for RoPE.

Materializing an N×N boolean mask works for short sequences but is memory-death at 8k or 32k context — that's why the production path uses variable-length kernels or a block-sparse abstraction instead of a dense mask.

How do I implement packing correctly?

Three approaches, roughly in order of how modern they are.

1. FlashAttention varlen with cu_seqlens. FlashAttention exposes an unpadded API that takes cumulative sequence lengths and never lets attention cross a boundary. You pass the document lengths and it computes each block independently — no dense mask ever exists.

from flash_attn import flash_attn_varlen_func
import torch

# Three packed documents of length 1000, 1000, 1072 -> 3072 tokens total.
seqlens = torch.tensor([1000, 1000, 1072], dtype=torch.int32)
cu_seqlens = torch.zeros(len(seqlens) + 1, dtype=torch.int32, device="cuda")
cu_seqlens[1:] = torch.cumsum(seqlens, dim=0)   # [0, 1000, 2000, 3072]
max_seqlen = int(seqlens.max())

# q, k, v are packed: shape (total_tokens, n_heads, head_dim), NO batch dim.
out = flash_attn_varlen_func(
    q, k, v,
    cu_seqlens_q=cu_seqlens,
    cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max_seqlen,
    max_seqlen_k=max_seqlen,
    causal=True,      # causal *within* each document block
)

cu_seqlens is the whole trick: the kernel treats [0,1000), [1000,2000), [2000,3072) as independent causal segments. Leakage is structurally impossible.

2. PyTorch FlexAttention with a mask_mod. FlexAttention lets you express the document mask as a small predicate and compiles it into a block-sparse kernel, so you get correctness without the O(N²) memory of a dense mask.

from torch.nn.attention.flex_attention import create_block_mask, flex_attention

# document_id[i] = which document token i belongs to, e.g. [0,0,...,1,1,...,2]
def doc_causal(b, h, q_idx, kv_idx):
    causal = q_idx >= kv_idx
    same_doc = document_id[q_idx] == document_id[kv_idx]
    return causal & same_doc

block_mask = create_block_mask(doc_causal, B=None, H=None,
                               Q_LEN=seq_len, KV_LEN=seq_len)
out = flex_attention(q, k, v, block_mask=block_mask)

The mask_mod reads like the math: causal and same document. FlexAttention skips fully-masked blocks, so a pack of many short documents costs close to block-diagonal FLOPs, not dense.

3. Let the model infer boundaries from position_ids. Recent Hugging Face transformers will build the correct block-diagonal mask for FlashAttention-2 when you pass a position_ids tensor that resets per document ([0,1,2,...,0,1,2,...]). The library detects the resets and constructs cu_seqlens internally. This is the least error-prone option if you're already in the HF trainer stack — but you must actually pass reset positions, not rely on defaults.

The common thread: you need both the boundary-aware mask and the reset positions. Fixing one without the other still leaves a distribution mismatch.

Why is this bug so hard to catch?

Because every signal you normally watch looks healthy. Training loss is computed per-token and averaged; a little cross-document leakage nudges each token's loss by a tiny amount, well inside normal noise. The curve is smooth. Throughput is better than the correct run because leaked attention and contiguous positions require no extra bookkeeping. Gradients are finite, there's no NaN, no shape error.

The damage shows up only in downstream evals, as a diffuse quality regression — slightly worse instruction following, more topic bleed, weaker long-context behavior — that's easy to blame on data or hyperparameters. Teams burn days tuning learning rate when the real bug is one missing mask.

A few things make it worse. The more aggressively you pack (more documents per sequence), the more contamination. Attention-sink behavior means early tokens of the sequence — document A — get disproportionately attended to by everyone downstream, so A silently contaminates B, C, and D more than a uniform model would suggest. And short documents packed near the start of long sequences are the most corrupted positionally.

How do I verify my packing is correct?

Three concrete checks:

1. Inspect the mask on a toy batch. Pack two known documents, dump the attention mask (or cu_seqlens), and confirm the off-diagonal block is zero. Ten minutes, catches the bug immediately.
2. Assert position resets. Log position_ids for a packed batch and verify each document restarts at 0. If they run 0…4095 straight through, RoPE is corrupted regardless of the mask.
3. Run a leakage probe. Take a packed sequence, then run the same documents unpacked one at a time. With correct masking, the logits for each document are identical (within float tolerance) in both cases. Any divergence means attention is crossing boundaries.

That third check is the gold standard: correct packing is provably equivalent to processing documents separately. If packed and unpacked logits don't match, you have leakage, full stop.

Direct answer: does sequence packing leak across documents?

Yes — sequence packing leaks attention across document boundaries whenever you use a plain causal mask and contiguous position IDs, and the leak is silent because loss and throughput both look fine. A standard causal mask lets tokens in one packed document attend to every earlier document, and running position_ids continuously across the pack corrupts RoPE's relative positions. The fix is a block-diagonal (intra-document) attention mask combined with per-document position ID resets, implemented via FlashAttention varlen with cu_seqlens, PyTorch FlexAttention with a document mask_mod, or by passing reset position_ids so the library infers boundaries. Verify it by asserting that packed and unpacked forward passes produce identical logits before you trust any fine-tuning run.