Packing, Masking, and Tokenization as a Model Interface

A batch of short and long documents can spend more time processing padding than text. Packing puts several documents in a row so the accelerator does less wasted work. To make it behave like the original batch, we also need to preserve attention boundaries, prediction targets, and loss weighting.

Prerequisites: the tokenization pipeline, especially special tokens and chat templates. An attention mask is a grid specifying which positions can read from which other positions; Arc 4 develops the calculation. The variable-length kernel section also builds on the FlashAttention example in vectors and spaces.

The padding waste problem

A common training layout uses a rectangular tensor: every row has the same allocated length. If a batch contains sequences of length 14, 87, 203, and 51, padding to 203 adds 189, 116, 0, and 152 padding positions respectively.

A dense implementation can still compute states for padding positions even when their losses are ignored. The wasted fraction depends on the length distribution and batching policy. Krell et al. study packing as a way to reduce it while preserving independent examples.

Length bucketing reduces padding when rows remain rectangular. Variable-length kernels can skip padded positions entirely.

Sequence packing

Packing concatenates multiple sequences 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, 125 wasted), you get one row of 512 with 7 wasted.

Choosing which sequences share a row is a bin-packing problem. Heuristics such as first-fit-decreasing can reduce unused positions, but their utilization depends on the length distribution and allowed row sizes. Token utilization is not the same as GPU utilization: kernels, memory traffic, and scheduling still determine throughput.

The next question is whether the packed documents should be allowed to attend to one another.

Document boundaries

With one causal mask over a packed row, every position can attend to earlier positions regardless of document boundaries. A token in Document B can therefore read tokens in Document A; a token in Document C can read both preceding documents.

Some pretraining pipelines intentionally concatenate documents under one causal mask. If the goal is to reproduce independently processed examples, though, that changes the objective: a prediction in Document B now has access to Document A.

An EOS token marks a boundary in the text, but does not enforce one in attention. The model can learn to use that marker; only the attention mask or the kernel's sequence boundaries guarantee that tokens from the preceding document are inaccessible.

Select tokens below to inspect which earlier positions are allowed. Switch to isolated documents to remove cross-document connections.

Packing-aware masking

To isolate documents, use a block-diagonal attention mask. Each document has its own triangular causal block, with connections outside the block masked to negative infinity so they receive zero softmax weight.

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.

Position IDs also deserve a check. A model with learned absolute position embeddings sees a different input if Document B starts at position 3 instead of 0. Resetting positions preserves the unpacked convention. With standard RoPE and isolated attention, a constant offset within a document cancels from relative rotation differences; resetting is not inherently required for that mathematical reason. Some kernels nevertheless use position resets to infer sequence boundaries.

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]):
    """
    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 = attend, False = block.
    """
    total_len = sum(seq_lengths)
    # Start with all zeros (block everything)
    mask = torch.zeros(total_len, total_len, dtype=torch.bool)
    
    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=torch.bool))
        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.

This Boolean mask follows PyTorch scaled-dot-product attention's convention. An additive mask instead needs 0 for allowed positions and negative infinity for blocked positions; a floating 0/1 matrix is not equivalent. Check the API because Boolean conventions differ.

There is also an allocation cost: 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 (which can be shared across heads).

Flash Attention and variable-length sequences

The FlashAttention implementation supports variable-length sequences through its cu_seqlens interface. It can enforce document boundaries without allocating a dense mask.

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.

Kundu et al. evaluate variable-length packing for training efficiency.

# 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.

Preserve the targets and the update

An attention mask alone is not enough. In next-token training, the last token of Document A must not be assigned the first token of Document B as its target if the documents are meant to be independent.

For an interface that takes unshifted labels and shifts them internally, mask the label at each document start:

labels = input_ids.clone()
doc_starts = [0, 3, 5, 10]
labels[doc_starts] = -100  # ignore these targets after the internal shift

Here input_ids is one packed row. Keep each document's EOS token as a target predicted from its own preceding token. If you construct shifted targets yourself, apply the mask to the preceding prediction position instead. Padding and any prompt-only loss mask must also be applied.

Packing need not change the gradient. If both layouts contain the same examples in one logical batch, isolate their attention, preserve positions and targets, sum the same token losses, and divide by the same denominator, they compute the same objective. Kernel arithmetic and dropout can introduce numerical differences.

Changing the number of examples per optimizer step, averaging per row instead of per valid token, or changing the sampling order is a separate training choice. Control those choices when comparing packed and unpacked runs.

Tokenization as a model interface

The vocabulary maps text to IDs; the embedding table assigns those IDs learned vectors. Templates supply conversation structure, attention boundaries control available context, and loss masks select prediction targets.

Replacing a tokenizer with one that assigns different meanings to the same IDs breaks the learned lookup. Vocabulary adaptation is possible, but it requires remapping or initializing weights and further training, rather than merely changing a preprocessing file.

Save the tokenizer and template alongside the weights. For training, also record the packing, position, and loss conventions so a later run can reproduce the input and objective.

Worked example

Consider "world" in Document B. With a causal mask over the whole row, its attention can include "The", "cat", and "sat" from Document A, as well as "Hello" and "world" from its own document. With isolated blocks, only the latter two positions are available.

Removing keys changes the softmax denominator and generally changes the attention output. The visual shows allowed connections, not measured attention weights.

Misconceptions

"Packing must change the gradient." It need not. Preserve the logical batch, boundaries, targets, positions, and loss normalization. Comparing three separate optimizer updates with one combined update is a change of training procedure, not a necessary consequence of packing.

"EOS tokens isolate documents." EOS is a learned marker. Isolation requires the attention mask or sequence boundaries passed to the attention kernel.

"You can swap tokenizers on a trained model." An arbitrary swap is not compatible. The embedding matrix is indexed by token ID; if an ID changes meaning, its row still represents the old token. 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." Attention masks control which positions can read one another. Loss masks select which prediction errors contribute to the objective. A response-only SFT recipe trains on assistant targets while retaining the prompt as input; those prompt positions can still receive gradients through predictions that attend to them. A packed SFT batch can use both kinds of mask.

Closing Arc 3

The pipeline connects what text actually is at the byte level to BPE and other tokenization algorithms. Vocabulary design determines the available IDs; the embedding table turns those IDs into vectors. Special tokens and chat templates encode conversation structure, while packing and masking determine how examples share a training batch.

To check a packed batch, compare its loss with the unpacked examples under the same positions, targets, and normalization.

What's next

Arc 4 starts with Self-Attention: Q, K, V from First Principles. It derives how each position combines information from the positions its attention mask allows.

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.

• Ali, M., et al. (2023). Tokenizer Choice For LLM Training: Negligible or Crucial?.

• PyTorch. Scaled dot-product attention. Boolean masks use True for allowed attention; floating masks are added to scores. Other attention APIs may use the opposite Boolean convention.