Packing, Masking, and Tokenization as a Model Interface

This is the last post in Arc 3, and I want to use it to tie together something I underestimated for a long time. The tokenizer and everything around it aren't preprocessing. They're part of the model. The vocabulary, the merge rules, the special tokens, the chat template, the way sequences get packed into batches, the masks that control which positions contribute to the loss. All of it forms a contract between raw text and the model's parameters. Change any piece of that contract and you've changed the model's input distribution. You might as well have changed the weights.

Before I get philosophical about it, let me start with a concrete problem: waste.

Prerequisites: the tokenization pipeline from the prior posts in this arc, especially the chat template formatting from special tokens and chat templates. We'll also reference FlashAttention from the in-the-wild section of vectors and spaces. We haven't formally met attention yet (that's Arc 4), so when I say "attention mask" here, picture a grid that says which positions are allowed to influence which other positions. The mechanics show up in the next arc.

The padding waste problem

When you train a language model, you process sequences in batches. Every sequence in the batch needs to be the same length (GPUs do parallel matrix multiplications across the batch dimension, and you can't multiply matrices of different shapes). So if your batch contains sequences of length 14, 87, 203, and 51, you pad them all to 203. That's adding 189 padding tokens to the first sequence, 116 to the second, none to the third, and 152 to the fourth.

Those padding tokens consume memory, move through the GPU, and produce gradients that you throw away. On natural language datasets where sequence lengths follow a power-law distribution (lots of short sequences, a few very long ones), this waste is staggering. Krell et al. measured up to 89% padding tokens on some datasets when using naive batching. That's nearly nine out of every ten tokens doing zero useful work.

Sorting (group similar-length sequences together) or dynamic batching (variable batch sizes) helps, but you're still fundamentally limited by the longest sequence in each batch. The GPU is doing real FLOPs on every padding token, and none of those FLOPs produce useful gradients.

Sequence packing

Here's the better idea: instead of padding, just concatenate multiple sequences end-to-end into a single fixed-length row. If your target length is 512 tokens and you have three sequences of lengths 140, 210, and 155, you pack them together into one row of 505 tokens with only 7 tokens of padding at the end. Instead of three rows padded to 210 each (630 tokens, 280 wasted), you get one row of 512 with 7 wasted. That's a huge difference.

The bin-packing problem here is NP-hard in general, but simple heuristics work well in practice. First-fit-decreasing (sort sequences by length, then greedily pack each into the first row with room) and shortest-pack-first (always add to the row with the most remaining capacity) both achieve near-100% utilization. The Hugging Face blog on efficient pretraining reports GPU utilization jumping from around 60% to 99%+ after switching to packing, and the Continuum Labs writeup walks through the same tradeoffs end-to-end.

So packing saves a ton of compute. But it introduces a nasty problem.

Cross-contamination

When you pack three unrelated documents into a single row, the model doesn't know where one ends and another begins. Under a standard causal rule (every position sees all earlier positions, nothing more, the transformer's default, which we'll unpack properly in Arc 4), the last token in Document B can see tokens in Document A. The first token in Document C has all of Document A and Document B in its context window.

This is cross-contamination. The model is computing attention weights between tokens that have nothing to do with each other, and the gradients from those attention weights flow back into the model parameters. You're training the model to learn spurious relationships between the end of a news article and the beginning of a Python function that happened to get packed next to it.

Separator tokens don't fix it. The natural reaction is, "well, I put an EOS token between Document A and Document B, so the model knows they're separate." But the attention mechanism doesn't have a concept of EOS as a boundary. It computes a dot product between every query and every key. EOS is just another token with an embedding. The model can learn to attend to it less, but that's a statistical regularity it has to discover through training, not a hard constraint. Meanwhile, every gradient update during that learning process is being polluted by cross-document attention.

Click on individual tokens below to see the difference. In naive mode, notice how tokens from one document attend to tokens from completely unrelated documents (the red cells). Toggle to block-diagonal to see the correct behavior.

Packing-aware masking

The fix is to modify the attention mask. Instead of a standard lower-triangular causal mask, you use a block-diagonal mask: each document gets its own triangular block along the diagonal, and everything outside those blocks is masked to negative infinity (so it gets zero weight after softmax).

Concretely, if Document A occupies positions 0-2, Document B occupies positions 3-4, and Document C occupies positions 5-9, you construct a mask where position 4 (last token of Doc B) can only attend to positions 3 and 4, not to positions 0-2. The attention for each document is completely self-contained.

There's a second, sneakier thing to fix. Position IDs. In an unpacked batch, the first token of every sequence has position 0. In a packed row, if you naively use range(N), Document B's first token gets position 3, and Document C's first token gets position 5. The model treats "position 3" and "position 5" as real structural information and feeds it through whatever positional encoding scheme is in use (we'll get into rotary encodings in Arc 4). Hand it a lie about where a document begins and the positional signal lies right back.

Fixing it is easy: generate position IDs that restart at 0 at every boundary. The harder part is remembering to do it at all; it's one of those details that only bites you quietly, in the form of slightly-wrong gradients, for a long time.

Here's what it looks like in PyTorch to build a block-diagonal mask from a list of sequence lengths:

import torch
 
def build_block_diagonal_mask(seq_lengths: list[int], dtype=torch.float32):
    """
    Build a block-diagonal causal attention mask for packed sequences.
    
    Args:
        seq_lengths: list of individual document lengths within the packed row.
                     e.g. [3, 2, 5, 2] for four documents packed into length 12.
    
    Returns:
        mask of shape (total_len, total_len) where True/1.0 = attend, False/0.0 = block.
    """
    total_len = sum(seq_lengths)
    # Start with all zeros (block everything)
    mask = torch.zeros(total_len, total_len, dtype=dtype)
    
    offset = 0
    for length in seq_lengths:
        # For each document, create a causal (lower-triangular) block
        causal_block = torch.tril(torch.ones(length, length, dtype=dtype))
        mask[offset:offset + length, offset:offset + length] = causal_block
        offset += length
    
    return mask
 
def build_position_ids(seq_lengths: list[int]):
    """
    Build position IDs that reset to 0 at each document boundary.
    
    e.g. seq_lengths=[3, 2, 5] -> [0, 1, 2, 0, 1, 0, 1, 2, 3, 4]
    """
    position_ids = []
    for length in seq_lengths:
        position_ids.extend(range(length))
    return torch.tensor(position_ids, dtype=torch.long)
 
# Example usage:
seq_lengths = [3, 2, 5, 2]  # four documents packed into one row
mask = build_block_diagonal_mask(seq_lengths)
position_ids = build_position_ids(seq_lengths)
 
print(f"Packed length: {sum(seq_lengths)}")
print(f"Position IDs: {position_ids.tolist()}")
print(f"Mask shape: {mask.shape}")
print(f"Mask:\n{mask.int()}")
# Output:
# Position IDs: [0, 1, 2, 0, 1, 0, 1, 2, 3, 4, 0, 1]
# Each document gets its own causal block along the diagonal.

That works. But there's a catch: the mask is a dense $N \times N$ matrix, where $N$ is the total packed sequence length. For a packed row of 8,192 tokens, that's a 67-million-element boolean matrix per attention head. This is where it gets expensive.

Flash Attention and variable-length sequences

This is exactly where FlashAttention becomes essential. Dao et al.'s FlashAttention doesn't just speed up attention by fusing operations and reducing memory traffic (which it does). It also provides native support for variable-length sequences through the cu_seqlens (cumulative sequence lengths) interface.

Instead of materializing a giant $N \times N$ mask, you pass FlashAttention a flat packed tensor and a list of cumulative sequence lengths. For our example with lengths 3, 2, 5, 2, that's cu_seqlens = [0, 3, 5, 10, 12]. FlashAttention internally computes attention within each document's boundaries without ever building the full mask matrix.

This matters a lot. Kundu et al. showed that packing with FlashAttention's varlen support achieves up to 2x training throughput compared to padded batching, with zero cross-contamination. The memory savings are even bigger because you never allocate the $N \times N$ mask at all.

# Using FlashAttention's variable-length interface (pseudocode)
from flash_attn import flash_attn_varlen_func
 
# q, k, v are packed tensors of shape (total_tokens, num_heads, head_dim)
# cu_seqlens marks document boundaries: [0, 3, 5, 10, 12]
# max_seqlen is the longest individual document
 
output = 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,
)
# Each document attends only within itself. No mask matrix allocated.

Without FlashAttention (or a similar fused kernel), you'd need to either materialize the block-diagonal mask explicitly or loop over documents sequentially. Both options are much slower. This is one of those cases where a systems-level optimization (fused CUDA kernels) directly enables a training technique (packing) that would be impractical without it. We talked about GPU fundamentals earlier in the series, and this is a concrete example of why they matter for training decisions.

The gradients are different

I want to pause and make something explicit, because I think it's easy to wave this off as "just an optimization." Packing changes what gradients your model sees.

In padded batching, each training example produces gradients from exactly one document. In packed batching, each row produces gradients from multiple documents. If you have three documents packed together and one of them is much harder than the others (say, a code snippet sandwiched between two short English sentences), the loss contribution from that code snippet now influences the gradient update alongside the other two documents. The effective batch composition per gradient step is different.

This isn't necessarily bad. In many cases it's fine, and the efficiency gains are worth it. But it means packing is not a transparent optimization you can swap in without thinking about it. The loss per token can change depending on what gets packed together, and if your packing is biased (say, all the short sequences are from one domain and the long ones from another), you might inadvertently change the domain mixture that the model sees within each gradient step.

Krell et al. found that with proper block-diagonal masking, the training dynamics with packing closely match unpacked training. But "closely" isn't "identically," and if you're doing careful ablations where you need reproducible training runs, packing is a variable you need to control for.

Tokenization as a model interface

Let me zoom out and say what I've been building toward across all of Arc 3.

The tokenizer is not preprocessing. I keep saying this because it's the most important takeaway from this arc.

When we choose a vocabulary size, we're choosing the dimensionality of the model's input and output (the embedding table has one row per token). When we choose merge rules, we're choosing what semantic units the model will learn to work with. When we choose special tokens, we're defining the model's understanding of conversation structure, tool calls, and document boundaries. When we choose how to pack and mask training data, we're choosing which attention patterns the model will learn and which spurious correlations it will pick up.

All of these are interface decisions. They're part of the model's contract with the world.

Liang et al. studied this directly and found that tokenizer choice has measurable effects on downstream task performance that interact with model scale in non-obvious ways. A tokenizer that's great for English text might fragment code badly, producing longer sequences that eat up context window and produce less coherent attention patterns. A tokenizer trained on code might waste vocabulary entries on language-specific keywords that dilute coverage of natural language.

You can't swap tokenizers after the fact. The embedding table is a learned lookup indexed by token ID. If you change the tokenizer, every token ID maps to a different subword, so every embedding is wrong. You'd need to retrain the embedding layer at minimum, and in practice the effects propagate through the entire network because every layer has adapted to the distributional properties of the original tokenization.

This is why every major model release specifies its tokenizer as carefully as its architecture. GPT-4 uses cl100k_base (100,256 tokens). Llama 3 uses a custom 128K BPE tokenizer. These aren't afterthoughts. They're design decisions that constrain everything downstream.

Worked example

Let me make this concrete with the visualization above. Consider the token "world" in Doc B (position 4 in the packed row). Click on it in the visualization.

In naive mode (standard causal mask), "world" attends to "The", "cat", "sat" (all from Doc A), plus "Hello" and "world" (from Doc B). That's five keys it computes attention weights against. The dot product between "world" and "cat" might not be zero. Softmax normalizes across all five, so even a small affinity between "world" and "sat" steals probability mass from the "Hello"-"world" connection that actually matters.

Now toggle to block-diagonal mode. "world" only attends to "Hello" and "world." Two keys, both from the same document. The softmax is now computed over just these two, so the full attention budget goes to the tokens that are actually relevant.

The attention weight for "Hello" given query "world" is literally a different number in the two modes. In naive mode it might be 0.25 (diluted across five keys). In block-diagonal mode it might be 0.75 (concentrated on two keys). Those different weights produce different context vectors, which produce different gradients, which nudge the model in different directions. Multiply that by billions of tokens over thousands of training steps and you get a meaningfully different model.

Misconceptions

"Packing is just an optimization." It's an optimization that also changes gradient composition. With proper masking the effect is small, and without masking it's significant, but either way packing isn't a transparent swap. You're changing what the optimizer sees per step.

"EOS tokens prevent cross-contamination." They don't. The attention mechanism computes scores between every query-key pair within the causal mask. An EOS token is just another entry in the embedding table. The model can learn to down-weight it, but that's a learned behavior, not a hard constraint. During early training when the model hasn't yet learned what EOS means, every gradient step has full cross-contamination. Karpathy made this point clearly in the llm.c discussion on intra-document masking.

"You can swap tokenizers on a trained model." You can't. The embedding matrix is a lookup table indexed by token ID. Changing the tokenizer changes the mapping from text to IDs, which means every embedding retrieval returns a vector that was trained for a different subword. You would need to retrain or at least adapt the embedding layer, and even then the rest of the network has co-adapted to the original tokenization's statistical properties. Some people have done vocabulary extension (adding new tokens with initialized embeddings for fine-tuning), but that's additive, not a swap.

"Loss masking and packing masks are the same thing." They solve different problems. Packing masks (block-diagonal attention masks) prevent cross-contamination between documents during attention. Loss masks control which tokens contribute to the gradient update. In SFT, for example, you mask out the prompt tokens from the loss so the model only trains on the response. You often need both at the same time: packing masks to keep documents separate in attention, and loss masks to train only on the right tokens within each document.

Closing Arc 3

We've covered a lot of ground in this arc. We started with what text actually is at the byte level, built up through BPE and other tokenization algorithms, explored vocabulary design tradeoffs, looked at how token IDs become vectors, examined special tokens and chat templates, and now we've seen how packing and masking close the loop between tokenization and training.

The through-line is this. Everything between raw text and the model's first matrix multiply is part of the model. The tokenizer, the special tokens, the chat template, the packing strategy, the mask. These are all interface decisions, and they all affect what the model learns.

What's next

Three arcs in, we've built every piece of the input pipeline. Raw bytes, a vocabulary, a tokenizer, an embedding table, special tokens, chat templates, packing, masking. Everything between the keyboard and the first matrix multiply.

Next up starts Arc 4: Transformers from First Principles. The opening post is Self-Attention: Q, K, V from First Principles, which sets up the residual-stream picture of a transformer (a per-token vector that every block reads from and writes back to) and then derives scaled dot-product attention as the operation that lets tokens talk to each other. We've spent three arcs on what happens before the first matmul. Now we finally go inside.

Additional reading (and watching)

• Krell, M., et al. (2021/2024). Efficient Sequence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance.
• Kundu, A., et al. (2024). Enhancing Training Efficiency Using Packing with Flash Attention.
• Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
• Hugging Face Blog. Packing with Flash Attention.
• Karpathy, A. llm.c Discussion #690: Intra-document masking.
• Liang, D., et al. (2023). Tokenizer Choice For LLM Training: Negligible or Crucial?.