LM5-T3 — Training Loops, Validation and Checkpointing

Learning objectives

After this lesson, you should be able to:

• Write the training loop with the correct order of operations.
• Run a validation pass correctly (eval mode, no gradients) and restore training mode.
• Compute tokens seen, epochs and effective batch size, including gradient accumulation.
• Save and restore a checkpoint that allows training to resume exactly.

The idea: the same five lines, thousands of times

One training step is five operations in a fixed order:

1. Draw a batch: inputs x and targets y of shape (B, n).
2. Forward pass and loss: the model produces logits and the cross-entropy against y gives a number.
3. Zero the old gradients. PyTorch accumulates gradients across backward() calls, so they must be cleared first.
4. Backward pass: loss.backward() fills the gradient of every parameter. Optionally clip the global gradient norm.
5. Optimizer step (and learning-rate update): the weights move.

Evaluation

The training loss is measured on batches the model is learning from, with dropout on, so it is noisy and optimistic. Every N steps, evaluate on validation data: switch the model to evaluation mode (dropout off), disable gradient tracking (torch.no_grad(), which saves memory and time), compute the average loss over several batches, and switch back to training mode. Forgetting to switch back is a classic bug: the rest of the run silently trains without dropout.

What to log

Log the step, the training loss, the validation loss, the learning rate, the gradient norm and the number of tokens processed. The curves of training and validation loss against step are the main instrument for reading a run (Lesson LM5-T4). The gradient norm and learning rate explain sudden changes. Averaging the training loss over the last few steps gives a smoother curve than a single batch.

Checkpoints

A run may take hours or days. A checkpoint is a file with everything needed to continue: the model’s weights, the optimizer’s state (Adam’s moving averages; without them the first steps after resuming are wrong), the step number, the scheduler state and the random-number state. If only the weights are saved you can run the model but you cannot resume training faithfully. Save to a temporary name and rename, so that a crash during saving does not destroy the last good checkpoint.

Gradient accumulation

If a batch of the desired size does not fit in memory, process several smaller micro-batches, summing (averaging) their gradients before one optimizer step. The effective batch is micro-batch size × number of accumulation steps, and the result is the same as one big batch if the loss is divided by the number of accumulation steps.

The mathematics: bookkeeping of tokens and steps

Let the batch have B sequences of n tokens and the run S steps. With A accumulation steps:

B, n

micro-batch size and sequence length.

A

gradient-accumulation steps (1 if none); the effective batch is B·A sequences.

S

the number of optimizer steps.

N_train

the number of tokens in the training set; epochs can be fractional.

Example: B = 32, n = 64, A = 1, S = 400 gives 2,048 tokens per step and 819,200 tokens seen. On a training stream of 34,869 characters (the notebook) that is 23.5 epochs: the toy model sees each character many times, which is fine for a toy text and would be called overfitting-prone for a real corpus. Large-scale pretraining mostly makes less than one pass over its data.

Accumulation. Averaging the gradients of A micro-batches equals the gradient of the average loss over A · B sequences. In code, compute each micro-loss divided by A, call backward() on each, and step the optimizer once.

Evaluation noise. The validation loss is itself an estimate: with m evaluation batches its standard error shrinks like 1/√m. Use enough batches (10 to 100 for small runs) that differences you care about are larger than the noise, and use the same batches each time (a fixed generator) when comparing checkpoints.

Time and throughput. A step costs about 6 · P · B · n FLOPs for a model with P parameters (Lesson LM4-T5). For the notebook model (about 107,000 parameters) that is 6 · 106,624 · 2,048 ≈ 1.3 · 10⁹ FLOPs per step. A CPU delivering some tens of GFLOP/s in practice therefore needs well under a second per step, in line with the measured 20 ms to 50 ms.

A worked example: a complete loop with a checkpoint round trip

The model below is a bigram language model (an embedding table that maps each character straight to logits) so that the example runs in a second. The loop is the same one you would use for a GPT.

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

torch.manual_seed(0)
text = "the robot picks up the red box. " * 60
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
cut = int(0.9 * len(data))
train_data, val_data = data[:cut], data[cut:]
V, n, B = len(chars), 16, 8

def get_batch(d, gen):
    s = torch.randint(0, len(d) - n - 1, (B,), generator=gen)
    return torch.stack([d[i:i + n] for i in s]), torch.stack([d[i + 1:i + n + 1] for i in s])

class Bigram(nn.Module):
    def __init__(self):
        super().__init__()
        self.table = nn.Embedding(V, V)          # logits for the next character given the current one
    def forward(self, idx):
        return self.table(idx)

def evaluate(model, d, batches=5):
    gen = torch.Generator().manual_seed(123)   # same batches every time
    model.eval()
    with torch.no_grad():
        total = sum(F.cross_entropy(model(x).flatten(0, 1), y.flatten()).item() for x, y in (get_batch(d, gen) for _ in range(batches)))
    model.train()
    return total / batches

model = Bigram()
opt = torch.optim.AdamW(model.parameters(), lr=0.05, weight_decay=0.0)
gen = torch.Generator().manual_seed(1)
print("step 0  val", round(evaluate(model, val_data), 3))
for step in range(1, 201):
    x, y = get_batch(train_data, gen)
    loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
    opt.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()
    if step % 100 == 0:
        print("step", step, "train", round(loss.item(), 3), "val", round(evaluate(model, val_data), 3))

# checkpoint round trip: weights, optimizer state and step
buf = io.BytesIO()
torch.save({"model": model.state_dict(), "opt": opt.state_dict(), "step": 200}, buf)
buf.seek(0)
ckpt = torch.load(buf)
model2 = Bigram(); model2.load_state_dict(ckpt["model"])
opt2 = torch.optim.AdamW(model2.parameters(), lr=0.05, weight_decay=0.0); opt2.load_state_dict(ckpt["opt"])
print("same weights:", all(torch.equal(a, b) for a, b in zip(model.parameters(), model2.parameters())))   # True
print("same val loss:", evaluate(model, val_data) == evaluate(model2, val_data))                          # True
print("optimizer step counter restored:", opt2.state_dict()["state"][0]["step"].item(), "== 200")

The printed validation loss falls from about ln(V) = 3.1 to about 0.7 within 200 steps. A model that sees only the previous character cannot do better on this text, because the same letter is followed by different letters in different places. Look at what evaluate does: eval mode, no_grad, a fixed generator so that the same validation batches are used each time, and a switch back to train mode.

Engineering implications and limitations

Checkpoint often enough. Choose the interval so that losing the work since the last checkpoint is acceptable. Keep the last few, plus the best by validation loss.

Save the optimizer. Resuming with a fresh optimizer restarts Adam’s statistics: the first updates are much too large. Save and load optimizer.state_dict().

Load with map_location and weights_only. When loading a checkpoint on a different device pass map_location. Checkpoints are pickles; load only files you trust, and prefer weights_only=True for plain tensors.

Reproducibility. Seed Python, NumPy and PyTorch, and pass an explicit generator to the data sampler. Exact bit-for-bit reproducibility across machines is not guaranteed, but runs that differ by more than noise point to a real difference in settings.

Do not evaluate too often. Evaluation costs forward passes. Every 50 to 1,000 steps is typical; the interval should give at least a dozen points on the curve.

Use loss.item() sparingly. It forces the GPU to finish all queued work. Log it every few steps, not every step, in large runs.

Early stopping. If validation loss has not improved for a long time while training loss falls, stop and keep the best checkpoint. For pretraining with a cosine schedule one usually trains to the end.

Log to a file. Write the curves to a file or an experiment tracker. Terminal output scrolls away, and you will want to compare runs later.

Common misconceptions

Engineering takeaways

1. Per step: batch, forward and loss, zero_grad, backward, clip, optimizer step, schedule update.
2. Evaluate every N steps in eval mode with no_grad on fixed validation batches, then return to train mode.
3. Tokens per step = B·n·A; tokens seen = S·B·n·A; epochs = tokens seen / training tokens.
4. A checkpoint holds model weights, optimizer state, step (and scheduler and RNG state) so that training can resume.
5. Log training and validation loss, learning rate, gradient norm and tokens; the curves are how you read a run.

References and further reading

See Chapter 5 of the book (training an LLM, and loading and saving model weights in PyTorch) and Appendix D (additions to the training loop).

• PyTorch documentation: saving and loading models, torch.nn.utils.clip_grad_norm_ and torch.optim.AdamW.
• Build a Large Language Model (From Scratch), Chapter 5: Pretraining on Unlabeled Data (primary book reference)
• Build a Large Language Model (From Scratch), Appendix D: Adding Bells and Whistles to the Training Loop (primary book reference)

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: reading the curves: overfitting, underfitting and what the data have to do with it.

Related content

Nothing below is required to finish this lesson.

• Before
  LM5-T2: Backpropagation and Gradient-Based Optimization

  Explain gradients, gradient descent and AdamW, and how learning rate and clipping shape training.

• Lab
  LM5-L3: Pretrain a Tiny GPT and Diagnose the Run

  Train a small GPT on a small corpus, plot training and validation loss, sample at different stages and diagnose a bad run.