How long a mistake stays quiet

The failures that get fixed are the ones that raise. Six of the pytorch museum's exhibits are step-time exceptions, and every one of them names its own cause in the message: a shape that does not fit the layer, class labels arriving as floats, a backward through a graph that was already freed. They cost minutes.

One step further out are the failures that produce a number and the wrong number. An optimizer state dict from a differently shaped model loads without complaint and fails inside the update, which chapter four's first lesson takes apart. A missing RNG entry gives a first resumed loss that is simply different, measured in the same arc.

Past that the signal gets thin. A missing scheduler is exact for two steps and then bends. A missing optimizer state is exact for one. A mid-epoch loader position never announces itself at all. The rest of this lesson lives at that end of the scale, because a failure that takes a hundred steps to show is a failure people attribute to the data.

Twenty steps of an Adam loop with dropout in it, killed at ten, restored with the model, the optimizer and the RNG. The resumed ten losses equal the last ten of the run that was never interrupted, as a list comparison. That is the shape of the proof, and it is the same shape LAB·P4 runs over a TPU bridge with its own numbers.

Now leave the model in eval mode, which is what a validation pass before the resume does if nobody calls train() afterwards. The first resumed loss reads 6.386841 where the correct run reads 6.504373. It is lower, it is plausible, and the run continues descending from there.

Two things happened at once and both are silent. Dropout stopped zeroing activations, so the model in front of the loss is a different model. And dropout stopped drawing from the generator, so every batch index after the first is a different batch. The loop has no way to notice either.

import torch
from torch import nn

torch.manual_seed(1)
DATA = torch.randn(512, 16)
TARGET = DATA.sum(dim=1, keepdim=True)


def build():
    torch.manual_seed(0)
    model = nn.Sequential(nn.Linear(16, 64), nn.ReLU(), nn.Dropout(0.1), nn.Linear(64, 1))
    return model, torch.optim.AdamW(model.parameters(), lr=1e-2)


def run(model, opt, n):
    out = []
    for _ in range(n):
        idx = torch.randint(0, 512, (32,))
        loss = nn.functional.mse_loss(model(DATA[idx]), TARGET[idx])
        opt.zero_grad()
        loss.backward()
        opt.step()
        out.append(round(loss.item(), 6))
    return out


model, opt = build()                       # the run that is never interrupted
torch.manual_seed(7)
straight = run(model, opt, 20)

model, opt = build()                       # the same run, killed at step 10
torch.manual_seed(7)
run(model, opt, 10)
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
            "step": 10, "rng": torch.get_rng_state()}, "ck.pt")


def resume(train_mode=True, restore_opt=True):
    state = torch.load("ck.pt")
    model, opt = build()
    model.load_state_dict(state["model"])
    if restore_opt:
        opt.load_state_dict(state["opt"])
    torch.set_rng_state(state["rng"])
    if not train_mode:
        model.eval()                       # the validation pass nobody turned off
    return run(model, opt, 10)


print(resume() == straight[10:])
print(straight[10:13])
print(resume(train_mode=False)[:3])
print(resume(restore_opt=False)[:3])

# True
# [6.504373, 9.049624, 2.990426]
# [6.386841, 9.498335, 2.94874]
# [6.504373, 9.063147, 3.139434]

Stretch the same experiment to two hundred steps with the kill at a hundred and the three resumes end up in places that all look like success. The whole-state resume equals the uninterrupted tail exactly. The eval-mode resume ends twelve times lower, at a last-twenty mean of 0.011072 against 0.137139. The resume that dropped the optimizer state ends at 0.139627, within two percent of the correct run.

Read the third row before the second. Dropping Adam's moments produced a first resumed loss of 0.262158, identical to the correct one, because the loss is measured before the update and a missing moment cannot reach it until the step after. By the end of a hundred steps the two runs are indistinguishable by any summary you would put in a report, while the worst individual step sits 1.5058 times its own value away from the correct one.

The second row is the more uncomfortable one. A broken resume that converges better than the correct run will survive every review that asks whether the loss went down, and this is the ordinary case rather than a contrived one: turning dropout off improves a training loss on a small model almost by definition.

Every one of these runs converged. Only one of them resumed.

from statistics import mean

model, opt = build()
torch.manual_seed(7)
long_run = run(model, opt, 200)

model, opt = build()
torch.manual_seed(7)
run(model, opt, 100)
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
            "step": 100, "rng": torch.get_rng_state()}, "ck.pt")

tail = long_run[100:]                      # resume() above, with 100 in place of 10
for name, got in [("whole state", resume()),
                  ("eval mode", resume(train_mode=False)),
                  ("no optimizer", resume(restore_opt=False))]:
    a, b = torch.tensor(got), torch.tensor(tail)
    print(f"{name:14} equal={got == tail} step101={got[0]} "
          f"last20={round(mean(got[-20:]), 6)} max_rel={((a - b).abs() / b.abs()).max():.4f}")

# whole state    equal=True step101=0.262158 last20=0.137139 max_rel=0.0000
# eval mode      equal=False step101=0.09149 last20=0.011072 max_rel=0.9674
# no optimizer   equal=False step101=0.262158 last20=0.139627 max_rel=1.5058

The bar the chapter sets is an overlay: the resumed curve on top of the uninterrupted one, the same curve rather than two curves that end up close. Turning that into something a machine can fail is a one-line change, and worth making, because a human looking at a plot cannot express a tolerance and will accept anything with the right shape.

Compared point by point, both broken resumes above are easy to catch. The correct one passes allclose down to a relative tolerance of one part in a million. The eval-mode resume needs a tolerance of 1.0 to pass, and the optimizer-free one fails at every tolerance in the list. So the numbers are not subtle. What is subtle is that nobody compares a hundred pairs of numbers by hand, and every summary of those same hundred numbers agrees.

The comparison worth writing down is therefore a list against a list, taken from a run you kept for the purpose. Save the uninterrupted losses once, resume, compare with ==, and record the index of the first disagreement rather than a verdict. On a run long enough that keeping the whole list is awkward, compare a fixed window after the restore and say in the report how many steps it covered.

LAB·P4 closes the TPU version of this on a v6e-1: a kill at step 99 with a loss of 0.01456 and a first resumed loss of 0.01437, published as the lab's reference block. Those two numbers are the lab's, not this page's, and the reason they are quoted rather than re-derived is that they come from hardware this machine does not have.

def survives(got, rtol):
    return torch.allclose(torch.tensor(got), torch.tensor(tail), rtol=rtol, atol=0)


for name, got in [("whole state", resume()),
                  ("eval mode", resume(train_mode=False)),
                  ("no optimizer", resume(restore_opt=False))]:
    print(name, [r for r in (1e-6, 1e-3, 1e-2, 1e-1, 1.0) if survives(got, r)])

# whole state [1e-06, 0.001, 0.01, 0.1, 1.0]
# eval mode [1.0]
# no optimizer []

One infinite target is enough. The loss comes back inf, the backward fills every gradient with inf or nan, and one optimizer step turns all 1,153 parameters of this model into nan. Nothing raises at any point.

Checkpoint at that moment and the file is a permanent record of a dead run. Restoring it gives [nan, nan, nan] for the next three losses, and it will give nan forever, because there is no arithmetic that brings a nan weight back. A resume from the checkpoint before it is the only recovery, which is the argument for keeping more than one.

A guard costs one line at the save site: refuse to write a checkpoint whose loss is not finite. Chapter four's second lesson covers the other half of this, the clip that returns inf and quietly zeroes every gradient except the poisoned one, and the flag that turns it into an error. Between the two of them a run stops rather than saving over its own last good state.

import math

model, opt = build()
torch.manual_seed(7)
run(model, opt, 5)

bad = nn.functional.mse_loss(model(DATA[:32]), TARGET[:32] * float("inf"))
opt.zero_grad()
bad.backward()
opt.step()
print(bad.item(), math.isfinite(bad.item()))          # inf False

weights = model.state_dict()
print(sum(int(v.isnan().sum()) for v in weights.values()),
      sum(v.numel() for v in weights.values()))       # 1153 1153

torch.save({"model": weights, "opt": opt.state_dict(), "step": 6}, "poison.pt")
state = torch.load("poison.pt")
model, opt = build()
model.load_state_dict(state["model"])
opt.load_state_dict(state["opt"])
torch.manual_seed(7)
print(run(model, opt, 3))                             # [nan, nan, nan]