Build checkpoints and resumable training

Problem

When you need this. A stopped training process loses its progress, and the final weights may be worse than an earlier version. Save enough state to resume the run and keep the version with the lowest development loss.

The previous build trained for 600 steps and kept only the final weights. Its development loss was lowest at step 450, so the best weights were lost. And if the process had crashed at step 300, all the work up to then would have been lost too. This build adds the two things every longer training run needs: checkpoints that let a run continue after an interruption, and a saved copy of the best weights.

The contract is run(tok, text, cfg, folder, resume=False). Every 50 steps it scores the development stories, writes the full training state to last.json and, when the development loss is the lowest so far, writes the parameters to best.json. With resume=True it loads last.json, verifies it and continues from the step after it.

Success has three checks, all recorded. The checkpointed run gives the same parameter checksum as the previous build’s run. A run interrupted at step 300 and resumed ends with exactly the same parameters as one that never stopped. And a damaged checkpoint, or one written with different settings, is refused instead of loaded.

Concepts

The previous build defined the training, development and test stories. Four ideas make training resumable:

Checkpoint: A file holding everything the next step depends on. That is more than the parameters: Adam’s two running averages and its step count, the shuffle generator’s state, the stories left in the current epoch, the history so far, and hashes of the corpus, tokenizer and numeric implementation. Leave any of them out and the resumed run takes a different path.
Best checkpoint: The parameters at the evaluation with the lowest development loss. This is model selection on development data. The comparison later scores both the selected and final parameters on the reserved test rows; those scores do not decide which checkpoint is saved.
Atomic write: The checkpoint is written to a temporary file in the same folder and then flushed to disk with os.fsync(), then renamed over the old one with os.replace(). On POSIX systems a successful rename is atomic, so if the process stops during the write, the old checkpoint stays, never half of a new one. A power failure is a stronger case: surviving it also needs the folder itself synced after the rename, which this build does not do.
Digest: A SHA-256 hash of the checkpoint’s contents, stored inside it. On load, the program recomputes the hash; any changed value gives a different hash and the file is refused.

Architecture

Checkpointed training · save, verify, resume

At each evaluation, training writes its complete resumable state to last.json and writes best.json when development loss improves. Saving writes a temporary file, flushes it and atomically replaces the target. Resuming verifies integrity, configuration, corpus, tokenizer and source identity, then restores parameters, optimizer state and shuffle state. Invalid checkpoints are refused.

Start inside checkpoint.run() at Adam step. Every 50 steps, dev loss scores the development stories, and the decision lower than best? has two outcomes. On yes, save() writes best.json with the step, the development loss, the configuration and the parameters. On either outcome, save() then writes the full state to last.json, and the path labeled next step continues training. On restart, the dashed hand-off carries last.json to load(), which checks integrity, configuration and training-input identity. If one check fails, load() raises ValueError with the recorded message. Otherwise run() continues from the restored step. The run comparison below reports development loss and parameter identity.

checkpoint.py holds save(), load() and run(). run() repeats the training loop from the previous build step for step, so with the same seeds it produces the same parameters, and adds the checkpoint writes. Two flags, restore_optimizer and restore_shuffle, exist only to reproduce two resume bugs in the tests. trace_checkpoints.py runs four jobs in a temporary folder and writes checkpoint-trace.json and best.json.

Working code

Runtime: Python 3.10 or newer. Dependencies: Python standard library only.

Download and extract the complete lab using Download the complete lab above. Open build-checkpoints-and-resumable-training/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

checkpoint.pySource · pythonDownload

"""Checkpointed, resumable training for the one-block model, standard library only.

A checkpoint holds everything the next step depends on: the parameters, Adam's
two running averages and step count, the shuffle generator's state, the stories
left in the current epoch and the recorded history. It is written to a temporary
file and renamed into place, and it carries a SHA-256 digest of its contents.
"""
import hashlib
import json
import math
import os
import random
from pathlib import Path

import model as m
import train as t

FORMAT_VERSION = 1


def _digest(body):
    return hashlib.sha256(json.dumps(body, sort_keys=True).encode()).hexdigest()


def params_sha256(params):
    """An exact fingerprint of every parameter value: equal hashes mean equal parameters."""
    return hashlib.sha256(json.dumps(params, sort_keys=True).encode()).hexdigest()


def save(path, body):
    """Write to a temporary file, flush it to disk, then rename it over the old checkpoint.

    If the process stops at any point, the old checkpoint or the complete new one is
    left, never half of one. Surviving a power failure would also need the folder
    itself to be flushed, which this function does not do.
    """
    path = Path(path)
    data = dict(body, format_version=FORMAT_VERSION)
    data["sha256"] = _digest(data)
    tmp = path.with_suffix(path.suffix + ".tmp")
    with open(tmp, "w") as f:
        f.write(json.dumps(data))
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def input_signature(tok, text):
    """Bind the data, token IDs and exact numeric implementation used by a run."""
    source = Path(__file__).parent
    return {
        "corpus_sha256": hashlib.sha256(text.encode()).hexdigest(),
        "tokenizer_sha256": _digest(tok.to_json()),
        "source_sha256": {name: hashlib.sha256((source / name).read_bytes()).hexdigest()
                          for name in ("model.py", "train.py", "bpe.py", "checkpoint.py")},
    }


def load(path, cfg, expected_inputs=None):
    """Read and verify a checkpoint. Raises ValueError if it is damaged or from another run."""
    try:
        data = json.loads(Path(path).read_text())
    except json.JSONDecodeError as error:
        raise ValueError(f"checkpoint {path} is not valid JSON: {error.msg}") from None
    if not isinstance(data, dict):
        raise ValueError(f"checkpoint {path} does not hold a JSON object")
    if data.get("format_version") != FORMAT_VERSION:
        raise ValueError(f"checkpoint format {data.get('format_version')!r}; expected {FORMAT_VERSION}")
    digest = data.pop("sha256", None)
    if digest != _digest(data):
        raise ValueError(f"checkpoint {path} failed its SHA-256 check")
    if data["config"] != cfg:
        changed = sorted(k for k in set(cfg) | set(data["config"]) if cfg.get(k) != data["config"].get(k))
        raise ValueError(f"checkpoint was written with a different config: {changed}")
    if expected_inputs is not None and data.get("training_inputs") != expected_inputs:
        raise ValueError("checkpoint training inputs differ: corpus, tokenizer or source changed")
    return data


def run(tok, text, cfg, folder, resume=False, stop_after=None, restore_optimizer=True, restore_shuffle=True):
    """Train with a checkpoint every eval_every steps.

    Writes folder/last.json (full state) and folder/best.json (parameters with the
    lowest development loss so far). stop_after simulates an interruption.
    restore_optimizer=False and restore_shuffle=False reproduce two resume bugs.
    """
    folder = Path(folder)
    folder.mkdir(parents=True, exist_ok=True)
    training_inputs = input_signature(tok, text)
    train_s, dev_s, _ = t.split_stories(text)
    tr = t.encode_all(tok, train_s, cfg["max_len"])
    dv = t.encode_all(tok, dev_s, cfg["max_len"])
    p = m.init_params(len(tok.vocab), cfg["d"], cfg["hidden"], cfg["max_len"], cfg["init_seed"])
    opt = t.Adam(p, cfg["learning_rate"], cfg["beta1"], cfg["beta2"], cfg["adam_eps"])
    rng = random.Random(cfg["data_seed"])
    order, history, dev_history, best, start = [], [], [], None, 0
    if resume:
        state = load(folder / "last.json", cfg, training_inputs)
        p.update(state["params"])
        if restore_optimizer:
            opt.m, opt.v, opt.t = state["adam"]["m"], state["adam"]["v"], state["adam"]["t"]
        if restore_shuffle:
            version, internal, gauss = state["rng_state"]
            rng.setstate((version, tuple(internal), gauss))
            order = state["order"]
        history, dev_history, best, start = state["history"], state["dev_history"], state["best"], state["step"]
    for step in range(start + 1, cfg["steps"] + 1):
        if len(order) < cfg["batch"]:
            order += rng.sample(range(len(tr)), len(tr))
        batch = [tr[i] for i in order[:cfg["batch"]]]
        del order[:cfg["batch"]]
        loss, g = t.batch_gradient(p, batch)
        opt.step(p, g)
        history.append(round(loss, 4))
        if step % cfg["eval_every"] == 0 or step == cfg["steps"]:
            dev_loss = round(t.mean_loss(p, dv), 4)
            dev_history.append({"step": step, "dev_loss": dev_loss})
            if best is None or dev_loss < best["dev_loss"]:
                best = {"step": step, "dev_loss": dev_loss}
                save(folder / "best.json", {"step": step, "dev_loss": dev_loss, "config": cfg, "params": p, "training_inputs": training_inputs})
            save(folder / "last.json", {
                "step": step, "config": cfg, "params": p, "training_inputs": training_inputs,
                "adam": {"m": opt.m, "v": opt.v, "t": opt.t},
                "rng_state": list(rng.getstate()), "order": order,
                "history": history, "dev_history": dev_history, "best": best,
            })
            if stop_after is not None and step >= stop_after:
                break
    return p, history, dev_history, best

save() computes the digest over the contents with sorted keys, so the same contents always give the same digest. load() checks the JSON, format version, digest and configuration. Resumption also compares training_inputs: the corpus bytes, tokenizer alphabet and merge order, and source hashes. A changed tokenizer may keep the same vocabulary size while changing every token ID; matching tensor shapes cannot detect that mistake. The round-trip test checks that loaded parameters are exactly the saved values, and the resume tests depend on that.

Read trace_checkpoints.py

trace_checkpoints.pySource · pythonDownload

"""Record uninterrupted, interrupted-and-resumed and incorrectly resumed training runs.

Run: python3 trace_checkpoints.py          (prints the trace; checkpoint-time.json records how long one run took)
     python3 trace_checkpoints.py --write  (also rewrite checkpoint-trace.json, best.json and checkpoint-time.json)
     python3 trace_checkpoints.py --check  (exit 1 if either file differs)
Checkpoints for each run go to a temporary folder that is deleted afterwards.
"""
import json
import platform
import time
import shutil
import sys
import tempfile
from pathlib import Path

import checkpoint as ck
import train as t
from bpe import Tokenizer

HERE = Path(__file__).parent
STOP_AT = 300


def scores(p, tok, test_seqs, test_stories):
    return {"test_loss": round(t.mean_loss(p, test_seqs), 4), "color_recall": t.color_recall(p, tok, test_stories),
            "params_sha256": ck.params_sha256(p), "checksum": t.checksum(p)}


def integrity_cases(folder, cfg):
    last = folder / "last.json"
    text = last.read_text()
    cases = {}
    bad = folder / "tampered.json"
    bad.write_text(text.replace('"step": ', '"step": 1', 1))
    truncated = folder / "truncated.json"
    truncated.write_text(text[: len(text) // 2])
    for name, path, config in (("one changed value", bad, cfg), ("truncated file", truncated, cfg),
                               ("different learning rate", last, dict(cfg, learning_rate=0.02))):
        try:
            ck.load(path, config)
            cases[name] = "accepted"
        except ValueError as error:
            cases[name] = str(error).replace(str(folder) + "/", "")
    return cases


def run():
    tok = Tokenizer.load(HERE / "tokenizer.json")
    text = (HERE / "corpus.txt").read_text()
    cfg = t.CONFIG
    _, _, test = t.split_stories(text)
    te = t.encode_all(tok, test, cfg["max_len"])
    work = Path(tempfile.mkdtemp())
    start = time.perf_counter()
    try:
        full, _, dev_history, best = ck.run(tok, text, cfg, work / "full")
        best_params = json.loads((work / "full" / "best.json").read_text())
        resumed = {}
        for name, flags in (("resumed", {}), ("resumed_without_optimizer_state", {"restore_optimizer": False}),
                            ("resumed_without_shuffle_state", {"restore_shuffle": False})):
            folder = work / name
            ck.run(tok, text, cfg, folder, stop_after=STOP_AT)
            p, _, dh, _ = ck.run(tok, text, cfg, folder, resume=True, **flags)
            resumed[name] = {"final_dev_loss": dh[-1]["dev_loss"], "params_sha256": ck.params_sha256(p)}
        record = {
            "config": cfg,
            "interrupted_at_step": STOP_AT,
            "uninterrupted": {"dev_loss": dev_history, "final": scores(full, tok, te, test)},
            "best_checkpoint": dict(best, **scores(best_params["params"], tok, te, test)),
            "resume_runs": resumed,
            "integrity": integrity_cases(work / "full", cfg),
        }
        timing = {"seconds_for_all_runs": round(time.perf_counter() - start, 1), "python": platform.python_version(),
                  "machine": platform.machine(), "system": platform.system()}
        return record, best_params, timing
    finally:
        shutil.rmtree(work)


if __name__ == "__main__":
    record, best, timing = run()
    text = json.dumps(record, indent=1) + "\n"
    best_text = json.dumps(best) + "\n"
    if "--write" in sys.argv:
        (HERE / "checkpoint-trace.json").write_text(text)
        (HERE / "best.json").write_text(best_text)
        (HERE / "checkpoint-time.json").write_text(json.dumps(timing, indent=1) + "\n")
    elif "--check" in sys.argv:
        same = (HERE / "checkpoint-trace.json").read_text() == text and (HERE / "best.json").read_text() == best_text
        sys.exit(0 if same else "checkpoint-trace.json or best.json is out of date")
    else:
        print(text, end="")

Run the four recorded jobs:

Terminal commandsTerminal · bash

python3 trace_checkpoints.py

python3 trace_checkpoints.py --check repeats them and exits with an error if checkpoint-trace.json or best.json differ. The recorded jobs took 333.3 seconds together, measured on one machine; your numbers will differ. That is one measurement, not a benchmark, and it is kept in checkpoint-time.json, which --check does not compare.

Compare the resumed runs

The previous build recorded only train.checksum(), the sum of the squares of every parameter value rounded to 9 decimals. The uninterrupted checkpointed run gives 501.972163829, the same as the previous build’s 501.972163829, so, as far as the checksum can show, writing checkpoints did not change training. A rounded sum can miss a small change, so the comparison below uses params_sha256(), a SHA-256 hash of every parameter value written out in full; the table shows its first 12 hexadecimal digits.

Run	Final development loss	Parameter SHA-256	Same parameters as the uninterrupted run
Uninterrupted	0.3078	`432cc8995907`	yes
Stopped at step 300, resumed	0.3078	`432cc8995907`	yes
Resumed without Adam’s state	0.2980	`b04c82ed5818`	no
Resumed without the shuffle state	0.2999	`e4f11e32d37f`	no

The correctly resumed run is identical: the full hash matches, so every parameter value is the same. The two incomplete resumes end with different parameters. Without Adam’s state, the running averages restart at zero, and the bias correction treats step 301 as if it were step 1. Without the shuffle state, the remaining stories of the epoch and every later shuffle come out in a different order.

Keep the best checkpoint

The development loss was lowest at step 450, 0.2941, so best.json holds those parameters. Scored once on the test stories:

Parameters	Test loss	Color recalled in the last sentence
Best checkpoint (step 450)	0.2628	36 of 60
Final (step 600)	0.2694	39 of 60

The selected checkpoint has lower test loss in this run. The original split still includes duplicate story text, as explained in the baseline lesson. It recalls the color in fewer stories. The lower average test loss in this run does not mean that every targeted check improved. The next build, Build text generation with sampling controls, and the builds after it use best.json and report the color check beside the loss.

Check your understanding

Look at the row for the run resumed without Adam’s state. Before you open the comparison, decide whether that row shows a bug, and how you would know.

Compare your reasoning

Its final development loss, 0.2980, is lower than the uninterrupted run’s 0.3078, so the loss alone would not raise an alarm. The parameter hash shows that it is a different run. A resume should reproduce the uninterrupted run exactly; if it does not, a piece of state is missing, whether or not the loss happens to look better this time.

Read the explanation

A resumable checkpoint holds everything the next step depends on. Compare the final parameters of a resumed and an uninterrupted run to prove it.

Inspect the recorded trace

checkpoint-trace.jsonRecorded output · jsonDownload

{
 "config": {
  "d": 16,
  "hidden": 64,
  "max_len": 64,
  "init_seed": 3,
  "data_seed": 5,
  "batch": 8,
  "steps": 600,
  "learning_rate": 0.01,
  "beta1": 0.9,
  "beta2": 0.999,
  "adam_eps": 1e-08,
  "eval_every": 50
 },
 "interrupted_at_step": 300,
 "uninterrupted": {
  "dev_loss": [
   {
    "step": 50,
    "dev_loss": 0.7212
   },
   {
    "step": 100,
    "dev_loss": 0.5114
   },
   {
    "step": 150,
    "dev_loss": 0.4558
   },
   {
    "step": 200,
    "dev_loss": 0.4143
   },
   {
    "step": 250,
    "dev_loss": 0.3454
   },
   {
    "step": 300,
    "dev_loss": 0.3116
   },
   {
    "step": 350,
    "dev_loss": 0.3082
   },
   {
    "step": 400,
    "dev_loss": 0.304
   },
   {
    "step": 450,
    "dev_loss": 0.2941
   },
   {
    "step": 500,
    "dev_loss": 0.297
   },
   {
    "step": 550,
    "dev_loss": 0.3141
   },
   {
    "step": 600,
    "dev_loss": 0.3078
   }
  ],
  "final": {
   "test_loss": 0.2694,
   "color_recall": {
    "correct": 39,
    "total": 60
   },
   "params_sha256": "432cc8995907c11af87dfecd4bf5914d206c9c02ed5ba219a75e10db39fdf756",
   "checksum": 501.972163829
  }
 },
 "best_checkpoint": {
  "step": 450,
  "dev_loss": 0.2941,
  "test_loss": 0.2628,
  "color_recall": {
   "correct": 36,
   "total": 60
  },
  "params_sha256": "8342f2b53060f3a56a340b87eaaf53348afa9f9c1b8697e4995ec06564d24677",
  "checksum": 455.712777262
 },
 "resume_runs": {
  "resumed": {
   "final_dev_loss": 0.3078,
   "params_sha256": "432cc8995907c11af87dfecd4bf5914d206c9c02ed5ba219a75e10db39fdf756"
  },
  "resumed_without_optimizer_state": {
   "final_dev_loss": 0.298,
   "params_sha256": "b04c82ed581898a15f00ff365b7e234f166b587630c0567bc5e16763fd178bd3"
  },
  "resumed_without_shuffle_state": {
   "final_dev_loss": 0.2999,
   "params_sha256": "e4f11e32d37f98bfa6286484ec27f2cd52987c588c7e600c643c1dd0292e6e4f"
  }
 },
 "integrity": {
  "one changed value": "checkpoint tampered.json failed its SHA-256 check",
  "truncated file": "checkpoint truncated.json is not valid JSON: Expecting ',' delimiter",
  "different learning rate": "checkpoint was written with a different config: ['learning_rate']"
 }
}

Failure cases

Changing a tokenizer can preserve its vocabulary size while assigning different IDs to the same text. Matching matrix shapes would then hide a different training task. A corpus edit can also preserve the number of stories while changing their content. The input identity check refuses both changes before the next optimizer step.

Run python inspect_inputs.py --check to exercise a matching bundle and four refusal cases without retraining. test_input_identity_refusals_match_the_record also checks a missing identity and a changed implementation. SHA-256 here detects a mismatch; it does not authenticate a file supplied by an untrusted party.

Each case has a named test in test_checkpoint.py.

Case	Result with this implementation	What it tells you
A checkpoint with one changed value	`checkpoint tampered.json failed its SHA-256 check`	Any changed value is caught by the digest
A truncated checkpoint	`checkpoint truncated.json is not valid JSON: Expecting ',' delimiter`	A half-written file is refused, not partly loaded
A checkpoint from a different configuration	`checkpoint was written with a different config: ['learning_rate']`	Continuing with other settings would mix two runs
A file that holds JSON null	`load()` raises `ValueError`	A file that is valid JSON but not a checkpoint is refused, not read
A write stopped before the rename	`save()` raises; the old checkpoint still loads	The rename is the only step that replaces the file
A checkpoint in another format version	`load()` raises `ValueError`	A file written by a changed program is refused
Resuming without the optimizer state	The final parameters differ from the uninterrupted run	Adam’s averages and step count are part of the state
Resuming without the shuffle state	The final parameters differ from the uninterrupted run	The data order is part of the state
Changed corpus	Refused before the next optimizer step	The corpus digest binds content, even when the story count stays the same
Changed tokenizer with the same vocabulary size	Refused before the next optimizer step	Equal matrix dimensions do not establish equal token IDs
Missing or changed input identity	Refused before the next optimizer step	A resume needs the corpus, tokenizer and implementation bindings recorded by this format

The complete lab includes these tests and their fixtures. Run the 20 standard-library tests:

Terminal commandsTerminal · bash

python3 -m unittest -v test_checkpoint.py

Read the tests

test_checkpoint.pySource · pythonDownload

"""Tests for checkpoint.py. Run: python3 -m unittest -v test_checkpoint.py

Short training jobs check resuming; the recorded runs are repeated
only by `python3 trace_checkpoints.py --check`.
"""
import copy
import json
import inspect_inputs
import tempfile
import unittest
from pathlib import Path
from unittest import mock

import checkpoint as ck
import train as t
from bpe import Tokenizer

HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")
TEXT = (HERE / "corpus.txt").read_text()
RECORD = json.loads((HERE / "checkpoint-trace.json").read_text())
SHORT = dict(t.CONFIG, steps=6, eval_every=2)


class Files(unittest.TestCase):
    def setUp(self):
        self.dir = Path(tempfile.mkdtemp())
        self.path = self.dir / "c.json"
        ck.save(self.path, {"step": 4, "config": SHORT, "params": {"w": [0.5, -1.25]}})

    def test_round_trip_keeps_exact_values(self):
        self.assertEqual(ck.load(self.path, SHORT)["params"], {"w": [0.5, -1.25]})

    def test_atomic_write_leaves_no_temporary_file(self):
        self.assertEqual(sorted(p.name for p in self.dir.iterdir()), ["c.json"])

    def test_failed_rename_keeps_the_old_checkpoint(self):
        # Stop save() at the last moment: the new file is written but never renamed into place.
        with mock.patch.object(ck.os, "replace", side_effect=OSError("stopped before the rename")):
            with self.assertRaises(OSError):
                ck.save(self.path, {"step": 6, "config": SHORT, "params": {"w": [9.0, 9.0]}})
        old = ck.load(self.path, SHORT)
        self.assertEqual((old["step"], old["params"]), (4, {"w": [0.5, -1.25]}))

    def test_changed_value_is_rejected(self):
        self.path.write_text(self.path.read_text().replace("-1.25", "-1.5"))
        with self.assertRaisesRegex(ValueError, "SHA-256"):
            ck.load(self.path, SHORT)

    def test_truncated_file_is_rejected(self):
        text = self.path.read_text()
        self.path.write_text(text[: len(text) // 2])
        with self.assertRaisesRegex(ValueError, "not valid JSON"):
            ck.load(self.path, SHORT)

    def test_different_config_is_rejected(self):
        with self.assertRaisesRegex(ValueError, r"different config: \['learning_rate'\]"):
            ck.load(self.path, dict(SHORT, learning_rate=0.02))

    def test_json_that_is_not_an_object_is_rejected(self):
        self.path.write_text("null")
        with self.assertRaisesRegex(ValueError, "does not hold a JSON object"):
            ck.load(self.path, SHORT)

    def test_other_format_version_is_rejected(self):
        data = json.loads(self.path.read_text())
        data["format_version"] = 2
        self.path.write_text(json.dumps(data))
        with self.assertRaisesRegex(ValueError, "format 2"):
            ck.load(self.path, SHORT)


class Resume(unittest.TestCase):
    def runs(self, **flags):
        work = Path(tempfile.mkdtemp())
        full = ck.run(TOK, TEXT, SHORT, work / "full")[0]
        ck.run(TOK, TEXT, SHORT, work / "cut", stop_after=2)
        resumed = ck.run(TOK, TEXT, SHORT, work / "cut", resume=True, **flags)[0]
        return json.dumps(full, sort_keys=True), json.dumps(resumed, sort_keys=True)  # every value

    def test_resume_refuses_changed_corpus(self):
        cfg=dict(SHORT,d=2,hidden=4,batch=1,steps=2,eval_every=1)
        with tempfile.TemporaryDirectory() as work:
            ck.run(TOK,TEXT,cfg,work,stop_after=1)
            with self.assertRaisesRegex(ValueError,"training inputs differ"):
                ck.run(TOK,TEXT.replace("red","blue"),cfg,work,resume=True)

    def test_resume_refuses_changed_tokenizer_with_same_vocab_size(self):
        cfg=dict(SHORT,d=2,hidden=4,batch=1,steps=2,eval_every=1)
        changed=Tokenizer(list(reversed(TOK.alphabet)),TOK.merges)
        self.assertEqual(len(changed.vocab),len(TOK.vocab))
        with tempfile.TemporaryDirectory() as work:
            ck.run(TOK,TEXT,cfg,work,stop_after=1)
            with self.assertRaisesRegex(ValueError,"training inputs differ"):
                ck.run(changed,TEXT,cfg,work,resume=True)

    def test_resume_matches_an_uninterrupted_run(self):
        full, resumed = self.runs()
        self.assertEqual(full, resumed)

    def test_resume_across_an_epoch_boundary_restores_the_shuffle(self):
        # With 56 stories per batch an epoch lasts 5 steps, so step 6 needs a fresh shuffle.
        cfg = dict(t.CONFIG, steps=6, eval_every=3, batch=56)
        work = Path(tempfile.mkdtemp())
        full = ck.run(TOK, TEXT, cfg, work / "full")[0]
        ck.run(TOK, TEXT, cfg, work / "cut", stop_after=3)
        resumed = ck.run(TOK, TEXT, cfg, work / "cut", resume=True)[0]
        self.assertEqual(json.dumps(full, sort_keys=True), json.dumps(resumed, sort_keys=True))

    def test_best_checkpoint_keeps_the_lowest_development_loss(self):
        # A deliberately high learning rate makes the development loss go up and down.
        cfg = dict(t.CONFIG, steps=6, eval_every=1, learning_rate=1.0)
        work = Path(tempfile.mkdtemp())
        _, _, dev, best = ck.run(TOK, TEXT, cfg, work)
        lowest = min(dev, key=lambda r: r["dev_loss"])
        self.assertNotEqual(lowest["step"], dev[-1]["step"])
        saved = ck.load(work / "best.json", cfg)
        self.assertEqual((saved["step"], saved["dev_loss"]), (lowest["step"], lowest["dev_loss"]))
        self.assertEqual(best, {"step": lowest["step"], "dev_loss": lowest["dev_loss"]})

    def test_resume_without_optimizer_state_differs(self):
        full, resumed = self.runs(restore_optimizer=False)
        self.assertNotEqual(full, resumed)

    def test_resume_without_shuffle_state_differs(self):
        full, resumed = self.runs(restore_shuffle=False)
        self.assertNotEqual(full, resumed)


class Recorded(unittest.TestCase):
    def test_input_identity_refusals_match_the_record(self):
        rows=inspect_inputs.inspect()
        self.assertEqual(rows,json.loads((HERE/'input-checks.json').read_text()))
        self.assertTrue(rows[0]['accepted'])
        self.assertTrue(all(not row['accepted'] for row in rows[1:]))

    def test_recorded_resume_matches_and_bugs_do_not(self):
        runs, full = RECORD["resume_runs"], RECORD["uninterrupted"]["final"]["params_sha256"]
        self.assertEqual(runs["resumed"]["params_sha256"], full)
        self.assertNotEqual(runs["resumed_without_optimizer_state"]["params_sha256"], full)
        self.assertNotEqual(runs["resumed_without_shuffle_state"]["params_sha256"], full)

    def test_best_checkpoint_is_the_lowest_development_loss(self):
        best = RECORD["best_checkpoint"]
        lowest = min(RECORD["uninterrupted"]["dev_loss"], key=lambda r: r["dev_loss"])
        self.assertEqual((best["step"], best["dev_loss"]), (lowest["step"], lowest["dev_loss"]))

    def test_params_hash_detects_one_changed_value(self):
        p = {"w": [0.5, -1.25]}
        q = {"w": [0.5, -1.2500000000000002]}
        self.assertNotEqual(ck.params_sha256(p), ck.params_sha256(q))

    def test_saved_best_checkpoint_reproduces_its_scores(self):
        data = ck.load(HERE / "best.json", t.CONFIG)
        _, _, test = t.split_stories(TEXT)
        te = t.encode_all(TOK, test, t.CONFIG["max_len"])
        self.assertEqual(ck.params_sha256(data["params"]), RECORD["best_checkpoint"]["params_sha256"])
        self.assertEqual(round(t.mean_loss(data["params"], te), 4), RECORD["best_checkpoint"]["test_loss"])
        self.assertEqual(t.color_recall(data["params"], TOK, test), RECORD["best_checkpoint"]["color_recall"])


if __name__ == "__main__":
    unittest.main()

test_resume_across_an_epoch_boundary_restores_the_shuffle fails if run() restores the remaining stories but not the generator. test_best_checkpoint_keeps_the_lowest_development_loss uses a deliberately high learning rate so that the development loss rises and falls, and fails if best.json is written at every evaluation.

Deployment

best.json is the artifact the next builds load; last.json is only for resuming and can be deleted when a run finishes. Keep the configuration with the checkpoint, as the files here do, because load() refuses a file written with other settings.

Scaling

This section is unmeasured. A full checkpoint holds the parameters and Adam’s two averages, about three times the parameters’ size, and this build writes it as JSON text at every evaluation. That is small here. For larger models, how often to write, how many old checkpoints to keep and which file format to use become real trade-offs between lost work, disk space and time spent writing.

Key takeaways

A resumable checkpoint holds everything the next step depends on: parameters, optimizer state, data order and random state.
Write checkpoints atomically and verify them on load; refuse a file that is damaged or was written with other settings.
Keep the best checkpoint by development loss, score it once on the test stories, and still report the targeted checks.

Best practices

Test resumption by comparing final parameters with an uninterrupted run, both within an epoch and across an epoch boundary.
Store the configuration and a format version inside every checkpoint.
Do not judge a resumed run by its loss; a different run can look as good or better.

References

Sources checked October 20, 2026.

Python 3.15: os.replace: a successful rename is atomic, a POSIX requirement; the rename can fail across file systems, which is why the temporary file sits beside the target.
Python 3.15: os.fsync: forces a file’s data to disk; call flush() first for a Python file object.
Python 3.15: random.getstate and random.setstate: the saved state restores the generator exactly.
Python 3.15: hashlib: sha256() and hexdigest().
PyTorch tutorials: saving and loading models, section on a general checkpoint for resuming training: save the optimizer’s state with the model’s, along with items such as the epoch and loss.