How model checkpoints are saved during training

Model checkpoint training is the safety net that lets you recover from hardware failures, spot preemption, or gradient explosions without restarting from scratch. Most engineers treat checkpoints as an afterthought until a 72-hour run dies at hour 71. This guide covers the formats, scheduling strategies, and recovery patterns that actually hold up in production.

Choose the right serialization format

PyTorch’s default torch.save pickles the entire object graph — model, optimizer, scheduler, RNG states, and any custom objects attached to the state dict. This works for single-GPU experiments but creates three problems at scale: files exceed filesystem limits, loading requires the exact same code version, and you cannot inspect weights without the training environment.

Use safetensors instead. It stores tensors in a flat, memory-mappable format with a JSON header describing shapes and dtypes. No pickle, no arbitrary code execution, and you can stream individual tensors without loading the whole file.

# Saving with safetensors
from safetensors.torch import save_file

state_dict = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scheduler": scheduler.state_dict(),
    "epoch": epoch,
    "global_step": global_step,
    "rng_state": torch.get_rng_state(),
    "cuda_rng_state": torch.cuda.get_rng_state_all(),
}
save_file(state_dict, f"checkpoint-{global_step}.safetensors")

Loading is equally straightforward and works across PyTorch versions:

from safetensors.torch import load_file

checkpoint = load_file("checkpoint-15000.safetensors")
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])

Pitfall: safetensors does not preserve tensor metadata like requires_grad or _backward_hooks. Re-attach these after loading if your training loop depends on them.

Shard checkpoints for multi-GPU training

When training across multiple GPUs with DistributedDataParallel or FSDP, each rank holds a shard of the model and optimizer state. Saving a single consolidated file requires gathering all shards to rank 0 — a synchronization bottleneck that wastes GPU memory and time.

Instead, save per-rank shards and consolidate only when needed (e.g., for inference or checkpoint transfer). FSDP’s StateDictType.SHARDED_STATE_DICT handles this automatically:

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import StateDictType

# Configure once before training
FSDP.set_state_dict_type(
    model,
    StateDictType.SHARDED_STATE_DICT,
    StateDictConfig(offload_to_cpu=True, rank0_only=False),
)

# Each rank saves its own shard
checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, f"checkpoint-{global_step}-rank{dist.get_rank()}.pt")

For inference, consolidate on demand using StateDictType.FULL_STATE_DICT with rank0_only=True. This avoids the consolidation cost during training entirely.

Tradeoff: Sharded checkpoints are faster to write but require the same world size and FSDP configuration to reload. If you change GPU count or sharding strategy, you must consolidate first.

Schedule checkpoints by time, not just steps

Step-based scheduling (every N steps) seems natural but creates two failure modes: you either checkpoint too frequently and saturate the filesystem, or too rarely and lose hours of work. Time-based scheduling adapts to variable step times — especially important when using dynamic batching, gradient accumulation, or mixed-precision with loss scaling.

A practical hybrid: checkpoint at least every 30 minutes and at most every 5000 steps, whichever comes first. Also checkpoint on validation improvement and before any learning rate schedule milestone.

import time
from pathlib import Path

class CheckpointScheduler:
    def __init__(
        self,
        save_dir: Path,
        min_interval_sec: int = 1800,      # 30 minutes
        max_steps: int = 5000,
        keep_last_n: int = 3,
        keep_best_n: int = 2,
    ):
        self.save_dir = save_dir
        self.min_interval_sec = min_interval_sec
        self.max_steps = max_steps
        self.keep_last_n = keep_last_n
        self.keep_best_n = keep_best_n
        self.last_save_time = time.time()
        self.last_save_step = 0
        self.best_metrics = []

    def should_save(self, step: int, metric: float | None = None) -> bool:
        now = time.time()
        time_elapsed = now - self.last_save_time
        steps_elapsed = step - self.last_save_step

        if time_elapsed >= self.min_interval_sec:
            return True
        if steps_elapsed >= self.max_steps:
            return True
        if metric is not None and self._is_best(metric):
            return True
        return False

    def _is_best(self, metric: float) -> bool:
        if len(self.best_metrics) < self.keep_best_n:
            return True
        return metric < max(self.best_metrics)  # lower is better

    def save(self, step: int, state_dict: dict, metric: float | None = None):
        # ... save logic ...
        self.last_save_time = time.time()
        self.last_save_step = step
        if metric is not None:
            self.best_metrics.append(metric)
            self.best_metrics.sort()
            self.best_metrics = self.best_metrics[:self.keep_best_n]
        self._prune_old_checkpoints(step)

Implement atomic writes and corruption guards

Partial writes corrupt checkpoints. A filesystem crash mid-write leaves a truncated file that torch.load or safetensors.load_file will fail to parse — or worse, silently load with missing tensors. Write to a temporary file, verify integrity, then atomic rename.

import tempfile
import os
import hashlib

def atomic_save(state_dict: dict, path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    
    with tempfile.NamedTemporaryFile(
        mode="wb", dir=path.parent, delete=False, suffix=".tmp"
    ) as tmp:
        tmp_path = Path(tmp.name)
        save_file(state_dict, tmp_path)
        
        # Verify by loading back
        _ = load_file(tmp_path)
        
        # Optional: compute checksum for later verification
        checksum = hashlib.sha256(tmp_path.read_bytes()).hexdigest()
        (tmp_path.with_suffix(".sha256")).write_text(checksum)
    
    # Atomic on POSIX; replace on Windows
    os.replace(tmp_path, path)
    os.replace(tmp_path.with_suffix(".sha256"), path.with_suffix(".sha256"))

Store the SHA-256 alongside the checkpoint. On load, verify before deserializing:

def verified_load(path: Path) -> dict:
    checksum_path = path.with_suffix(".sha256")
    if checksum_path.exists():
        expected = checksum_path.read_text().strip()
        actual = hashlib.sha256(path.read_bytes()).hexdigest()
        if actual != expected:
            raise RuntimeError(f"Checksum mismatch for {path}")
    return load_file(path)

Pitfall: Network filesystems (NFS, EFS, S3 via FUSE) may not honor os.replace atomicity. For cloud storage, upload the verified temp file directly via SDK (e.g., boto3 multipart upload) rather than relying on POSIX semantics.

Prune aggressively but keep recovery anchors

Disk fills fast. A 7B parameter model with optimizer states in FP32 consumes ~56 GB per checkpoint. Keep three recent checkpoints, two best-by-metric, and one “anchor” every 50k steps for long-run recovery. Delete everything else.

def _prune_old_checkpoints(self, current_step: int):
    all_ckpts = sorted(self.save_dir.glob("checkpoint-*.safetensors"))
    if len(all_ckpts) <= self.keep_last_n + self.keep_best_n + 1:
        return

    # Identify protected checkpoints
    protected = set()
    # Recent
    protected.update(all_ckpts[-self.keep_last_n:])
    # Best (track separately in scheduler)
    # Anchors: every 50k steps
    for ckpt in all_ckpts:
        step = int(ckpt.stem.split("-")[1])
        if step % 50000 == 0:
            protected.add(ckpt)

    for ckpt in all_ckpts:
        if ckpt not in protected:
            ckpt.unlink(missing_ok=True)
            ckpt.with_suffix(".sha256").unlink(missing_ok=True)

Anchor checkpoints matter because optimizer state (momentum buffers, Adam second moments) diverges over time. Restoring from a 100k-step-old checkpoint with a fresh optimizer often destabilizes training. Anchors let you resume with matching optimizer state.

Handle optimizer state correctly

The optimizer state dict contains per-parameter tensors (exp_avg, exp_avg_sq for AdamW) that must align exactly with the model’s parameter order and shapes. Two common mismatches break resumption:

1. Parameter reordering: Adding/removing layers or using torch.compile with dynamic shapes can change parameter registration order. Always save model.state_dict() and optimizer.state_dict() in the same forward pass without intervening model modifications.

2. FSDP parameter flattening: FSDP flattens parameters into a single contiguous buffer. The optimizer state dict keys are flattened parameter IDs, not module paths. When consolidating for inference, use FSDP.optim_state_dict_to_load to map correctly.

# Correct pattern: capture both atomically
with torch.no_grad():
    model_state = model.state_dict()
    optim_state = optimizer.state_dict()

save_file({"model": model_state, "optimizer": optim_state}, path)

Pitfall: Gradient accumulation steps do not affect optimizer state — the optimizer only updates on optimizer.step(). But if you checkpoint mid-accumulation, you must also save accumulation_step and the accumulated gradients if using gradient clipping that depends on accumulated norm.

Resume deterministically

Resuming requires restoring not just weights but the entire stochastic state: Python random, NumPy, PyTorch CPU and CUDA RNGs, and the data loader sampler position. Missing any of these introduces non-determinism that compounds.

def load_checkpoint(path: Path, model, optimizer, scheduler, train_loader):
    checkpoint = verified_load(path)
    
    model.load_state_dict(checkpoint["model"])
    optimizer.load_state_dict(checkpoint["optimizer"])
    scheduler.load_state_dict(checkpoint["scheduler"])
    
    torch.set_rng_state(checkpoint["rng_state"])
    torch.cuda.set_rng_state_all(checkpoint["cuda_rng_state"])
    # If using numpy/random in data pipeline
    np.random.set_state(checkpoint["numpy_rng_state"])
    random.setstate(checkpoint["python_rng_state"])
    
    # Restore data loader position
    train_loader.sampler.set_epoch(checkpoint["epoch"])
    # For IterableDataset with sharding, restore iterator state
    if hasattr(train_loader.dataset, "set_state"):
        train_loader.dataset.set_state(checkpoint["dataset_state"])
    
    return checkpoint["epoch"], checkpoint["global_step"]

For DistributedSampler, call set_epoch(epoch) before creating the iterator each epoch. For IterableDataset with worker sharding, implement __getstate__/__setstate__ to persist the current shard and offset.

Offload to object storage asynchronously

Local NVMe fills; network storage adds latency. The pattern that works: write locally, verify, then enqueue an async upload to S3/GCS. Use a background thread or separate process so training never blocks on network I/O.

import queue
import threading
import boto3

upload_queue = queue.Queue()
s3 = boto3.client("s3")

def uploader_worker():
    while True:
        local_path, s3_key = upload_queue.get()
        try:
            s3.upload_file(str(local_path), "my-checkpoint-bucket", s3_key)
            # Optionally delete local after confirmed upload
        except Exception as e:
            logging.error(f"Upload failed for {local_path}: {e}")
            # Re-queue or alert
        finally:
            upload_queue.task_done()

threading.Thread(target=uploader_worker, daemon=True).start()

# In checkpoint save:
atomic_save(state_dict, local_path)
upload_queue.put((local_path, f"checkpoints/{run_name}/{local_path.name}"))

Tradeoff: Async upload means the latest checkpoint may not be in object storage when a node fails. Accept this risk or synchronously upload anchor checkpoints only.

Test recovery before you need it

Add a CI job that: (1) trains for 100 steps, (2) kills the process, (3) restarts from the latest checkpoint, (4) trains 100 more steps, (5) verifies loss matches a non-interrupted baseline within tolerance. This catches RNG state bugs, data loader misalignment, and optimizer state corruption before they waste a real run.

# test_recovery.sh
set -e
python train.py --max-steps 100 --checkpoint-dir /tmp/ckpt-test
python train.py --resume-from /tmp/ckpt-test --max-steps 200
python compare_losses.py baseline.log resumed.log --tolerance 1e-5

Run this on every dependency upgrade (PyTorch, transformers, CUDA driver). Checkpoint format compatibility breaks silently.

Summary checklist

• Use safetensors for all new checkpoints; migrate legacy .pt files
• Save sharded state dicts with FSDP; consolidate only for inference
• Schedule by time (30 min minimum) and steps (5k maximum), plus best-metric and anchors
• Write atomically with temp-file + rename; verify with SHA-256
• Prune to 3 recent + 2 best + anchors every 50k steps
• Capture model, optimizer, scheduler, all RNG states, data loader position in one atomic snapshot
• Offload to object storage asynchronously; sync-upload anchors
• Test kill-and-resume in CI on every environment change

Model checkpoint training is infrastructure, not experimentation. Treat it with the same rigor as your data pipeline — because when a training run fails, the checkpoint is the only thing that stands between you and starting over.