PyTorch, section 6 of 6, 3h

Checkpoints and reproducibility

By the end: You can save and resume a training run exactly where it stopped, and make runs as repeatable as the hardware allows.

  • What a checkpoint must contain
  • Loading on a different device
  • Resuming mid-run
  • Seeds and deterministic settings

New to pytorch? PyTorch is the Python library almost everyone uses to build and train neural networks. It hands you two things you would otherwise have to write: arrays of numbers that can live on a GPU instead of the main processor, and automatic differentiation — it watches the arithmetic you perform and works out on its own how every weight should move to reduce the error. Layers, optimisers and data loading come with it. Why it matters

A long run will be interrupted, by a preempted instance or a full disk or a crash at hour nine, and what you saved decides whether you resume or start again. This section is what a checkpoint has to contain to resume exactly: not only the weights, but the optimiser state, the scheduler, the step count and the random number generator states. It also covers how far repeatability actually goes, which is less far than setting seeds suggests, and why comparing several seeds beats trusting one run.

More than the weights#

import torch

torch.save({
    "model": model.state_dict(),
    "optimizer": opt.state_dict(),       # Adam's moments: resuming without them restarts learning dynamics
    "scheduler": sched.state_dict(),
    "step": step,
    "rng": torch.get_rng_state(),
    "cuda_rng": torch.cuda.get_rng_state_all(),
}, f"ckpt_{step}.pt")
What                 Why it must be there                          Left out, you get
model state_dict     the weights, and buffers such as BN stats     a model, missing its statistics
optimizer            Adam's two moment estimates, per parameter    a different loss curve, see below
scheduler            where in the warmup or decay you were         the learning rate resets to step 0
step / epoch          where to resume counting                     logs that overlap and a wrong schedule
RNG states           shuffling, dropout, augmentation              a different data order from here on
data position        which examples have been seen this epoch      some examples twice, some never
config               the hyperparameters this run used             a checkpoint nobody can reproduce
The first two are the ones people remember. Rows 3 to 6 are why a resumed run diverges from an uninterrupted one even when the weights are right.

The optimiser row is the one worth measuring, because the cost of omitting it is invisible. AdamW holds two running averages per parameter, exp_avg and exp_avg_sq, plus a step count used for bias correction. Discard them and the first few steps after a resume take badly scaled updates while those averages refill.

# Train 40 steps straight through, and compare against stopping at 20 and resuming.
#
# step | uninterrupted | resumed with opt | resumed without opt
#   20 |    1.039943   |     1.039943     |    1.039943
#   21 |    1.006236   |     1.006236     |    0.999241      <- already different
#   25 |    0.871761   |     0.871761     |    0.848588
#   39 |    0.432757   |     0.432757     |    0.399392
#
# max deviation with    optimiser state: 0.0          bit-identical
# max deviation without optimiser state: 0.0338

print(list(opt.state_dict()["state"][0].keys()))   # ['step', 'exp_avg', 'exp_avg_sq']
With the optimiser state the resumed run is identical to the nanosecond. Without it, it diverges on the very first step.

Note which way it diverged. The run without optimiser state reached a lower loss, because emptying Adam's second moment produces larger effective steps for a while. That is the trap: the broken resume looks better, so nobody investigates, and the experiment you thought you were continuing is now a different experiment whose curve cannot be compared with anything.

Loading#

ckpt = torch.load("ckpt_4000.pt", map_location="cpu")   # load anywhere, move later
model.load_state_dict(ckpt["model"])
opt.load_state_dict(ckpt["optimizer"])
sched.load_state_dict(ckpt["scheduler"])
start_step = ckpt["step"] + 1

weights_only

Recent PyTorch versions load checkpoints with weights_only=True by default, refusing arbitrary pickled Python objects. It is a security feature, since loading a pickle can run code. Keep checkpoints to tensors and plain values, and never load untrusted checkpoints with it turned off.

torch.save({"w": torch.ones(2), "obj": SomeCustomClass()}, "x.pt")
torch.load("x.pt")
# UnpicklingError: Weights only load failed. This file can still be loaded, to do so
#                  you have two options ...

torch.load("x.pt", weights_only=False)     # runs whatever the pickle says to run
torch.save({"w": torch.ones(2), "step": 7}, "y.pt")
torch.load("y.pt")                         # fine: tensors and plain values only
The fix is to store plain values, not to pass weights_only=False. Saving a whole config object instead of a dict is what usually causes this.

Two further habits. Save to a temporary name and rename afterwards, because a process killed midway through writing leaves a truncated file that loads as a corrupt checkpoint rather than as no checkpoint. And keep two: the latest, for resuming, and the best validation so far, for shipping. They are rarely the same file, and the run that overwrites one with the other has lost the model it was trying to produce.

import os

def save_atomic(state, path):
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)          # atomic on the same filesystem

save_atomic(state, "last.pt")
if val_loss < best_val:
    best_val = val_loss
    save_atomic(state, "best.pt")
os.replace is atomic, so a reader sees either the old file or the new one, never half of either.

Repeatable runs#

import random, numpy as np

def seed_everything(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)

seed_everything(0)
torch.use_deterministic_algorithms(True)   # errors on ops without a deterministic version

Even with every seed set, some GPU operations are not bit-for-bit repeatable, and different hardware or library versions differ. Aim for runs that agree closely, record versions, and compare results across several seeds rather than trusting one.

What                              Repeatable?   Why not
same machine, same versions       usually yes   with seeds set and deterministic algorithms on
DataLoader with num_workers > 0   needs care    each worker needs its own seed; use worker_init_fn
atomicAdd on GPU (scatter, some
  pooling and index operations)   no            float addition is not associative and the
                                                thread order is not fixed between runs
different GPU model               no            different kernels chosen for the same operation
different library version         no            kernel selection and defaults change
across a cuDNN benchmark cycle    no            it times algorithms and picks the fastest,
                                                which varies with machine load
use_deterministic_algorithms(True) turns the second row into an error rather than a silent difference, which is the point of it.

The practical conclusion is not to chase bit-exactness but to stop treating a single run as a measurement. A difference of 0.3 per cent between two configurations is noise unless you have run each three times and the gap is larger than the spread within each. Record the seed alongside the result, report the mean and the range, and reserve determinism for debugging, where you genuinely need two runs to differ in only the thing you changed.

The versions are part of the checkpoint

Write torch.__version__, the CUDA version, the Python version and your own commit hash into the checkpoint dict. It costs a few bytes and is the difference between reproducing a six-month-old result and guessing at it. Loading a checkpoint saved by a much older PyTorch usually works; loading one into a much older PyTorch usually does not.

Where people go wrong

Saving only the model weights every epoch, then discovering that resuming a crashed run changes the loss curve because the optimiser started from scratch.

If this comes up in an interview

What do you need to save to resume a training run exactly?

Model weights, optimiser state, scheduler state, the step or epoch counter, random number generator states, and the position in the data. Missing the optimiser state is the common mistake: Adam restarts with empty moment estimates and training behaves differently.

Build this: Train your small transformer for 2,000 steps, kill it at 1,000, resume from the checkpoint, and confirm the loss curve matches an uninterrupted run closely.

Watch it explained

7 min

How to save and load models in Pytorch

Aladdin Persson

This is one section of PyTorch, 30 hours across 6 sections. Read it in full to get the whole topic in order.

Previous sectionWhen the loss will not go down