Checkpointing

Problem

A large training run lasts days or weeks on many machines, and something will stop it: a hardware fault, a preempted spot instance, a crashed process, a bug fix you need to deploy. Without a checkpoint all of that compute is lost. With a checkpoint that holds only the weights, training restarts, but as a different run: the optimizer forgets its momentum, the learning-rate schedule starts again from warmup, and the data order changes. Losses spike, results stop being reproducible, and a bug introduced by the resume is hard to tell apart from a real training problem.

Checkpoints also leave the machine that made them. People download weights, and PyTorch's default file format is a Python pickle, which can run code when it is loaded. And at scale a checkpoint is many gigabytes split across hundreds of ranks, so writing it takes real time, and the job that reads it back may not have the same number of GPUs.

The same word also names an unrelated memory technique. Activation checkpointing keeps only some of the activations from the forward pass and recomputes the rest during backward, so a deeper model or a bigger batch fits in memory. It is covered at the end of Mechanics.

Intuition

A checkpoint is a save file for a game with randomness in it. To continue exactly where you stopped, the save must hold the map and your inventory (the weights), but also how fast you are moving and in which direction (the optimizer's moments), where you are on the level's timer (the schedule), the state of the dice (the random number generators) and which cards you have already drawn from the deck (the position in the data). Leave any one out and you can still keep playing, but not the same game.

Mechanics

State dicts

PyTorch saves state as plain dictionaries, never as live objects. Module.state_dict() maps each parameter's and buffer's dotted name (layers.0.attn.q_proj.weight) to its tensor; load_state_dict(sd, strict=True) copies tensors back by name and, with strict=False, returns the missing and unexpected keys instead of raising. Because the file holds names and tensors, not code, the model class must be rebuilt by your code before loading.

Optimizer.state_dict() has two entries: state, the per-parameter tensors (for Adam: exp_avg, exp_avg_sq and step), and param_groups, the hyperparameters of each group with parameters replaced by integer ids. Its docstring is explicit that "the parameter itself is NOT saved"; load_state_dict matches ids to the new optimizer's parameters by order, so the parameters must be passed in the same order as when the checkpoint was made. A learning-rate scheduler's state_dict is its __dict__ minus the optimizer; restoring it (load_state_dict) restores how many steps have been taken. See optimizers and learning-rate schedules for what those states mean.

Random state and data position

Shuffling, dropout, data augmentation and random cropping all draw from generators. torch.get_rng_state() returns the CPU generator's state as a byte tensor and set_rng_state restores it; torch.cuda.get_rng_state_all() and set_rng_state_all do the same for every GPU. Python's random and NumPy have their own generators, saved separately.

The position in the data is the piece most often lost. torch.utils.data.DataLoader at this pin has no state_dict, so the epoch and the number of batches already consumed must be saved by your loop. The order itself can be made a pure function of saved numbers: DistributedSampler.__iter__ shuffles with a fresh generator seeded with self.seed + self.epoch, so after set_epoch(epoch) the order is reproduced exactly, and resuming mid-epoch means skipping the batches already seen. The plain RandomSampler, given no generator, draws a fresh seed from the global generator on every iter(), so its order depends on the RNG state at the start of the epoch. The worked example uses the seed + epoch scheme for the order and a separate stateful generator for input noise, so both pieces have to be restored.

Saving and loading files

torch.save(obj, f) pickles the object with protocol DEFAULT_PROTOCOL = 2 and writes tensor storages as separate entries in a zip archive (_use_new_zipfile_serialization=True by default). Write to a temporary file and rename it over the old one, so a crash mid-write never leaves a half-written file under the name your resume logic will read; keep the last few checkpoints rather than one, in case the newest is bad.

torch.load takes map_location (load GPU tensors onto the CPU or another device), mmap (map the file instead of reading it into memory) and weights_only. Unpickling can call any importable function, because a pickle can say "call this global with these arguments" (Python's __reduce__). Since 2.6 weights_only defaults to True; the error message says setting it to False "can result in arbitrary code execution". With weights_only=True a restricted unpickler resolves each GLOBAL opcode against an allowlist (tensor rebuild functions, dtypes, torch.Size, nn.Parameter, collections.OrderedDict, collections.Counter, and others), refuses the modules sys, os, posix and nt outright, and raises "Unsupported global" for anything else. Your own trusted classes can be added with torch.serialization.add_safe_globals. Formats such as safetensors avoid pickle entirely by storing only a header and raw tensor bytes.

How big, and how often

With mixed precision and Adam, each parameter carries a bf16 copy for compute (2 bytes), an fp32 master copy (4) and two fp32 moments (4 + 4): 14 bytes to resume, against 2 bytes for an inference export. The ZeRO paper counts the same 12 bytes of fp32 optimizer state plus 2 bytes each of half-precision weights and gradients, 16 in all, in memory; gradients are recomputed every step and are not saved.

Saving often wastes time writing; saving rarely wastes time redoing lost work. If writing a checkpoint stops training for C seconds and failures arrive on average every M seconds, saving every T seconds loses about C/T of the time to writing and, on a failure, half an interval of work on average. Young (1974) minimized that first-order waste:

waste(T) ≈ C / T + T / (2M)   ⇒   T* = √(2 C M)

At T* the two terms are equal. Making writes cheaper shortens the best interval and lowers the waste, which is why large jobs save asynchronously: dcp.async_save first copies the state to staging storage (CPU memory by default) and then runs the save in a separate thread, so training is blocked only for the copy. The formula assumes failures are independent and C is much smaller than M; a job on N machines has roughly N times the failure rate of one.

Distributed checkpoints

When FSDP or tensor parallelism shards parameters and optimizer state across ranks, gathering everything onto rank 0 to call torch.save is slow and may not fit in its memory. torch.distributed.checkpoint (DCP) instead has every rank write its own shards in parallel: dcp.save "handles ShardedTensor, and DTensor by having each rank only save their local shards", and a .metadata file records which piece of which tensor lives where. On dcp.load "each rank will try to read the least amount of data necessary"; the default load planner "handles resharding by issuing multiple read requests against storage", so a checkpoint saved on 64 GPUs can be loaded on 48. get_state_dict(model, optimizers) returns model and optimizer state dicts keyed by canonical names that "can be resharded with a different number of trainers and/or different parallelisms", and set_state_dict puts them back. DCP loads in place: tensors in the state dict you pass must already be allocated with their destination layout.

Activation checkpointing

Backpropagation needs the activations of every layer, so by default the forward pass keeps all of them, and activation memory grows with depth × batch × sequence length. torch.utils.checkpoint.checkpoint(fn, *args) runs fn without saving its intermediate tensors, keeps only its inputs, and runs fn again during backward to rebuild them; its docstring describes this as trading compute for memory. With preserve_rng_state=True (the default) it stashes and restores the RNG state so dropout masks in the recomputation match the original forward. Pass use_reentrant=False; the docstring recommends it, and it works with torch.autograd.grad and keyword arguments.

checkpoint_sequential(functions, segments, input) splits a sequence of n layers into k segments of ⌊n/k⌋ layers; all segments except the last are checkpointed and the last absorbs the remainder. At the peak, memory holds the k − 1 saved segment inputs plus one segment's activations, about k + n/k layers' worth, which is smallest near k = √n: memory grows with √n instead of n, for about one extra forward pass. That is the sublinear-memory result of Chen et al. (2016). Transformer libraries usually checkpoint each block, saving one tensor per block and recomputing attention and MLP intermediates.

Worked example

Implementation

import os, random
import torch
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
from torch.utils.checkpoint import checkpoint

def save(path, step, model, opt, sched, scaler, epoch, batches_done):
    state = {
        "model": model.state_dict(),
        "optimizer": opt.state_dict(),
        "scheduler": sched.state_dict(),
        "scaler": scaler.state_dict(),
        "rng": {"torch": torch.get_rng_state(),
                "cuda": torch.cuda.get_rng_state_all(),
                "python": random.getstate()},
        "data": {"epoch": epoch, "batches_done": batches_done},
        "step": step,
    }
    tmp = path + ".tmp"
    torch.save(state, tmp)
    os.replace(tmp, path)                       # atomic: never a half-written file at path

def resume(path, model, opt, sched, scaler, sampler):
    state = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(state["model"])        # model already built and on its device
    opt.load_state_dict(state["optimizer"])      # same parameter order as when saved
    sched.load_state_dict(state["scheduler"])
    scaler.load_state_dict(state["scaler"])
    torch.set_rng_state(state["rng"]["torch"])
    torch.cuda.set_rng_state_all(state["rng"]["cuda"])
    random.setstate(state["rng"]["python"])
    sampler.set_epoch(state["data"]["epoch"])    # DistributedSampler: order = seed + epoch
    return state["step"], state["data"]["epoch"], state["data"]["batches_done"]

# Sharded models (FSDP, tensor parallel): every rank calls these.
def save_sharded(ckpt_dir, model, opt):
    model_sd, optim_sd = get_state_dict(model, opt)
    dcp.save({"model": model_sd, "optim": optim_sd}, checkpoint_id=ckpt_dir)

def load_sharded(ckpt_dir, model, opt):        # world size may differ from the save
    model_sd, optim_sd = get_state_dict(model, opt)
    state = {"model": model_sd, "optim": optim_sd}
    dcp.load(state, checkpoint_id=ckpt_dir)      # fills the tensors in place
    set_state_dict(model, opt, model_state_dict=state["model"], optim_state_dict=state["optim"])

# Activation checkpointing: keep each block's input, recompute its insides in backward.
class Stack(torch.nn.Module):
    def __init__(self, blocks):
        super().__init__()
        self.blocks = torch.nn.ModuleList(blocks)

    def forward(self, x):
        for block in self.blocks:
            x = checkpoint(block, x, use_reentrant=False)
        return x

In the loop, skip batches_done batches of the resumed epoch before training, and call save every T* seconds of wall time rather than every fixed number of steps if step times vary. The skip costs reading those batches again; a loader that saves its own position avoids that.

Tradeoffs

Connections

The state being saved is defined on other pages: the loop that produces it on training loop, Adam's moments on optimizers, warmup and cosine on learning-rate schedules, the grad scaler and master weights on mixed precision, and DistributedSampler on data loading and batching. Sharded checkpoints follow the layouts of distributed training. Activation checkpointing trades against the memory that backpropagation keeps in the computation graph, and competes with gradient accumulation as the way to fit a bigger effective batch.

References

• Young. A First Order Approximation to the Optimum Checkpoint Interval. Communications of the ACM 17(9), 1974. doi:10.1145/361147.361115
• Chen, Xu, Zhang and Guestrin. Training Deep Nets with Sublinear Memory Cost. 2016. arXiv:1604.06174
• Rajbhandari, Rasley, Ruwase and He. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. SC 2020. arXiv:1910.02054
• Kingma and Ba. Adam: A Method for Stochastic Optimization. ICLR 2015. arXiv:1412.6980
• Python documentation: pickle (including its warning that unpickling can execute arbitrary code) and pickletools.
• PyTorch v2.15.0: torch.save, torch.load, the weights_only allowlist, Module.state_dict, Optimizer.state_dict, DistributedSampler, dcp.save, dcp.load, get_state_dict, checkpoint, checkpoint_sequential.