Preparing Training Sequences: Shift, Pad, and Pack

Suppose you feed a sequence of four tokens ["A", "B", "C", "D"] into a causal transformer. In Lesson 1, we saw that the forward pass outputs a matrix of logits of shape [B, T, V]—one vocabulary score vector for each of the four positions.

Which target token belongs to each prediction position?

Let’s walk through each position explicitly before introducing shifted tensor notation:

• Position 0 (reads ["A"]): Its causal attention mask allows it to see only "A". It must predict the token that follows: "B".
• Position 1 (reads ["A", "B"]): Having attended to "A" and "B", it must predict "C".
• Position 2 (reads ["A", "B", "C"]): Having attended to "A", "B", and "C", it must predict "D".
• Position 3 (reads ["A", "B", "C", "D"]): Having attended to all four tokens, it computes logits to predict what follows "D". But our input has no subsequent token — there is no "E" in this sequence. So position 3 has no target in this training example.

This reveals the fundamental alignment of autoregressive training: An input sequence of length $T$ produces $T$ prediction vectors, but provides ground-truth next-token targets for only the first $T-1$ positions. The prediction at the final position cannot be scored because we do not have its label.

To train the model, we align predictions with their targets:

Predictions (positions 0, 1, 2)  ←→  Targets (tokens 1, 2, 3 = "B", "C", "D")

In tensor terms, this means:

• The inputs evaluated are input_ids[:, :-1] → ["A", "B", "C"] (the first T - 1 tokens)
• The corresponding targets are labels[:, 1:] → ["B", "C", "D"] (the last T - 1 tokens)

This makes the loss-side shift immediately intuitive: we drop the last logit (no next token to predict) and drop the first label (no preceding token to predict it).

How padding affects this example

Now imagine batching this sequence with a shorter one: ["A", "B", "<pad>", "<pad>"] (length 2, padded to 4).

• Position 0 (reads ["A"]): predicts "B" ✓
• Position 1 (reads ["A", "B"]): the next real token would be at position 2, but position 2 is <pad>. If we penalize the model for predicting <pad> here, we teach it to emit padding tokens. So we mask this target with -100 (CrossEntropyLoss(ignore_index=-100)).
• Positions 2 and 3: are padding positions — their logits are computed but their targets are -100, so they contribute zero loss and zero gradient.

This gives us two distinct masking mechanisms that protect different phases:

Mechanism	Phase	What it prevents
Attention mask (key padding mask)	Forward pass	Real queries attending to padding keys — would contaminate representations with <pad> embeddings
Loss ignore index (labels = -100)	Backward pass	Padding queries from updating weights — would teach the model to predict <pad>

The attention mask zeroes out pad keys during attention so real tokens never attend to padding tokens. The loss ignore index zeroes out loss and gradient contributions so padding queries do not update model weights.

Both are necessary and neither replaces the other.

Spotting the double-shift bug

A silent but catastrophic bug occurs when both the data collator and the loss function shift the targets:

• Collator shifts manually: labels = input_ids[:, 1:] (or similar)
• Loss function shifts again: shift_labels = labels[:, 1:]

With our ["A", "B", "C", "D"] example:

• After collator shift: labels = ["B", "C", "D"]
• After loss shift: shift_labels = ["C", "D"]
• Position 0 logits (which saw only "A") are evaluated against "C" instead of "B"!

The model is penalized for failing to predict token $t+2$ from prefix $0..t$ without ever having seen token $t+1$. Loss starts near $\ln(V)$ and never drops.

Diagnostic check: On a toy sequence [10, 20, 30, 40], assert that shift_labels[0, 0] == input_ids[0, 1] (i.e., 20). If it equals 30, you have a double shift.

Canonical rule: Never shift in the collator. Collate input_ids and labels with identical shapes [B, T], replace padding in labels with -100, and let the loss function execute the single canonical shift.

---

Now imagine training on batches of sentences with different lengths—some 3 tokens long, some 512 tokens long. How do we batch them together without padding tokens leaking into the attention layers or penalizing the loss?

In this lesson, you will trace training sequence preparation from input–target alignment through padded batches to sequence packing:

• We will construct the three essential training tensors: input_ids, attention_mask, and labels.
• We will clearly separate what the attention key padding mask controls (the forward pass) from what label masking (-100) controls (the backward loss).
• We will walk through diagnosing real-world bugs, including the catastrophic ‘double-shift’ bug that stalls pretraining.

The one-sentence version

A causal LM batch consists of three [B, T] tensors — input_ids, attention_mask, and labels — where labels is a clone of input_ids with pad positions set to -100; the next-token shift logits[:, :-1] vs labels[:, 1:] happens inside the loss function, not in data preparation.

Everything else in this lesson makes this convention concrete.

Lab: Run the companion CPU exercise at /labs/track1/02-tiny-decoder/ (padding, loss-side shift, label masking).

See the preparation pipeline first: a two-sequence batch preview

Before examining padding masks and sequence packing in detail, trace one concrete batch of variable-length sequences to see the exact tensor shapes and alignments produced:

• Small input: Two tokenized sequences of unequal length:
  • Sequence A: "a b c" → [1, 2, 3] (length 3)
  • Sequence B: "a a" → [1, 1] (length 2)
  • Batch parameters: B=2, pad_token_id=0, target batch length T=5.
• Action:
  1. Right-pad both sequences to length 5 with pad_token_id=0.
  2. Build attention_mask (1 for real tokens, 0 for padding).
  3. Clone input_ids to labels and set all pad positions to -100 (without manual shifting in the collator).
  4. Form the loss inputs by slicing model output logits[:, :-1, :] against targets labels[:, 1:].
• Observable result:
  • input_ids [2, 5]:

    [[1, 2, 3, 0, 0],
     [1, 1, 0, 0, 0]]

  • attention_mask [2, 5]:

    [[1, 1, 1, 0, 0],
     [1, 1, 0, 0, 0]]

  • labels [2, 5]:

    [[1, 2, 3, -100, -100],
     [1, 1, -100, -100, -100]]

  • Aligned loss targets labels[:, 1:] [2, 4]:

    [[2, 3, -100, -100],
     [1, -100, -100, -100]]

  • Evaluation: The loss evaluates exactly 3 active positions (logits[0, 0] against target 2, logits[0, 1] against target 3, and logits[1, 0] against target 1). All -100 positions are ignored by cross-entropy loss, ensuring padded steps neither leak gradients nor penalize the model.

Convention used throughout this lesson

We follow the Hugging Face transformers convention (used by DataCollatorForLanguageModeling with mlm=False and the built-in causal LM loss). This is the standard convention for decoder-only causal LM training.

Tensor	Shape	Content
input_ids	[B, T]	Right-padded token IDs; pad_token_id at pad positions
attention_mask	[B, T]	1 for real tokens, 0 for padding (key padding mask)
labels	[B, T]	Copy of input_ids, with pad positions set to -100. No manual shift.
Model logits	[B, T, V]	Output of the transformer for every position
Loss inputs	—	Computed as logits[:, :-1, :] (shape [B, T-1, V]) vs labels[:, 1:] (shape [B, T-1]) inside the loss function

Why this convention: It is the default in transformers, it avoids an extra data-collation step, and it makes the collator’s job a simple clone-and-mask. The shift is a loss-computation detail, not a data-preparation detail.

1. The next-token shift (loss-side, not data-side)

A decoder-only model produces logits of shape [B, T, V] for every position. For next-token prediction, the logits at position t score the token at position t+1. The loss function uses:

• logits[:, :-1, :] → shape [B, T-1, V] (drop the last position, which has no next token to predict)
• labels[:, 1:] → shape [B, T-1] (drop the first position, which has no preceding token to predict it)

This shift is the defining operation of causal language modelling. The transformer-architecture lesson shows the logits shape [2, 4, 16] becoming shifted logits [2, 3, 16] and targets [2, 3] inside the loss. The data collator does not perform this shift.

Key point: Do not shift labels in the collator. The HF DataCollatorForLanguageModeling(mlm=False) simply does labels = input_ids.clone(); labels[labels == pad_token_id] = -100. The shift is applied by the loss function.

2. Padding variable-length sequences

Sequences in a batch rarely have the same length. The standard approach is right-padding to the longest sequence in the batch (or to max_length):

# DataCollatorForLanguageModeling (mlm=False) logic, simplified
input_ids = pad_batch(tokenized_sequences, pad_token_id, padding_side="right")
attention_mask = (input_ids != pad_token_id).long()
labels = input_ids.clone()
labels[labels == pad_token_id] = -100

• input_ids: padded with tokenizer.pad_token_id. GPT-2 has no pad token: its id 50256 is <|endoftext|>, the EOS token. Many pipelines therefore set pad_token_id = eos_token_id, which quietly breaks this recipe, because labels[labels == pad_token_id] = -100 then erases every real EOS target and the key padding mask hides every EOS key. If you reuse EOS as pad, mask pad positions by position (from the attention mask) rather than by id, or add a dedicated pad token.
• attention_mask: 1 for real tokens, 0 for padding. Shape [B, T_padded]. This is the key padding mask passed to the model; the model expands it to [B, 1, T, T] and combines with a causal mask.
• labels: copy of input_ids, then pad positions set to -100. Shape [B, T_padded].

PyTorch CrossEntropyLoss(ignore_index=-100) ignores any target position equal to -100; its gradient is zero and it does not contribute to the mean loss.

Why two guards? Attention mask vs. loss label masking

Engineers learning LLM pretraining often ask: ‘If we mask the loss with labels = -100, why do we need an attention mask? And if we have an attention mask, why do we need labels = -100?’

They protect two completely different phases of computation:

1. The Attention Mask (Key Padding Mask) controls the FORWARD pass:

   • In self-attention, each query position calculates dot products with all key positions across the sequence.
   • If keys at positions 3 and 4 are <pad> tokens and you do not mask them, position 2 ('c') will compute attention scores with those padding keys. Softmax will allocate non-zero probability mass to them!
   • As a result, position 2’s representation will be contaminated by the embedding of <pad>, degrading its ability to predict the next real token.
   • Rule: The attention mask prevents real queries from attending to padding keys.

2. The Loss Label (ignore_index = -100) controls the BACKWARD pass:

   • Even when padding keys are masked out, padding queries (positions 3 and 4) still pass through the network and compute logits [B, T, V].
   • If their targets remained pad_token_id = 0, the cross-entropy loss would penalize the model whenever it failed to predict <pad> at the end of a sentence! The model would waste valuable parameter capacity learning to predict <pad> tokens.
   • Setting labels[pad] = -100 tells PyTorch to skip those loss terms completely. Their loss is 0, their gradient is 0, and they contribute nothing to weight updates.
   • Rule: The loss label prevents padding queries from penalizing the model or updating weights.

Causal attention mask (key padding mask broadcast)

The model receives attention_mask as a key padding mask of shape [B, T] (1=real, 0=pad). Internally, the model constructs a combined mask:

combined_mask = causal_mask & attention_mask[:, None, None, :]

• causal_mask: lower-triangular matrix of ones (position t attends to 0..t).
• attention_mask[:, None, None, :]: broadcasts the [B, T] padding mask over query and head dimensions, zeroing out keys that are padding.

Result: Real positions attend only to real positions at or before themselves. Pad key positions are masked out. Pad query positions are not explicitly zeroed in this mask (they still attend to real keys up to their position), but their loss contribution is zero because labels[pad] = -100. For maximum efficiency, variable-length kernels (FlashAttention varlen, FlexAttention) exclude pad queries entirely via cu_seqlens.

3. Causal attention mask (detailed)

The combined mask for a single batch item with T=5 and key padding mask [1,1,1,0,0]:

Causal (tril):         Key padding mask (broadcast to keys):   Combined:
[[1,0,0,0,0],          [[1,1,1,0,0],                           [[1,0,0,0,0],
 [1,1,0,0,0],    &       [1,1,1,0,0],          =              [1,1,0,0,0],
 [1,1,1,0,0],           [1,1,1,0,0],                        [1,1,1,0,0],
 [1,1,1,1,0],           [1,1,1,0,0],                        [1,1,1,0,0],
 [1,1,1,1,1]]           [1,1,1,0,0]]                        [1,1,1,0,0]]

• Rows 0-2 (real queries): attend to real keys at or before themselves (cols 0-2). Cols 3-4 (pad keys) are zeroed by the key padding mask.
• Rows 3-4 (pad queries): the combined mask keeps their causal view of the real keys (cols 0-2). The key padding mask zeroes pad keys; it does not zero pad queries. Those rows hold no gradient because labels=-100 removes pad positions from the loss, and varlen kernels skip them entirely.

For key padding mask [1,1,0,0,0] (Seq B example):

Combined mask:
[[1,0,0,0,0],
 [1,1,0,0,0],
 [1,1,0,0,0],
 [1,1,0,0,0],
 [1,1,0,0,0]]

• Rows 0-1 (real queries): attend to real keys at or before themselves; cols 2-4 (pad keys) are zeroed.
• Rows 2-4 (pad queries): attend to the real keys in cols 0-1. Only pad keys are masked; the loss is zero via labels=-100.

The shift and the mask in one picture

Shift and combined mask

Panel A: the loss drops the last logit and the first label. Panel B: the combined causal and key-padding mask for a 5-position row with 3 real tokens. The panels sit side by side on wide screens and stack on narrow ones.

Read this diagram as text

Text equivalent. Panel A repeats the row from the worked example. input_ids is 1, 2, 3, 0, 0 and labels is 1, 2, 3, -100, -100. The loss keeps logits positions 0 to 3 and label positions 1 to 4, so logit 0 is scored against label 2, logit 1 against label 3, and logits 2 and 3 against -100, which contributes no loss. Each arrow between the two bottom rows is one cross-entropy term, so this row contributes two terms. Panel B is a five by five grid. A cell is 1 only when the key is at or before the query and the key is a real token. The upper triangle is 0 because those keys are in the future. Columns 3 and 4 are 0 for every row because they are padding keys. Rows 3 and 4 are padding queries: they still read keys 0 to 2, and their loss disappears because their labels are -100.

4. Sequence packing (concatenating multiple examples)

Instead of padding, multiple short sequences are concatenated into one long sequence up to the model’s context length, separated by EOS tokens:

seq1 + [EOS] + seq2 + [EOS] + seq3 + [EOS] ...

Critical additions to prevent cross-sequence contamination:

1. Modified causal mask (block-diagonal): Standard lower-triangular mask plus blocking attention from any token to tokens in previous sequences. The Hugging Face blog (Lukas, 2024) shows the algorithm: find EOS indices, compute sequence lengths, build a repeated_idx tensor marking the last valid position for each token, then mask = tril() & (mask_indices <= repeated_idx). This produces a block-diagonal causal mask where each block corresponds to one sequence.

2. Reset position IDs: Each sequence starts at position 0. pos_ids = arange(T) - repeat_interleave(sequence_start_indices, sequence_lengths). For RoPE, this is required so relative positions are computed within each sequence, not across the packed context.

3. Labels at sequence boundaries (Convention A — adopted in this lesson): Do not mask the first token of a packed sequence. The EOS token in the context serves as the separator; the model learns to predict the first token of the next sequence given the EOS. Only pad positions are -100.

   Convention B (some implementations): Mask the first token of each packed sequence (set its label to -100) because there is no valid intra-sequence predecessor.

   Note: This lesson uses Convention A throughout. Convention B is defensible: it avoids training the model to predict a token whose only context is another document. The cost of A is one extra, weakly conditioned target per packed boundary; the cost of B is a slightly smaller effective batch. Pick one deliberately and keep it consistent between data preparation and evaluation.

NVIDIA NeMo (2024) uses FlashAttention’s variable-length kernels (cu_seqlens) to avoid the O((∑s_i)²) attention cost, keeping complexity at O(∑s_i²).

Packed example (Convention A)

Setup: Two sequences from Section 2, packed into context length 8 with EOS.

• Seq A: [1, 2, 3] (len 3) + EOS 4
• Seq B: [1, 1] (len 2) + EOS 4
• Concatenated: [1, 2, 3, 4, 1, 1, 4, 0] (pad 0 fills remaining 8 positions)
• EOS indices: 3, 6
• Sequence lengths: 4 (positions 0-3), 3 (positions 4-6)
• Position IDs (reset at each sequence): [0, 1, 2, 3, 0, 1, 2, -1] (pad position ID = -1, unused)

Modified causal mask (block-diagonal):

• Standard tril() for T=8.
• Block 1 (positions 0-3): full lower-triangular.
• Block 2 (positions 4-6): full lower-triangular.
• Cross-block: no position in Block 2 may attend to any position in Block 1, and no position in Block 1 may attend forward into Block 2. Each block sees only its own earlier positions.
• Position 7 (pad): all zeros here because the block rule caps the row at the last position of its own sequence, and position 7 belongs to no sequence. This is the packed mask, not the §3 padding mask, where pad queries keep their view of real keys.

Labels (Convention A): Clone of input_ids, pad → -100: labels = [1, 2, 3, 4, 1, 1, 4, -100]

After loss shift labels[:, 1:] → [2, 3, 4, 1, 1, 4, -100]:

• Valid targets: positions 0-5 (targets 2, 3, 4, 1, 1, 4). Position 6 is -100 → ignored.
• logits[3] (EOS of seq1) → labels[4] = 1 (first token of seq2, Convention A: valid target)
• logits[6] (EOS of seq2) → labels[7] = -100 (ignored)

5. Fully worked numeric example

Setup (toy vocabulary, tiny model):

• Vocabulary: a=1, b=2, c=3, <pad>=0, <eos>=4
• Model: d=4, V=5, max_len=8
• Batch size B=2, two sequences of different lengths.

Sequences (post-tokenisation, pre-padding):

• Seq A: "a b c" → IDs [1, 2, 3] (length 3)
• Seq B: "a a" → IDs [1, 1] (length 2)

Step 1: Pad to max length in batch (5) — right padding, pad_id=0

	input_ids	attention_mask
A	[1, 2, 3, 0, 0]	[1, 1, 1, 0, 0]
B	[1, 1, 0, 0, 0]	[1, 1, 0, 0, 0]

Step 2: Create labels (clone input_ids, mask pad positions to -100) — NO manual shift

	labels
A	[1, 2, 3, -100, -100]
B	[1, 1, -100, -100, -100]

Detail: labels has the same shape as input_ids. The loss function will apply logits[:, :-1] vs labels[:, 1:] internally.

Step 3: Model forward pass → logits [B, T, V] = [2, 5, 5]

For loss we use:

• logits[:, :-1, :] → shape [2, 4, 5] (drop last position)
• labels[:, 1:] → shape [2, 4] (drop first position)

PyTorch CrossEntropyLoss expects input [N, C] or [N, C, ...] and target [N] or [N, ...]. Flattened:

• logits[:, :-1, :].reshape(-1, V) → [8, 5]
• labels[:, 1:].reshape(-1) → [8]

Step 4: Causal + padding attention mask

For Seq A (T=5, key padding mask [1,1,1,0,0]):

Combined mask (what position t can attend to):
[[1,0,0,0,0],
 [1,1,0,0,0],
 [1,1,1,0,0],
 [1,1,1,0,0],
 [1,1,1,0,0]]

• Rows 0-2 (real queries): attend to real keys at or before themselves. Cols 3-4 (pad keys) zeroed.
• Rows 3-4 (pad queries): they still attend to the real keys at or before them. The mask zeroes pad keys, not pad queries; labels=-100 removes these rows from the loss.

For Seq B (T=5, key padding mask [1,1,0,0,0]):

Combined mask:
[[1,0,0,0,0],
 [1,1,0,0,0],
 [1,1,0,0,0],
 [1,1,0,0,0],
 [1,1,0,0,0]]

• Rows 0-1 (real queries): attend to real keys at or before themselves; cols 2-4 zeroed.
• Rows 2-4 (pad queries): attend to real keys in cols 0-1; their loss is zero via labels=-100.

Step 5: Loss computation (conceptual)

For Seq A, after labels[:, 1:]:

• labels[0, 1:] = [2, 3, -100, -100]
• Only positions 0 and 1 are valid (targets 2 and 3). Positions 2,3 are -100 → ignored.
• loss = mean(CE(logits[0,0], 2), CE(logits[0,1], 3)).

For Seq B:

• labels[1, 1:] = [1, -100, -100, -100]
• Only position 0 is valid (target 1).
• loss = CE(logits[1,0], 1).

Packed version of the same two sequences (Convention A, context=8):

Concatenated: [1, 2, 3, 4, 1, 1, 4, 0] (EOS=4, pad=0 to fill 8)

• EOS indices: 3, 6
• Sequence lengths: 4, 3
• Position IDs: [0, 1, 2, 3, 0, 1, 2, -1]

Modified causal mask (block-diagonal): standard tril(), Block 1 (pos 0-3) full lower-triangular, Block 2 (pos 4-6) full lower-triangular, cross-block attention blocked, pos 7 (pad) all zeros.

Labels (Convention A): [1, 2, 3, 4, 1, 1, 4, -100] After loss shift: labels[:, 1:] → [2, 3, 4, 1, 1, 4, -100]

• Valid targets for logits positions 0-5: 2, 3, 4, 1, 1, 4.
• logits[3] (EOS of seq1) → labels[4] = 1 (first token of seq2, valid under Convention A)
• Position 6 (-100) ignored.

6. Failure modes readers must learn to detect

Failure mode	Symptom	Root cause	Test
Future-token leak	Validation perplexity unrealistically low; model “cheats”	Causal mask not applied, or applied incorrectly (e.g., upper-triangular, or missing cross-sequence blocks in packing)	Assert attn_mask[i, j] == 0 for all j > i; for packed, assert no attention across EOS boundaries
Misaligned labels	Loss doesn’t decrease; perplexity near V (random)	Shift off by one (e.g. comparing logits[:, 1:] with labels[:, :-1], which drops the loss-side shift), or a collator that shifts before the loss does	Unit test on a known sequence: for every valid t, the loss must pair logits[t] with input_ids[t+1]
Padding keys not masked	Wasted compute; garbage states at pad positions	Attention mask doesn’t zero pad keys, so real queries attend to pad tokens	Assert attention_mask[:, pad_positions] == 0 (key padding mask); verify labels[pad_positions] == -100
Fully masked attention row	NaN loss or NaN gradients on the first step	Every key in a row is masked, so the softmax normaliser is zero (e.g. a pad query with no real key before it, or an all-pad row)	Assert every row of the combined mask has at least one unmasked key, or drop all-pad rows before the model
Packed cross-contamination	Model learns spurious correlations between unrelated sequences	Packed attention mask missing the “block previous sequences” logic	Assert attn_mask[seq2_start:, :seq1_end] == 0
Position ID drift in packing	Quality degrades on long packed sequences	Position IDs not reset at sequence boundaries (RoPE)	Assert pos_ids[seq_start] == 0 for each sequence; check pos_ids strictly increasing within each sequence

Hands-on diagnosis: two bugs that break pretraining

Bug 1: The Catastrophic Double-Shift

A common mistake when writing custom data pipelines is attempting to prepare “next-token labels” inside the PyTorch Dataset or collate_fn:

# BUGGY collator:
def collate_fn(batch):
    input_ids = pad_sequence([item['input_ids'] for item in batch], pad_value=0)
    # Developer manually shifts labels:
    labels = input_ids[:, 1:].clone()
    labels[labels == 0] = -100
    return {'input_ids': input_ids[:, :-1], 'labels': labels}

The developer then passes this dictionary to Hugging Face’s AutoModelForCausalLM or a standard causal loss wrapper. Inside transformers/loss/loss_utils.py, the model’s loss function applies its standard next-token shift:

# Inside standard CausalLM loss computation:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()  # SHIFTS A SECOND TIME!
loss = F.cross_entropy(shift_logits.view(-1, V), shift_labels.view(-1))

What happens to the math?

• Position 0’s logits (which saw only token 0) are evaluated against shift_labels[0].
• Because labels were shifted in collation and in the loss function, shift_labels[0] corresponds to token 2!
• The model is penalized for failing to predict token t + 2 from prefix 0..t without ever having seen token t + 1.
• Symptom: The loss begins around ln(V) (for GPT-2 with V = 50,257, ln(50257) ≈ 10.82) and never drops below ~9.5. Perplexity remains near random chance, and the model never learns grammar or vocabulary.
• Diagnostic assertion:

# Unit test to catch shift bugs on a toy sequence [10, 20, 30, 40]:
test_input = torch.tensor([[10, 20, 30, 40]])
# With canonical loss-side shift:
# Logit at position 0 must evaluate against target 20:
assert (
    shift_labels[0, 0].item() == test_input[0, 1].item()
), f"Alignment bug! Expected target {test_input[0, 1].item()}, got {shift_labels[0, 0].item()}"

• The canonical rule: Never shift in the collator. Collate input_ids and labels with identical shapes [B, T], replace padding in labels with -100, and let the loss function execute the shift.

Bug 2: Pad Token Leakage into the Loss

A developer configures a causal LM with pad_token_id = 0, but forgets to replace padding positions with -100 in labels:

# BUGGY collator:
labels = input_ids.clone()  # Pad tokens remain 0!

• Symptom: Training loss plummets suspiciously fast to near zero. But when the model is prompted to generate text, it immediately outputs an endless stream of <pad> <pad> <pad>....
• Root cause: In typical pretraining batches with fixed context lengths, 60–80% of the tokens in variable-length batches are padding. Because <pad> is given a valid label ID (0), the easiest way for the model to minimize cross-entropy loss is to predict <pad> at every position.
• The fix: Replace pad tokens by position: labels[input_ids == pad_token_id] = -100. Verify with:

assert (labels == pad_token_id).sum() == 0, "Found unmasked pad tokens in training labels!"

7. One practical engineering decision: packing vs. padding + loss masking

Scenario: Small corpus (e.g., 10k sequences), average length 512 tokens, max context 2048, batch size 8. Assume padding to max_length=2048 (fixed-length padding, not dynamic batch padding).

	Padding + loss masking	Sequence packing
Tokens/step	8 × 2048 = 16,384 slots, of which 8 × 512 = 4,096 are real. Padding waste is 75%.	8 × 2048 = 16,384 slots, of which roughly 12,800-14,300 are real: packing removes the per-sequence padding, but each row keeps one partial tail sequence and one EOS per document, so expect about 80-88% utilisation, not 100%.
Implementation complexity	Low (HF DataCollatorForLanguageModeling handles it)	High (custom collator, block-diagonal mask, position IDs, boundary label decision)
Cross-sequence risk	None (sequences independent)	Must verify mask blocks previous sequences; bug = silent quality loss
Position embedding compatibility	Works with learned, sinusoidal, RoPE	RoPE requires position IDs reset; learned/sinusoidal need max_position_embeddings ≥ pack_length
Debugging	Straightforward (per-sequence loss visible)	Harder (loss aggregated across packed sequences)
Convergence	Baseline	No assumed difference. Measure it: packing changes the sequence-length distribution, and any effect on validation loss is empirical.

Decision: For this scenario, use padding + loss masking. The corpus is small, the padding overhead (75%) is acceptable for a first training run, and the implementation is standard and well-tested. Switch to packing only when:

• Fixed-max_length padding waste is high for your configuration, not as a general threshold. It is a property of max_length versus your length distribution: padding to 2048 when the average document is 512 wastes 75% by construction, while dynamic batch padding wastes far less. Compute the ratio for your own lengths before switching.
• You have validated the block-diagonal attention mask and position ID logic with unit tests
• You are using a framework (NeMo, Megatron, PyTorch FlexAttention) that provides verified packed kernels (cu_seqlens / document masking)

Trade-off statement: Packing trades engineering complexity and debuggability for compute efficiency. The break-even point depends on length distribution skew and framework support. Do not pack “by default” without a measured padding waste problem.

8. Prerequisite links (earlier lessons that must be complete)

1. transformer-architecture (llm-internals/transformer-architecture): Provides the [B, T, V] logits shape trace, the causal mask concept, and the loss-side shift from logits [B, T, V] to shifted logits [B, T-1, V] and targets [B, T-1]. This lesson assumes the reader knows why we shift in the loss.
2. tokenisation-and-embeddings (llm-internals/tokenisation-and-embeddings): Provides the token ID vocabulary, special tokens (pad, EOS), embedding lookup producing [B, T, d], and the output projection to logits. This lesson assumes the reader knows where pad_token_id and eos_token_id come from.

Both are in the same module (architecture-and-inputs) and must be published before this lesson.

---

9. Where this lesson stops

Topic	Deferred to / Covered elsewhere
Optimiser choice (AdamW, etc.), learning rate schedules, weight decay	training-and-debugging (order 10)
Training loop, gradient accumulation, mixed precision, checkpointing	training-and-debugging
Evaluation metrics (perplexity, eval harness), validation splits	training-and-debugging / evaluating-reasoning
Distributed training (DDP, FSDP, tensor parallel)	inference-optimisation track / advanced modules
Data loading, dataset streaming, tokenisation on-the-fly	Out of scope for this curriculum (assume pre-tokenised datasets)
Masked language modelling (BERT-style)	Not covered; this curriculum is decoder-only
Instruction tuning / chat template formatting	fine-tuning track
RLHF / preference optimisation	reinforcement-learning track

This lesson ends with the tensors (input_ids, attention_mask, labels) ready for model(**batch). The next lesson picks up at loss = model(**batch).loss: the optimiser, the schedule, and the debugging loop.

---

Module 1 synthesis: connecting the three lessons

With this third lesson complete, you can now connect the entire pipeline of how a decoder-only language model processes text, generates continuations, and trains on data:

1. Lesson 1 — Transformer Architecture: A decoder-only model maintains a residual stream of width d across $N$ identical blocks. Causal self-attention contextualizes representations by allowing tokens to mix across time, while MLPs act as per-position non-linear feature extractors and associative memories. The final hidden state is projected to vocabulary logits. During generation, only the last position’s logits are sampled.
2. Lesson 2 — Tokenisation and Embeddings: Language models do not take text or ASCII bytes directly; text is chunked into subwords via byte-level BPE, which assigns each token an arbitrary vocabulary integer ID. Embedding table $E$ looks up a continuous static vector of width d for each ID. At the other end of the network, the output projection computes the dot product between the contextual hidden vector and vocabulary prototype vectors (stored as [V, d] in PyTorch) to produce logits.
3. Lesson 3 — Preparing Training Sequences: To train the network, input sequences of length $T$ produce $T$ prediction vectors, but only the first $T-1$ positions have ground-truth targets. The collator right-pads variable-length sequences to form input_ids [B, T] and clones them into labels [B, T]. The attention key padding mask blocks real tokens from attending to padding keys in the forward pass, while -100 label masking removes padding queries from the loss in the backward pass. The loss function executes the single canonical shift: logits[:, :-1] versus labels[:, 1:].

In Module 2, we leave high-level component blocks behind and open up the mathematical heart of the transformer: deriving scaled dot-product attention, queries, keys, and values from scratch in PyTorch.

---

10. Sources

Source	Relevance
PyTorch CrossEntropyLoss docs — ignore_index parameter	Defines -100 as default ignore index; confirms padded labels are excluded from loss and gradient.
HF DataCollatorForLanguageModeling (source: transformers/src/transformers/data/data_collator.py, torch_call, mlm=False branch)	Shows labels = input_ids.clone(); labels[labels == pad_token_id] = -100 — the canonical padding + loss-masking pattern for causal LM.
HF causal LM loss (source: transformers/src/transformers/loss/loss_utils.py, ForCausalLMLoss)	Shows the loss-side shift: shift_logits = logits[..., :-1, :]; shift_labels = labels[..., 1:]; loss = cross_entropy(shift_logits, shift_labels).
Hugging Face blog: “Efficient LLM Pretraining: Packed Sequences and Masked Attention” (Lukas, 2024)	Complete algorithm for packed causal mask (EOS-index → sequence lengths → repeated_idx → tril() & (mask_indices <= repeated_idx)) and position ID reset.
NVIDIA NeMo Framework User Guide: Sequence Packing (2024)	Complexity analysis: naive packing O((∑s_i)²) vs. FlashAttention variable-length O(∑s_i²); 10× FLOPs / 6× time claims on Llama 7B + Dolly.
GPT-3 paper (Brown et al., 2020) — “Language Models are Few-Shot Learners”	Uses EOS as sequence separator in packed pre-training; does not mask first token of next sequence (Convention A).
Raffel et al., “Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer” (T5, 2020) — arXiv:1910.10683	Introduced sequence packing for encoder-decoder pre-training; cited as origin of modern packing.
PyTorch FlexAttention blog (2024) — “Document Masking / Jagged Sequences”	Official PyTorch API for packed attention without materialising full T×T mask; confirms packing is a supported pattern.
Transformer-architecture lesson (published) — shift logic	logits: [B, T, V] → shifted logits [B, T-1, V], targets [B, T-1], per-token loss [B, T-1] (inside loss).
Tokenisation-and-embeddings lesson (published, llm-internals/tokenisation-and-embeddings)	Where pad_token_id and eos_token_id come from, and how ids map through the embedding table.
Hugging Face PreTrainedTokenizer docs — pad_token_id, eos_token_id	Confirms that GPT-2 ships without a pad token and that reusing EOS as pad changes which positions are ignored.

---

Module exit check: sequence preparation and alignment

1. 01Given a batch of 3 sequences with IDs [[1, 2], [3], [4, 5, 6]] and pad token 0, write the input_ids, attention_mask, and labels tensors after dynamic batch padding (pad to longest in batch). Then write the slices logits[:, :-1] and labels[:, 1:] that the loss function uses, and count how many active cross-entropy loss terms are evaluated.

   Show a hint

   Longest sequence length is 3, so T=3. input_ids = [[1, 2, 0], [3, 0, 0], [4, 5, 6]] (shape [3, 3]); attention_mask = [[1, 1, 0], [1, 0, 0], [1, 1, 1]] (shape [3, 3]); labels = [[1, 2, -100], [3, -100, -100], [4, 5, 6]] (shape [3, 3]). Slicing for loss drops the last logit and first label: logits[:, :-1, :] has shape [3, 2, V] and labels[:, 1:] has shape [3, 2] = [[2, -100], [-100, -100], [5, 6]]. Active loss terms: 1 from sequence 0 (target 2), 0 from sequence 1 (single-token input cannot train next-token prediction), and 2 from sequence 2 (targets 5 and 6). Exactly 3 active terms are evaluated.

2. 02For the packed version of the same 3 sequences (context length 10, EOS=7, pad=0), write the concatenated input_ids, the reset position IDs, and the labels tensor under Convention A. Does the EOS token have an active loss target?

   Show a hint

   Sequences with EOS: [1, 2, 7], [3, 7], [4, 5, 6, 7]. Concatenated input_ids padded to 10: [1, 2, 7, 3, 7, 4, 5, 6, 7, 0]. Position IDs reset to 0 at each sequence boundary: [0, 1, 2, 0, 1, 0, 1, 2, 3, -1]. labels = [1, 2, 7, 3, 7, 4, 5, 6, 7, -100]. After loss shift labels[1:]: [2, 7, 3, 7, 4, 5, 6, 7, -100]. Yes: under Convention A, EOS tokens do have active targets. Logits at position 2 (EOS of seq 1) predict 3 (first token of seq 2), and logits at position 4 (EOS of seq 2) predict 4 (first token of seq 3). Only the pad position at index 9 (-100) is ignored.

3. 03Identify which of the following label configurations causes training failure and explain the failure mechanism: (a) labels = input_ids with no pad masking; (b) labels = input_ids with pad set to -100 but no loss-side shift; (c) labels cloned from input_ids with pad set to -100 and loss comparing logits[:, :-1] vs labels[:, 1:].

   Show a hint

   Configuration (a) leaks unmasked pad tokens into the loss, so the model wastes capacity learning to predict padding. Configuration (b) fails because comparing unshifted logits with labels evaluates predicting the current token rather than the next token, breaking the autoregressive objective. Configuration (c) is the correct canonical Hugging Face setup: pad tokens contribute zero loss via ignore_index=-100, and the loss-side shift accurately pairs position t logits with target t+1.

4. 04Module 1 Capstone Connection: A user prompts a model with "The capital of France is" (5 tokens). (a) How does the tokenizer turn this string into tensor input_ids? (b) Through the transformer blocks, does the tensor shape change? What changes semantically? (c) During generation, which logit vector is sampled to produce the word "Paris"? (d) During training on this 5-token sequence, how many next-token loss terms are calculated, and which tokens serve as input and target for the first loss term?

   Show a hint

   (a) The tokenizer splits the string with regex (absorbing leading spaces into " capital", " of", " France", " is") and maps subwords to vocabulary IDs, producing input_ids of shape [1, 5]. (b) Shape remains strictly invariant at [1, 5, d] across the residual stream. Semantically, the representation at position 4 (" is") transforms from a static word embedding into a contextual summary meaning "the entity whose capital is sought is France". (c) The last vector logits[:, -1, :] of shape [1, V] is sampled. (d) Exactly 4 loss terms (T - 1 = 5 - 1 = 4). For the first term, position 0 ("The") is the input and token 1 (" capital") is the target.

---

Related Lessons

Related lessons

• LLM Internals & Architecture · Transformer architecture and model inputsTransformer Architecture Explained: Trace a Decoder-Only LLM14 min
• LLM Internals & Architecture · Transformer architecture and model inputsTokenisation and Embeddings: From Text to Tensors12 min