Your PyTorch job resumed from a checkpoint and the loss jumped — what did it omit?

## Two different artifacts "Save the model" means two incompatible things. A **serving** checkpoint needs only the weights: `torch.save(model.state_dict(), path)`. A **resume** checkpoint needs everything required to make step N+1 identical to what it would have been had the job never stopped. Conflating them is why resumed runs show a loss spike, a visible discontinuity in the curve at the restart point, and sometimes permanently worse final quality. ## What a resume checkpoint contains - `model.state_dict()` — parameters **and** buffers. Buffers matter: BatchNorm's `running_mean` and `running_var` are in here and are not owned by the optimizer. - `optimizer.state_dict()` — two parts: `param_groups` (the learning rate, weight decay, betas as they currently stand) and `state` (the per-parameter buffers). - `scheduler.state_dict()` — mostly `last_epoch`, the counter the schedule's formula is evaluated at. - `scaler.state_dict()` — the current AMP loss scale and growth tracker. - `epoch` and `global_step` — plain integers, needed to resume the schedule of anything driven by step count, and to keep logging aligned. - Optionally RNG states from `torch.get_rng_state()`, `torch.cuda.get_rng_state()`, plus the Python and NumPy generators, if you want bitwise reproducibility of augmentation and dropout. ## Why the optimizer state matters most Adam and AdamW keep, per parameter, an exponential moving average of the gradient (`exp_avg`) and of its square (`exp_avg_sq`), plus a step count used for bias correction. Those buffers are the accumulated knowledge of thousands of iterations about the curvature and scale of each coordinate. Drop them and the optimizer restarts with zeroed moments. The bias correction term makes the very first steps after such a restart unusually large, and the per-coordinate scaling is wrong until the moving averages refill — which takes on the order of `1/(1-beta2)` steps, thousands at the usual `beta2 = 0.999`. The result is exactly the reported symptom: the loss jumps at the resume point and takes a while to recover. Plain SGD with momentum has a smaller version of the same problem via its `momentum_buffer`. ## The scheduler and the scaler A scheduler is stateless apart from its counter, so restoring `last_epoch` is all it takes — but forgetting it is expensive. A cosine schedule restarted from zero jumps the learning rate back to its peak; a warmup schedule restarted from zero re-runs warmup on a model that is well past that phase. Both produce a loss spike that looks identical to the optimizer-state one, so check both. The AMP scaler is a smaller effect: without its state the scaler restarts from its default scale and may skip a handful of steps while it re-probes. Not catastrophic, but free to fix. ## Loading order and devices The robust idiom is to load onto CPU and let each object place its own tensors: 1. Construct the model with the same architecture, then `model.load_state_dict(ckpt["model"])`, then move the model to the device. 2. Construct the optimizer **over the moved parameters**, then `optimizer.load_state_dict(ckpt["optimizer"])` — it casts its state tensors to match each parameter's device. 3. Construct the scheduler over that optimizer and load its state. `map_location="cpu"` in `torch.load` avoids the classic failure of a checkpoint saved from `cuda:3` refusing to load on a box with two GPUs. `load_state_dict` is strict by default and raises on any missing or unexpected key; passing `strict=False` returns a named tuple of `missing_keys` and `unexpected_keys` instead, which is the right tool when deliberately loading a backbone into a model with a new head — and the wrong tool for silencing a mismatch you have not understood, because it will happily load nothing at all. ## weights_only in torch 2.6 and later `torch.load` used to unpickle arbitrarily, which meant loading an untrusted checkpoint could execute code. Since PyTorch 2.6 the `weights_only` argument defaults to `True`, restricting deserialization to tensors and a safe set of primitive containers. Dictionaries of state dicts, ints, floats and lists load fine under that default. What breaks is a checkpoint into which someone pickled a custom class, an argparse namespace, or a whole `nn.Module` object — those now raise, and the fixes are to allowlist the type or, for a checkpoint you trust and control, to pass `weights_only=False` explicitly. The durable fix is not to pickle live objects into checkpoints in the first place. ## What you still cannot restore Even a complete checkpoint does not make a resumed run bit-identical unless you also restore the position within the epoch's data order. The usual pragmatic compromise is to checkpoint on epoch boundaries so the sampler restarts cleanly, and to accept that mid-epoch resume replays some samples. Nondeterministic kernels and different GPU counts are further sources of divergence — worth stating explicitly when someone asks for "exact" resume.