Skip to content
AI360Xpert
Data Concepts

Sequence Packing

Padding short samples to the longest one wastes most of your compute on nothing. Packing concatenates them instead, and the attention mask is the catch.

Padded rows spend most of their token slots on nothing, and packing the same samples into full rows recovers that compute at the price of a mask that must stop attention at every document boundary
Padded rows spend most of their token slots on nothing, and packing the same samples into full rows recovers that compute at the price of a mask that must stop attention at every document boundary

Why Does This Exist?

You're fine-tuning on 200,000 support tickets. The context length is 512 tokens and the average ticket runs 114.

Every batch is padded to 512, so roughly three quarters of the token slots your accelerator processes contain a pad token — no information, no gradient, full cost. You're paying for four times the compute you're using, and the loss curve looks completely normal while you do it.

Packing concatenates samples end to end and cuts the stream into full rows, so nearly every slot holds a real token. It's one of the largest free speedups in language-model training, and it comes with exactly one correctness condition that everyone gets wrong once.

The condition, stated up front

Two documents now share a row. Attention is all-to-all within a row, so by default token 300 — the middle of ticket two — can attend to token 50, the middle of ticket one. The model reads across a boundary that doesn't exist, and it will learn from it happily.

Think of It Like This

Sixty notes in envelopes built for a manuscript

You have sixty short notes to post and a box of envelopes sized for a manuscript. One note per envelope means paying manuscript postage sixty times to move sixty postcards' worth of paper.

So you put ten notes in each envelope. Six envelopes instead of sixty, and almost all the paper you're paying to move is writing.

Now the part that matters. Whoever opens the envelope has to know where each note ends. Miss that and they read ten notes as one rambling document — every sentence individually correct, the whole thing meaningless, and the last line of note three answering a question from note four.

That divider is your attention mask. Packing without one is putting ten notes in the envelope loose.

How It Actually Works

Where the waste comes from

Batches are rectangular. Every row in a batch has to be the same length, so the batch is padded to its longest member and every shorter row carries pad tokens through the whole forward and backward pass.

Two things make it worse than the average suggests. Padding is set by the longest sample in the batch, not the mean, so one 500-token outlier pads the other 31 rows. And the cost of attention is quadratic in row length, so a padded row is disproportionately expensive rather than merely wasteful.

The mask is the whole correctness story

After packing, a row holds several documents. What you need is attention restricted to within each document, which means the attention mask becomes block diagonal — one block per document, zero everywhere between blocks.

Get it wrong and there's no error. The model trains, the loss falls, and it has learned that documents continue into unrelated documents. The damage is subtle and real: at generation time the model has a mild tendency to drift into a new topic mid-answer, because that's what its training data did.

Modern attention kernels take a cu_seqlens-style argument — the cumulative boundary offsets of each document in the row — and handle this natively, which is both faster and less error-prone than materialising a mask. Use that path if it exists.

Two related details. Position ids must restart at each document, or document two believes it starts at position 300. And the loss mask has to exclude pad tokens and, usually, prompt tokens, which is a separate mask from the attention one and is easy to conflate.

Splitting versus dropping

Some documents are longer than the context. Two options: truncate, losing the tail, or split across rows, which preserves the text and severs the dependency at the cut. For pretraining, splitting is standard and the loss of context at boundaries is accepted. For instruction data, a split example is usually incoherent, so filtering by length beats splitting.

The alternative, and when it's enough

Length bucketing sorts samples by length and batches similar lengths together, so padding within a batch is small. Simpler than packing, needs no mask changes, and it works well when your length distribution is bimodal rather than long-tailed. The cost is that batches are no longer randomly composed, which correlates the gradient within a batch — so shuffle within buckets and keep buckets wide.

Packing is strictly better on throughput. Bucketing is strictly safer. On a first run, bucketing is the sensible default; on a long pretraining run, packing pays for the care it needs.

Worked example

Eight samples in a batch, token lengths 120, 90, 300, 40, 55, 30, 210 and 65. That's 910 real tokens, and the context length is 512.

Padded, one sample per row: 8×512=4,0968 \times 512 = 4{,}096 token slots processed. Utilisation is 910/4096=22.2%910/4096 = 22.2\%, so more than three quarters of the compute goes to pad tokens.

Packed: concatenate to 910 tokens and cut into full rows, giving 910/512=2\lceil 910/512 \rceil = 2 rows and 1,024 slots. Utilisation is 910/1024=88.9%910/1024 = 88.9\%.

Four times fewer token slots, from the same eight samples. And the attention mask on those two rows needs five and three blocks respectively, because the concatenation crossed document boundaries inside each row.

Show Me the Code

import numpy as np
lengths: np.ndarray = np.array([120, 90, 300, 40, 55, 30, 210, 65])context: int = 512
def padded_slots(lens: np.ndarray, ctx: int) -> int:    """One row per sample, every row padded out to the context length."""    return int(len(lens) * ctx)
def packed_slots(lens: np.ndarray, ctx: int) -> int:    """Concatenate, then cut into full rows. The floor is the real token count."""    return int(np.ceil(lens.sum() / ctx) * ctx)
print(int(lengths.sum()), padded_slots(lengths, context))            # -> 910 4096print(round(lengths.sum() / padded_slots(lengths, context), 3))      # -> 0.222print(packed_slots(lengths, context))                                # -> 1024print(round(lengths.sum() / packed_slots(lengths, context), 3))       # -> 0.889print(padded_slots(lengths, context) // packed_slots(lengths, context))  # -> 4

np.cumsum(lengths) gives you the boundary offsets the mask needs, which is the other half of doing this correctly.

Watch Out For

Packing without a document-boundary mask

You implement packing, throughput jumps four times, and the loss curve looks better than before because each step now sees four times as many real tokens.

Nothing tells you the mask is wrong. Attention across a boundary is a valid computation on valid tensors, and the model learns from it: every packed row teaches that one document flows into the next. What you get is a model with a slight, persistent tendency to change subject mid-response, and no test that names the cause.

Position ids are the twin bug. Leave them running continuously across the row and document two starts at position 300, so the model sees a document that begins in the middle of a context window — which never happens at inference.

Verify rather than trust. Build one packed row from two documents, run it, then run each document alone, and compare the logits for the second document's tokens. If packing is correct they match. If they don't, your mask is leaking across the boundary, and that check takes five minutes.

Confusing the attention mask with the loss mask

They're different masks doing different jobs, and the names in most codebases don't help.

The attention mask controls what a token may look at. The loss mask controls which positions contribute to the gradient. Pad tokens need excluding from both. Prompt tokens, in instruction tuning, need excluding from the loss and including in attention, because the answer has to be able to read the question.

Two failures follow. Include pad positions in the loss and the model spends part of its capacity learning to predict padding, which dilutes every real gradient. Exclude prompt tokens from attention rather than from the loss, and the model is being asked to produce answers without reading the questions.

Keep them as separately named tensors, and assert their sums: the number of loss positions should equal your real answer-token count, and it's a one-line check that catches both.

The Quick Version

  • Batches are rectangular, so every row is padded to the longest member and short samples carry pad tokens through the full pass.
  • Eight samples totalling 910 tokens at a context of 512: padded is 4,096 slots at 22.2% utilisation, packed is 1,024 at 88.9%. Four times less compute.
  • Padding is set by the longest sample, not the mean, and attention cost is quadratic in row length, so the waste compounds.
  • Packing puts several documents in a row, so the attention mask must be block diagonal — one block per document.
  • Get the mask wrong and nothing errors. The model learns that documents continue into unrelated documents.
  • Position ids restart at every document boundary. Prefer a kernel that takes cumulative sequence lengths over building a mask by hand.
  • The attention mask and the loss mask are different tensors. Pads leave both; prompts leave only the loss.
  • Length bucketing is the safer alternative: less throughput, no mask changes, and a good default for a first run.
  • Verify by running one document alone and inside a packed row, then comparing logits.

Related concepts