Packing documents: how not to waste compute on padding

Build a Small GPT From Scratch · Part 2 of 8

Documents come in all lengths, but a GPU wants rectangles. How you turn one into the other decides how much of your compute is spent on real text and how much on nothing.

The padding approach

The obvious way to build a batch is to give each document its own row and fill the unused space with a special padding token. It is simple, and wasteful. Here are eight short documents of 37, 12, 90, 25, 61, 8, 48 and 30 tokens, in rows that hold 128 tokens each:

Eight rows hold 1024 token slots, but only 311 of them are text. That is 30% useful, and the model spends the rest of its effort on padding that teaches it nothing.

The packing approach

Packing joins the documents into one long stream, puts an end-of-text token between neighbours, and cuts the stream into rows of exactly 128 tokens. The same eight documents plus seven separators come to 319 tokens and fit in 3 rows: 83% useful, and the only waste is the empty tail of the last row.

The catch

A document can now start in the middle of a row, and the model could in principle attend to the unrelated text just before it. Careful implementations either mask attention at document boundaries or accept the small leakage, because the end-of-text token already tells the model that a new document has begun. Either way, the model should never be asked to predict across a boundary as if it were ordinary text.

Real runs add one more step: the stream is shuffled in row-sized chunks, so a batch mixes many unrelated documents.

Takeaways

• Padding gives each document its own row and wastes the unused space.
• Packing joins documents with end-of-text tokens and cuts the stream into equal rows.
• In the worked example, 30% of slots held text with padding against 83% with packing.
• Document boundaries need either attention masking or an end-of-text signal.

Next in the series (Part 3): Warmup then cosine: the learning-rate shape of big training runs