Sequence packing with block-diagonal masks

Summary

Sequence packing concatenates multiple training examples back-to-back into a single fixed-length sequence and uses a block-diagonal attention mask so each example only attends within itself, eliminating the FLOPs and memory wasted on padding tokens.

Most NLP corpora have heavily skewed length distributions: many short examples, few long ones. With naive padding to the longest example in the batch, the wasted-token ratio is

$$1-\frac{\bar{\ell }}{{\ell }_{\max }}$$

For C4-like web text this is often 50–80%. Padded positions cost full FLOPs and memory but contribute nothing to the loss. Sequence packing recovers nearly all of that throughput.

The mechanism

1. Pick a fixed packed length $L$ (e.g., 8192).
2. Concatenate examples $e_1,e_2,\dots$ until adding the next would exceed $L$. Record the boundaries (cumulative sequence lengths, often called cu_seqlens).
3. Build a block-diagonal attention mask: token $i$ in example $e_a$ cannot attend to any token in $e_b\ne e_a$.
4. Compute attention with a kernel that respects cu_seqlens (FlashAttention-2 supports this natively via the varlen API).
5. Apply the loss only on response tokens within each example (mask the boundaries and any prompt tokens for SFT).

Position IDs reset at each example boundary so position 0 is the start of each packed example.

Learning objective: map packed-example boundaries to reset position IDs and a causal block-diagonal attention mask that permits no cross-example attention.

Learning objective

How can one token row still behave like three separate examples?

1 · PACK TOKENS, KEEP BOUNDARIES

slot	0	1	2	3	4	5
example	e1	e1	e1	e2	e2	e3
position	0	1	2	0	1	0

Boundaries: `cu_seqlens = [0, 3, 5, 6]`
The intervals [0, 3), [3, 5), and [5, 6) recover the three examples without padding.

2 · MASK BY EXAMPLE AND TIME

q \ k	e1:0	e1:1	e1:2	e2:0	e2:1	e3:0
e1:0	A	F	F	X	X	X
e1:1	A	A	F	X	X	X
e1:2	A	A	A	X	X	X
e2:0	X	X	X	A	F	X
e2:1	X	X	X	A	A	X
e3:0	X	X	X	X	X	A

A = attend · F = future · X = other example
Only the three lower-triangular diagonal blocks participate in attention.

Numbers

For pretraining on web text packed at $L=8192$:

• Wasted tokens drop from ~50% (naive batching) to <2% (just the slack at the end of the packed sequence).
• Throughput per GPU roughly doubles.
• Quality is unchanged when the mask is correct.

Most modern training stacks pack by default. Llama, Mistral, Qwen, and major SFT toolkits (axolotl, TRL) all support it.

Common implementation pitfalls

• Forgetting to reset positions. If position IDs continue across boundaries, attention learns to treat packed boundaries as long-range dependencies and quality drops.
• Wrong loss masking. Loss must not flow from one example to another; mask boundaries explicitly.
• Mixing prompt and response in SFT without masking. For SFT, mask out prompt tokens from the loss within each packed example.
• Using a kernel that doesn’t support varlen. Without FlashAttention-2 varlen (or equivalent), the block-diagonal mask materializes the full $L\times L$ matrix and you lose the speedup.

When not to pack

• Very long single examples that fill or exceed $L$ on their own (no concatenation possible; padding is already minimal).
• When examples have inter-document context that should attend across boundaries (rare).
• During inference, where examples come one at a time.

Related

• FlashAttention. The underlying kernel that makes varlen attention efficient.
• Gradient accumulation. Orthogonal way to grow effective batch size.