Training and Debugging: Next-Token Loss and Optimiser Steps

The one-sentence version

Training a causal decoder-only transformer means computing next-token cross-entropy loss against shifted logits, updating parameters with an optimiser step, and using gradient accumulation when the effective batch size exceeds the hardware limit; perplexity = exp(loss) measures prediction confidence.

Everything else in this lesson makes this concrete.

Lab: Run the companion CPU exercise at /labs/track1/02-tiny-decoder/ (loss computation, optimizer step, gradient accumulation, overfit test).

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.

HF DataCollatorForLanguageModeling convention

The data collator simply does: labels = input_ids.clone(); labels[labels == pad_token_id] = -100 The shift logits[:, :-1, :] vs labels[:, 1:] happens inside the loss function. Do not shift labels in the collator.

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.

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.

3. Cross-entropy loss and perplexity

For a single position t with logits logits_t of shape [V] and target token_id:

p_t = softmax(logits_t)   # [V], probabilities over vocab
loss_t = -log(p_t[target_t])   # scalar cross-entropy

For a full sequence of length T, averaged over T-1 positions (since position T-1 has no next token, and position 0 has no preceding token):

loss = (1 / (T - 1)) * sum_{t=0}^{T-2} -log(softmax(logits_t)[token_{t+1}])

Perplexity

perplexity = exp(loss)

• Interpretation: the effective vocabulary size the model seems to be choosing from.
• A random-guess model on a V-token vocab has perplexity ≈ V.
• An ideal model has perplexity ≈ 1.
• Lower is better; no units.

4. Optimizer step (SGD / Adam)

The optimizer computes parameter updates based on gradients. For SGD with learning rate lr:

θ_new = θ_old - lr * g

where g = ∂loss/∂θ is the gradient of the loss with respect to parameter θ.

Worked example: one SGD step

A minimal runnable example using PyTorch nn.Linear:

import torch
import torch.nn as nn
import torch.nn.functional as F

d, V = 4, 6
torch.manual_seed(42)
x = torch.randn(1, d)
target = torch.tensor([2])  # target class index

model = nn.Linear(d, V, bias=True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Record parameter before
w_before = model.weight.clone()

logits = model(x)
loss = F.cross_entropy(logits, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()

# Parameter has updated
assert not torch.allclose(model.weight, w_before)

After the step, the model parameters have changed in the direction that reduces the loss for the given input-target pair.

5. Gradient accumulation

Hardware has a limited GPU memory for activations. If the desired effective batch size is larger than fits in memory, gradients can be accumulated across multiple micro-batches before an optimizer step.

Without accumulation (one micro-batch)

for batch in dataloader:      # one batch only
    loss = model(batch).loss
    loss.backward()
    optimizer.step()          # update with this batch's gradients
    optimizer.zero_grad()     # reset gradients

With accumulation (N micro-batches)

optimizer.zero_grad()
for i, batch in enumerate(dataloader):
    loss = model(batch).loss / N     # scale loss by 1/N
    loss.backward()                  # accumulate gradients
    if (i + 1) % N == 0:             # after N micro-batches
        optimizer.step()              # update once with accumulated grads
        optimizer.zero_grad()         # reset for next cycle

The division by N ensures the gradient magnitude is approximately the same as a single micro-batch of size N × actual_batch_size, assuming loss averages over the batch.

Why accumulate

• Memory constraint: a single batch of size B may OOM, but N micro-batches of size B/N fit.
• Effective batch size: accumulated_gradient ≈ gradient_from_batch_size_effective.
• This is the standard pattern in GPT-2, GPT-3, and other LLM training code.

6. Overfit diagnostic (tiny-batch run)

A quick way to check that the training pipeline is working is to train on a very small dataset (e.g., 1-8 sequences) for a few steps and verify that loss decreases.

Mini-overfit test

• Dataset: 8 sequences, each T=8 tokens, from a tiny toy corpus (e.g., [“hello world”, “foo bar”, …]).
• Model: GPT-2 small or a 2-layer decoder-only with d=64, 2 heads.
• Optimizer: AdamW, lr=3e-4, one step.
• Expected: loss decreases from its initial value (around ln(V) ≈ 10.82 nats for V=50257, or ≈ 50257 perplexity) to a lower value.
• Pass condition: loss_after < loss_before and no NaN/inf.

If loss stays flat or increases, check:

• Learning rate too small
• Gradients vanishing/exploding (check gradient norms)
• Loss/shape mismatch (logits vs labels dimensions)
• Data collator bug (wrong shift, wrong -100 mask)
• Optimizer state corruption (reset if training longer)

7. Engineering trade-offs

Choice	When to prefer	Trade-off
Loss-side shift	Default for causal LM training	Avoids extra data-collation step; shift inside loss
Padding with EOS-as-pad	Quick experiments, no dedicated pad token	Risk of erasing real EOS targets; mask by position instead
Gradient accumulation	Effective batch size > memory limit	Adds training overhead; gradient scaling must be correct
Loss scaling in accumulation	N micro-batches, accumulate grads	Forgetting /N causes effectively larger lr
Overfit diagnostic	Tiny-batch run, 1–3 steps	Expected loss decrease depends on model/dataset/lr; not a fixed target

8. Failure modes

Failure mode	Symptom	Root cause	Test
Loss NaN	loss = NaN	Gradients exploded; often from too-large lr or unclipped grads	torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
Loss not decreasing	Loss flat or increasing after many steps	lr too small; wrong loss/shape; data collator bug	Reduce lr, check loss.item() history
Perplexity diverges	Perplexity → ∞ or → 0	Loss sign wrong; log probs of negative numbers; labels not -100 for padding	Verify labels[labels != -100] are valid token ids
Gradient accumulation off by N	Effective lr different from expected	Forgetting to scale loss by 1/N in accumulation	Verify loss = model_loss / accumulation_steps
Optimizer step never fires	Training loops forever without updating weights	accumulation_steps not dividing evenly; missing if condition	Ensure step() called inside (i+1) % accumulation == 0
Wrong pad_token_id	Erased EOS targets, wrong padding mask	Reusing EOS as pad without position-based masking	Use dedicated pad token or mask by position

9. Related lessons

Related lessons

• LLM Internals & Architecture · Transformer architecture and model inputsTransformer Architecture Explained: Trace a Decoder-Only LLM14 min
• LLM Internals & Architecture · Attention mechanisms from scratchSelf-Attention: Queries, Keys, and Values14 min
• LLM Internals & Architecture · Attention mechanisms from scratchCausal and Padding Masks: Implementation and Tests15 min
• LLM Internals & Architecture · Attention mechanisms from scratchMulti-Head Attention: Splitting, Merging, and Projecting14 min
• LLM Internals & Architecture · Building and training a language modelPositional Information: Absolute Embeddings and RoPE13 min
• LLM Internals & Architecture · Building and training a language modelInside the Transformer Block: Residuals and Normalisation12 min
• LLM Internals & Architecture · Transformer architecture and model inputsPreparing Training Sequences: Shift, Pad, and Pack15 min

10. Sources

| Source | Relevance | | —————————————————————————————————– | –––––––––––––––––––––––––––––––– | ———————————————————————————————————— | | Vaswani et al., “Attention Is All You Need” (arXiv:1706.03762), §6 | Training procedure, optimizer, loss minimization | | Radford et al., “Language Models are Unsupervised Multitask Learners” (OpenAI Technical Report, 2019) | Training objectives, next-token prediction, dataset construction | PDF | | PyTorch docs, torch.optim | SGD, Adam update rules, zero_grad, step | | Hugging Face DataCollatorForLanguageModeling source | Scaling effective batch size beyond memory limit | | Wolf et al., “Hugging Face Transformers” (2020) | DataCollatorForLanguageModeling, loss-side shift convention |

11. Exercises

Try it yourself

1. 01Implement a training step for a single linear model with cross-entropy loss: forward pass, loss computation, backward pass, and SGD update. Use toy data with V=10, d=5, one input-target pair.

   Show a hint

   model = nn.Linear(d, V, bias=True); x = torch.randn(1, d); target = torch.randint(0, V, (1,)); logits = model(x); loss = F.cross_entropy(logits, target); loss.backward(); with torch.no_grad(): model.weight -= lr * model.weight.grad; model.bias -= lr * model.bias.grad

2. 02Given a loss of 2.5 nats, what is the perplexity? What does this mean in terms of effective vocabulary size?

   Show a hint

   perplexity = exp(2.5) ≈ 12.2. The model's predictions are as if it were uniformly choosing from ~12 tokens out of the full vocabulary.

3. 03Modify the gradient-accumulation loop to accumulate over 4 micro-batches. What change must you make to the loss computation to keep the effective gradient magnitude the same?

   Show a hint

   Divide the loss by 4 before calling backward(): `loss = model(batch).loss / 4`. This ensures the accumulated gradient ≈ gradient from one batch of 4× the size.

4. 04Write a tiny overfit diagnostic: train a 2-layer decoder-only model with d=64 on 8 sequences of T=8 tokens for 3 steps with AdamW(lr=3e-4). Report initial loss, final loss, and whether loss decreased.

   Show a hint

   Use a toy corpus like ["hello world", "foo bar", "baz qux", "abc def", "ghi jkl", "mno pqr", "stu vwx", "yz 12"]. Track loss.item() before and after the first optimizer step.

---