promptdojo_

Checkpoints and reproducibility — step 1 of 7

ch 43 · training loops, backprop, optimizers, and schedulers1/7

Checkpoints and reproducibility: training you can trust

Training runs crash, get preempted, and — worse — succeed without you being able to say exactly what produced the weights. Chapter 37 gave pipelines checkpoints and idempotency; training gets the same discipline.

What a checkpoint must contain

Run the editor: a checkpoint is enough state to resume as if nothing happened. For a real run that's more than weights:

torch.save({
    "step": step,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),   # Adam's momentum lives here!
    "scheduler": scheduler.state_dict(),
    "config": config,                      # lr, batch size, data version
}, f"ckpt_{step}.pt")

The one people forget is the optimizer state — Adam's running averages are training state; resuming weights-only with a fresh optimizer changes the trajectory. Save on a cadence, keep the last few plus best-on-validation, and test the resume path once before the 12-hour run needs it — an unrehearsed restore is a hope, not a plan (chapter 17's backup lesson, ML edition).

Reproducibility: honest, not absolute

Same code + same data + same seed should mean same-ish run:

  • Seed everything you control (random, numpy, torch) and log the seed with the run.
  • Pin the data version. "The dataset" is a moving target unless the checkpoint records exactly which build it trained on — chapter 37's manifests are the ML training story too.
  • Log the config — every hyperparameter, in the checkpoint and wherever you track experiments (chapter 45 systematizes this).

Bit-perfect reproducibility across GPUs/library versions is genuinely hard (parallel floating-point sums don't commute exactly); the working standard is statistical reproducibility — re-runs land within normal seed-to-seed wobble. Which connects to chapter 38: two runs differing by a hair may be noise. Vary the seed before declaring victory.

Where AI specifically gets this wrong

  • Weights-only "checkpoints," restored on faith. Resume silently changes optimizer behavior — save the whole state dict bundle. And test the load path before you need it: the save code ran for weeks; the load code runs for the first time during the outage.
  • "Improvements" measured across changed seeds and data. The generated experiment changed three things and credits the model. One change per comparison — the tracker chapter makes this mechanical.