The checkpoint is a file

A checkpoint is a message to a process that does not exist yet. That process will build a model, an optimizer and a scheduler from scratch, overwrite their state with what it reads, and carry on as though nothing happened. Everything it needs has to be in the message, and nothing in the message is checked against the run that wrote it.

The contents of that message were settled in chapter four's arc. Five things go in, four of them serialize, and the third lesson there measured what each omission costs over the six steps after a restore. None of that is repeated here, and the itemized table is worth having open beside this page.

What is left is the artifact. A file has a size, a write has a duration, a process has a set of random streams that a single saved block does not cover, and none of those facts appear in a dictionary of tensors. They appear later, as a checkpoint that is a hundred times too large, as a resume that diverges on step one, or as a directory holding a file that loads and is wrong.

A step reads randomness from more places than a loop makes obvious. The batch index comes from the global generator. A DataLoader built with shuffle=True and a generator of your own reads that one instead. An augmentation written with the standard library reads Python's random. On an accelerator, an op that draws numbers reads that device's generator, not the CPU's.

torch.get_rng_state() returns one of them. Save it, restore it, and the two draws that did not come from the default CPU generator carry straight on from where they were: the private generator gives -0.293429 where it first gave 1.540996, and random.random() gives 0.757954 where it first gave 0.844422. Neither raises, and neither is visible in a loss until the numbers have already gone somewhere else.

The count that matters for your own loop is the number of streams, not the number of lines. Chapter five's arc derives worker seeding from a base seed the main process draws out of whichever generator the loader was handed, so a stream you restore is one the workers inherit. Everything else on the list is yours to save by hand or to design out of the loop.

A checkpoint with one RNG entry in it is a claim that the loop has one random stream.

import random
import torch

torch.manual_seed(0)
loader_gen = torch.Generator().manual_seed(0)
random.seed(0)
saved = torch.get_rng_state()          # the whole "rng" entry of a typical checkpoint

first = (torch.randn(1).item(), torch.randn(1, generator=loader_gen).item(), random.random())
torch.set_rng_state(saved)             # the whole restore
again = (torch.randn(1).item(), torch.randn(1, generator=loader_gen).item(), random.random())

print([round(v, 6) for v in first])
print([round(v, 6) for v in again])
print([a == b for a, b in zip(first, again)])
print(torch.cuda.is_available(), torch.cuda.get_rng_state_all())

# [1.540996, 1.540996, 0.844422]
# [1.540996, -0.293429, 0.757954]
# [True, False, False]
# False []

The docstring says it outright, in four lines that have survived every release since. set_rng_state works for CPU. For CUDA the note points you at manual_seed, which is a reseed rather than a restore, and a reseed puts the stream somewhere defined rather than somewhere continuous.

The body of manual_seed explains why the note has to exist. One call fans out to every device family torch knows about, cuda first, then mps, then xpu, then whatever a custom backend registered, before it finally seeds the default CPU generator on the last line. Seeding is device-wide; state capture is not.

So a bitwise resume on an accelerator needs the device's own state in the dictionary, torch.cuda.get_rng_state_all() for a multi-device host, and its matching setter on the way back. This machine has no such device, torch.cuda.get_rng_state_all() returns an empty list above, and that empty list is exactly what a CPU-authored checkpoint quietly carries onto a GPU run.

def set_rng_state(new_state: torch.Tensor) -> None:
    r"""Sets the random number generator state.

    .. note: This function only works for CPU. For CUDA, please use
             torch.manual_seed(seed), which works for both CPU and CUDA.
    ...
    """
    default_generator.set_state(new_state)

...

    seed = int(seed)
    import torch.cuda

    if not torch.cuda._is_in_bad_fork():
        torch.cuda.manual_seed_all(seed)

    import torch.mps
    if not torch.mps._is_in_bad_fork():
        torch.mps.manual_seed(seed)

    if hasattr(torch, 'xpu') and not torch.xpu._is_in_bad_fork():
        torch.xpu.manual_seed_all(seed)

    _seed_custom_device(seed)

    return default_generator.manual_seed(seed)

A state dict holds parameters and buffers. Put a model in eval(), print the keys, and the batch-norm running statistics are there along with num_batches_tracked, while the thing that decides how those statistics get used is not.

model.training is a plain Python attribute on the module and every submodule, and it never enters serialization in either direction. Load the eval-mode dictionary above into a freshly built model and that model reports training as True, because construction set it and the load never touched it.

The direction that costs you is the other one. A resume script that runs a validation pass before continuing leaves every module in eval mode, dropout stops firing, batch norm switches to its running estimates, and the loop trains on. Lesson two measures what that does to a hundred steps of loss.

import torch
from torch import nn

model = nn.Sequential(nn.Linear(4, 4), nn.Dropout(0.5), nn.BatchNorm1d(4))
model.eval()
print(sorted(model.state_dict()))
# ['0.bias', '0.weight', '2.bias', '2.num_batches_tracked', '2.running_mean',
#  '2.running_var', '2.weight']
print(model.training, any("training" in k for k in model.state_dict()))   # False False

fresh = nn.Sequential(nn.Linear(4, 4), nn.Dropout(0.5), nn.BatchNorm1d(4))
fresh.load_state_dict(model.state_dict())
print(fresh.training)                # True: the flag did not travel

Serialization keeps storage sharing, which is a good property with an expensive corner. Save a ten-element slice of a million-element buffer and the file is 4,001,101 bytes. Clone it first and the same ten elements cost 1,170.

The zip listing shows where the size went. A torch checkpoint is a zip archive, and the record named view/data/0 holds 4,000,000 bytes, the whole storage the slice was reading through. The slice itself is 40 bytes of that, and the archive has no way to express a partial storage.

Any parameter carved out of one flat buffer has this shape, and so does any metric tensor sliced from a longer log. The symptom is a checkpoint many times larger than the parameter count justifies, and the fix is a clone on the way into the dictionary. The serialization note upstream documents the same behaviour with a nine-element example.

import os
import zipfile

import torch

flat = torch.zeros(1_000_000)          # one buffer the parameters were carved from
head = flat[:10]

torch.save({"head": head}, "view.pt")
torch.save({"head": head.clone()}, "clone.pt")
print(os.path.getsize("view.pt"), os.path.getsize("clone.pt"))   # 4001101 1170
print(head.numel() * head.element_size(), head.untyped_storage().nbytes())
# 40 4000000
with zipfile.ZipFile("view.pt") as z:
    print([(i.filename, i.file_size) for i in z.infolist() if i.file_size > 100])
# [('view/data.pkl', 168), ('view/data/0', 4000000)]

torch.save opens its destination for writing before it knows whether the write will succeed, so the previous checkpoint stops existing the moment the new save begins. Hand it a file object that fails on the fourth write, the way a full disk does, and the 3,120-byte checkpoint holding step 10 is a 64-byte fragment by the time the exception arrives.

Reading it back raises, which is the merciful half. PytorchStreamReader failed reading zip archive: failed finding central directory is a torn file announcing itself, because the zip central directory is written last and a partial archive has none. What you have lost is not detectable from this process at all: the run that could have resumed from step 10 no longer can.

Writing to a temporary name and renaming afterwards is the whole fix, and os.replace is the call that swaps the two names in one step within a filesystem. The same failing writer now wrecks safe.pt.partial and leaves safe.pt at 3,152 bytes and step 10, still loadable. Keeping the last two files rather than one is the other half, for the case where the failure is in the run rather than in the disk.

A checkpoint that exists is not a checkpoint that loads, and a save in progress is a checkpoint you do not have.

import os

import torch
from torch import nn


class Dying:
    """a disk that fills up on the fourth write"""

    def __init__(self, path):
        self.f = open(path, "wb")
        self.n = 0

    def write(self, chunk):
        self.n += 1
        if self.n > 3:
            raise OSError(28, "No space left on device")
        return self.f.write(chunk)

    def flush(self):
        self.f.flush()

    def close(self):
        self.f.close()


torch.manual_seed(0)
model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 1))


def ckpt(step):
    return {"model": model.state_dict(), "step": step}


torch.save(ckpt(10), "run.pt")                       # the last good checkpoint
print(os.path.getsize("run.pt"), torch.load("run.pt")["step"])   # 3120 10

try:
    torch.save(ckpt(20), Dying("run.pt"))            # the save that does not finish
except Exception as err:
    print(type(err).__name__, err)
print(os.path.getsize("run.pt"))
try:
    torch.load("run.pt")
except Exception as err:
    print(type(err).__name__, str(err)[:64])

# RuntimeError [enforce fail at inline_container.cc:595] . unexpected pos 64 vs 0
# 64
# RuntimeError PytorchStreamReader failed reading zip archive: failed finding c

Six lines, and the ordering inside them is what earns the promise. Write into the temporary file, flush Python's buffer, fsync so the bytes have left the page cache for the disk, close, and only then rename. Skip the fsync and a machine that loses power shortly after the rename can come back with the new name pointing at data that never landed.

The rename is the moment the checkpoint becomes real, and no reader can observe a half-renamed file. That is why the pattern works across processes as well as across crashes, and why the ranks in the next lesson can be made to agree with it.

def save_atomic(obj, path, opener=lambda p: open(p, "wb")):
    tmp = path + ".partial"
    f = opener(tmp)
    try:
        torch.save(obj, f)
        f.flush()
        os.fsync(f.fileno())
    finally:
        f.close()
    os.replace(tmp, path)          # one syscall: it either happened or it did not


save_atomic(ckpt(10), "safe.pt")
try:
    save_atomic(ckpt(20), "safe.pt", Dying)
except Exception as err:
    print(type(err).__name__, str(err)[:48])
print(torch.load("safe.pt")["step"], os.path.getsize("safe.pt"),
      os.path.getsize("safe.pt.partial"))

# RuntimeError [enforce fail at inline_container.cc:595] . unex
# 10 3152 64

Chapter four's arc showed the one piece of a loop with no state dict at all, and where a saved generator does and does not recover an order. A distributed run adds a number to that same gap, and it is easy to miss because nothing about it looks like state.

DistributedSampler shuffles from a seed plus an epoch counter it holds internally, and the counter only moves when you call set_epoch. Iterate the sampler twice without calling it and both epochs hand back [4, 7, 2, 1], the identical order. Call set_epoch(1) and the order becomes [5, 2, 7, 1]. Two ranks at the same epoch partition the dataset between them without overlap, which is the property that makes the shuffle correct in the first place.

So the epoch belongs in the checkpoint next to the step, and every rank has to restore the same one. A resume that restarts the counter at zero replays the first epoch's order for the rest of the run, and a resume where the ranks disagree about the epoch hands the same rows to two ranks and never shows the rest.

The mid-epoch position remains unsolved by anything in torch.utils.data itself. StatefulDataLoader, in the torchdata package, is the answer upstream has been building: a drop-in loader with state_dict and load_state_dict that records how far into an epoch the iteration got. It is a separate install and it is not on this machine, so it is named here rather than measured.

import torch
from torch.utils.data import DistributedSampler, TensorDataset

ds = TensorDataset(torch.arange(8.))
sampler = DistributedSampler(ds, num_replicas=2, rank=0, shuffle=True)

print(list(sampler), list(sampler))       # two epochs, no set_epoch call
sampler.set_epoch(1)
print(list(sampler))

r0 = DistributedSampler(ds, num_replicas=2, rank=0)
r1 = DistributedSampler(ds, num_replicas=2, rank=1)
r0.set_epoch(3)
r1.set_epoch(3)
print(list(r0), list(r1))

# [4, 7, 2, 1] [4, 7, 2, 1]
# [5, 2, 7, 1]
# [2, 3, 1, 6] [4, 5, 0, 7]