Learn AI Create the Future — with Pranay
Main site

How Models Are Actually Trained

Packing documents into fixed-length batches

Documents come in every length but GPUs want fixed-size blocks, so pretraining glues documents end to end and slices the stream into equal pieces.

Read these first

On this page 8
  1. The short answer
  2. The analogy you have already lived
  3. Why it exists
  4. How it works
  5. The catch nobody mentions
  6. Where you have already seen this
  7. Remember this
  8. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

The short answer

Pretraining glues all the documents into one long ribbon of text, then cuts that ribbon into equal-length pieces.

The analogy you have already lived

Think of a shared auto-rickshaw at a station. It seats six. The driver could take one family of two and leave, then come back for the next family. Four empty seats, every trip, all day.

Instead the driver fills every seat before leaving, mixing families. Same number of trips, far more people moved.

A GPU is that auto-rickshaw. Its seats are token slots, and it charges you for empty ones exactly as much as full ones.

Why it exists

Documents are not the same length. A tweet is twelve words. A Wikipedia article is nine thousand.

But the hardware wants rectangles. Every row of a batch must be the same length, because that is what a matrix is.

The old fix was padding — filling the leftover slots with a meaningless filler token. It works, and it wastes a shocking amount of money. Half your electricity bill can go to multiplying zeros by zeros.

How it works

   PADDING                        PACKING
   one document per row           one long ribbon, cut evenly

   ###.....                       ###|########|##|
   ########                       ↓ cut every 8 slots
   ##......                       ###0#####
   #####...                       ###00###0
   #.......                       ...
   (# is a token, . is waste)     (every slot carries a real token)

The gain is not small. On a realistic mix of document lengths, padding wastes somewhere between a third and two-thirds of every batch. Packing wastes almost none.

The catch nobody mentions

If you glue a cricket report onto a cooking recipe, the model can read across the join. It might start predicting the recipe using the cricket score.

So packing needs two repairs.

A separator. A special end-of-document token goes between documents. The model learns from it where a text ends.

A wall. The model is told not to look backwards past a document boundary. Each document in the row is walled off from its neighbours.

Both repairs are cheap. Skipping them is a real and common bug.

Where you have already seen this

  • Loading a bus fully before it leaves, instead of one family per bus.
  • Fitting oddly-shaped items into a suitcase so no gaps remain.
  • A printer laying out several small pages onto one big sheet before cutting.

Remember this

  • Documents have different lengths, hardware wants equal-length rows.
  • Padding is easy and wastes a huge fraction of your compute.
  • Packing wastes nothing, but you must add a separator and a wall between documents.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Runs in under a second on a CPU.

Padding against packing, measured

packing.py
import torch

# Nine tokenised documents of wildly different lengths (token ids are fake but the
# length distribution is realistic: mostly short, a few long).
docs = [
    [11, 12, 13], [21, 22, 23, 24, 25, 26, 27, 28],
    [31, 32], [41, 42, 43, 44, 45], [51],
    [61, 62, 63, 64, 65, 66], [71, 72, 73, 74],
    [81, 82, 83, 84, 85, 86, 87], [91, 92, 93],
]
EOS = 0
L = 8  # context length of our toy model

# --- strategy A: one document per row, padded to L ---
padded = [d[:L] + [EOS] * (L - len(d[:L])) for d in docs]
real = sum(min(len(d), L) for d in docs)
print("PADDING")
for row in padded:
    print("   ", "".join("#" if t else "." for t in row))
print(f"    rows {len(padded)}  cells {len(padded) * L}  real tokens {real}"
      f"  wasted {100 * (1 - real / (len(padded) * L)):.0f}%")

# --- strategy B: concatenate everything, then cut every L tokens ---
stream, doc_of = [], []
for i, d in enumerate(docs):
    stream += d + [EOS]
    doc_of += [i] * len(d) + [i]          # which document each slot came from
rows = [stream[i:i + L] for i in range(0, len(stream) - L + 1, L)]
owner = [doc_of[i:i + L] for i in range(0, len(doc_of) - L + 1, L)]
print("\nPACKING (concatenate and cut)")
for r, o in zip(rows, owner):
    print("   ", " ".join(f"{t:>2}" for t in r), "   docs:", sorted(set(o)))
print(f"    rows {len(rows)}  cells {len(rows) * L}  real tokens {len(rows) * L}  wasted 0%")

# --- the cost: rows now mix documents, so masks and positions must be rebuilt ---
row, own = rows[0], owner[0]
seg = torch.tensor(own)
block = seg[:, None] == seg[None, :]                  # same document?
causal = torch.tril(torch.ones(L, L, dtype=torch.bool))
print("\nblock-diagonal causal mask for row 0 (1 = allowed to attend):")
print((block & causal).int().numpy())

pos = torch.zeros(L, dtype=torch.long)
for i in range(1, L):
    pos[i] = 0 if own[i] != own[i - 1] else pos[i - 1] + 1
print("position ids that restart at every document boundary:", pos.tolist())
Output
PADDING
    ###.....
    ########
    ##......
    #####...
    #.......
    ######..
    ####....
    #######.
    ###.....
    rows 9  cells 72  real tokens 39  wasted 46%

PACKING (concatenate and cut)
    11 12 13  0 21 22 23 24    docs: [0, 1]
    25 26 27 28  0 31 32  0    docs: [1, 2]
    41 42 43 44 45  0 51  0    docs: [3, 4]
    61 62 63 64 65 66  0 71    docs: [5, 6]
    72 73 74  0 81 82 83 84    docs: [6, 7]
    85 86 87  0 91 92 93  0    docs: [7, 8]
    rows 6  cells 48  real tokens 48  wasted 0%

block-diagonal causal mask for row 0 (1 = allowed to attend):
[[1 0 0 0 0 0 0 0]
 [1 1 0 0 0 0 0 0]
 [1 1 1 0 0 0 0 0]
 [1 1 1 1 0 0 0 0]
 [0 0 0 0 1 0 0 0]
 [0 0 0 0 1 1 0 0]
 [0 0 0 0 1 1 1 0]
 [0 0 0 0 1 1 1 1]]
position ids that restart at every document boundary: [0, 1, 2, 3, 0, 1, 2, 3]

Deterministic — no randomness anywhere. Written against PyTorch 2.5.1.

What that output is telling you

9 rows became 6, with nothing thrown away. A third fewer forward passes for the same tokens. On a real corpus the padding waste is commonly 40–60%, so this ratio is not a toy exaggeration.

The zeros inside packed rows are end-of-document markers, not padding. They are real training targets. Learning "this document has ended" is how a model learns to stop generating.

The mask is block-diagonal, not triangular. Look at row 4 of the matrix: it is allowed to see position 4 and nothing before. Document 1 starts there, and the cricket report cannot leak into the recipe.

Position ids restart at every boundary. With rotary or learned position encodings, a document that lands at slot 4 must still see itself as starting at position 0. Forget this and the second document in every row is trained at positions it will never see at inference.

How this is done in practice

Real pipelines do the concatenate-and-cut trick at tokenisation time, writing one giant memory-mapped array of token ids. Batching then becomes slicing, with no per-example work at all.

python
# the entire "dataloader" for a packed pretraining corpus
import numpy as np
tokens = np.memmap("corpus.bin", dtype=np.uint16, mode="r")   # billions of ids

def get_batch(batch_size, block_size, rng):
    ix = rng.integers(0, len(tokens) - block_size - 1, size=batch_size)
    x = np.stack([tokens[i:i + block_size] for i in ix])
    y = np.stack([tokens[i + 1:i + 1 + block_size] for i in ix])
    return x, y

No output block — this needs a real corpus.bin you have built. The point is the shape of the code: random offsets into one flat array, and nothing else. This is what nanoGPT does, and what most from-scratch pretraining code looks like.

Two variants you will meet:

  • Greedy / first-fit bin packing. Instead of cutting blindly, sort documents and fit whole documents into rows without splitting them. Better for fine-tuning, where cutting an instruction in half is destructive. See formatting an SFT dataset.
  • Padding-free / varlen. Flatten the batch into one long sequence plus a cu_seqlens array of boundaries. FlashAttention's varlen kernels consume exactly this, and it is what padding_free=True turns on in TRL. Details in FlashAttention and SDPA.

Common mistakes

Packing without a block mask. Everything trains, loss looks normal, and the model quietly learns to condition across unrelated documents. The damage shows up as odd topic drift in generation. Test it: pack two documents you know are unrelated, and check the attention mask matrix.

Packing without resetting position ids. Less obvious than the mask bug, and equally real. Print position_ids for one packed row before trusting the pipeline.

Packing an instruction dataset by cutting. Splitting a question from its answer teaches the model to answer questions it cannot see. Use whole-document bin packing for supervised fine-tuning, never blind cutting.

Counting padded tokens in your throughput number. "80,000 tokens per second" including 46% padding is a marketing number. Report real tokens per second.

Forgetting that the loss denominator changed. With packing, every position contributes a loss term. With padding and -100 labels, only real tokens do. Compare like with like when your loss curve moves after a pipeline change.

Try it yourself

Change L from 8 to 4 and re-run. Count how many documents get split across two rows. Then work out why pretraining tolerates that, and instruction tuning does not.

What to learn next

Researcher — Mathematics and papers.

The waste, quantified

Let document lengths $\ell_1, \dots, \ell_n$ be drawn from distribution $P$ and let $L$ be the context length. Padding each document to $L$ gives utilisation

$$ U_{\text{pad}} = \frac{\mathbb{E}[\min(\ell, L)]}{L} $$

Document lengths on web corpora are heavy-tailed and approximately log-normal, so $\mathbb{E}[\ell] \ll L$ for the long contexts now standard. At $L = 8192$ against a median web document of a few hundred tokens, $U_{\text{pad}}$ falls below 0.1. Concatenative packing gives $U = 1$ by construction.

Since attention cost is quadratic in sequence length while the padded region is masked, naive padding also wastes $O(L^2)$ attention work per row unless a varlen kernel is used.

The three regimes

1. Concatenative packing (pretraining). Documents are joined with EOS and the stream is cut at fixed $L$. Utilisation is exactly 1. Documents are split freely; over a full epoch each document appears at many offsets, so no information is lost in expectation.

2. Bin packing (supervised fine-tuning). Documents must not be split. This is 1-D bin packing, NP-hard, so heuristics are used. Krell et al., 2022 (Efficient Sequence Packing without Cross-contamination) introduced shortest-pack-first histogram packing, reaching >99% utilisation on Wikipedia-length distributions with an $O(n)$ pass — a 2× BERT pretraining speed-up purely from packing.

3. Padding-free / varlen. No packing into rows at all. The batch is one flat sequence of total length $\sum \ell_i$ with cumulative offsets cu_seqlens. FlashAttention-2 (Dao, 2023) exposes flash_attn_varlen_func, which computes exact block-diagonal attention with no masked-out work whatsoever. This is strictly better than a masked block-diagonal implementation, which still pays for the masked quadrant.

Cross-contamination

Naive packing computes standard causal attention over the concatenated row. Position $i$ in document $B$ attends to all of document $A$. The formal fix is a segment-aware mask:

$$ M_{ij} = \begin{cases} 0 & \text{if } j \leq i \text{ and } s_i = s_j\ -\infty & \text{otherwise} \end{cases} $$

where $s_i$ is the segment (document) id of position $i$.

How much does omitting it matter? Less than intuition suggests for pretraining, and a great deal for fine-tuning. Zhao et al., 2024 (Analysing the Impact of Sequence Composition on Language Model Pre-training) found intra-document causal masking plus retrieval-based document ordering measurably improved in-context learning and reduced distraction; the effect is real but of the order of a point or two, not catastrophic. For instruction data, where a packed row contains several complete question–answer pairs, contamination directly corrupts the supervised signal and masking is mandatory.

The complementary trick from the same line of work: order the corpus so that documents landing in the same context are related (via retrieval), turning packing from a necessary evil into a source of longer-range signal.

Position ids

Three conventions coexist:

  • Reset per document — each segment starts at 0. Matches inference; standard for masked packing.
  • Continuous across the row — cheaper, and with RoPE the relative offsets within a document are still correct, so it is defensible when the block mask is present.
  • Continuous with no mask — the buggy case, and unfortunately the default in more than one popular fine-tuning script.

For RoPE specifically, resetting is the safer choice, because relative distance is all RoPE encodes and a reset guarantees the distances a document sees match inference exactly.

Interaction with loss normalisation

Under packing, loss = sum(token_losses) / num_tokens with num_tokens = batch × L, a constant. Under padding with -100 labels, the denominator varies per batch. In distributed training this creates a subtle bug: per-device mean-of-means is not the global mean when devices hold different token counts. Transformers' average_tokens_across_devices and TRL's equivalent exist precisely to correct it, and getting it wrong silently reweights your gradient by device.

Papers

What to learn next