DL S055 · Checkpointing and Resuming

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 55 / 80 · Module M09 · ~2 hours

The story we're starting with

Checkpointing sounds boring until the first time it saves you $200 of GPU time. Then it becomes the most important thing you learn all year.

Here's the constraint: modern LLM pretraining is long. Even our tiny 100M run is 2.7 hours. A 1B run is a day. Real 7B runs are a week or more. Somewhere in that window, something will fail: your spot pod gets reclaimed, your disk fills up, PyTorch's memory allocator hiccups, the network partitions during a multi-node run, or you kick the power strip. If a failure means "start from step 0", you can never train anything serious.

The fix is checkpointing — periodically saving the training state to disk in a form that resumes cleanly. Sounds simple. It's not. What exactly is "state"? Just the weights? No — you also need the optimizer's momentum, the LR scheduler's step count, the dataloader's position in the current shard, and (subtle bug source) the random state so your data ordering resumes deterministically. Miss any one of these and your resumed run either crashes, produces a subtly different loss curve, or drifts from the intended schedule.

This session builds a checkpoint system that gets every corner right. By the end you'll have a save_checkpoint + load_checkpoint pair that can survive any kill and resume within a few seconds of where you were.

Prerequisites

• S053 (100M pretraining) — you have a training loop that calls save_checkpoint().
• S051 (DDP/FSDP) — for the sharded-checkpoint discussion at the end.

1 · What does "training state" actually contain?

Four things. Miss any, resume is broken.

Plus one thing that's nice to save:

• Dataloader position — if you're doing sequential iteration through shards, save which shard and byte offset. For our random-sampling dataset (S053), this is implicit in the RNG state.

2 · The 30-line save/load pair

import torch, os, glob, json, numpy as np
 
def save_checkpoint(model, optimizer, step: int, cfg, path: str, final: bool = False):
    payload = {
        "step": step,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "torch_rng": torch.get_rng_state(),
        "cuda_rng": torch.cuda.get_rng_state(),
        "numpy_rng": np.random.get_state(),
        "config": vars(cfg),
    }
    tmp = path + ".tmp"
    torch.save(payload, tmp)
    os.rename(tmp, path)        # atomic on POSIX — no half-written checkpoints
    # Also write a symlink so 'latest.pt' is always current
    latest = os.path.join(os.path.dirname(path), "latest.pt")
    if os.path.islink(latest):
        os.unlink(latest)
    os.symlink(os.path.basename(path), latest)
 
 
def load_checkpoint(path: str, model, optimizer) -> int:
    if not os.path.exists(path):
        return 0
    payload = torch.load(path, map_location="cpu")
    model.load_state_dict(payload["model"])
    optimizer.load_state_dict(payload["optimizer"])
    torch.set_rng_state(payload["torch_rng"])
    torch.cuda.set_rng_state(payload["cuda_rng"])
    np.random.set_state(payload["numpy_rng"])
    return payload["step"] + 1

Three things worth pausing on.

The atomic rename. torch.save can take seconds for a 100M model. If your process gets killed mid-save, you get a half-written .pt file that later crashes torch.load. Writing to .tmp and then os.rename is atomic on POSIX — either the file exists complete or it doesn't exist at all. Never write directly to the final path.

The latest.pt symlink. Your training loop only needs to know one path — latest.pt. Save creates it fresh every time; resume just loads latest.pt. No timestamped filename juggling.

Map to CPU on load. Even if resuming on GPU. Loading to CPU first avoids a spike of 2× model VRAM usage (once on CPU as it deserializes, once on GPU as it copies). Then .to(device) after.

3 · Rolling window — don't fill the disk

If you checkpoint every 500 steps and keep everything, a 100M model at ~1.5 GB per checkpoint × 15 checkpoints = 22 GB. Manageable. A 1B model at 15 GB × 15 = 225 GB. Not manageable.

Keep the last K checkpoints + one "milestone" checkpoint every N steps:

def cleanup_old_checkpoints(ckpt_dir: str, keep: int = 3, milestone_every: int = 5000):
    files = sorted(glob.glob(f"{ckpt_dir}/step-*.pt"))
    to_keep = set(files[-keep:])                       # last K
    for f in files:
        step = int(os.path.basename(f).split("-")[1].split(".")[0])
        if step % milestone_every == 0:                # milestone
            to_keep.add(f)
    for f in files:
        if f not in to_keep:
            os.remove(f)

Call after every save. Now disk stays bounded: ~keep × ckpt_size + (total_steps / milestone_every) × ckpt_size.

For our 100M @ 7630 steps, keep=3, milestone_every=2500:

• Rolling: 3 × 1.5 GB = 4.5 GB
• Milestones: ~3 × 1.5 GB = 4.5 GB
• Total: ~9 GB. Comfortable.

4 · Save what to disk?

Two dtype tradeoffs:

Save in fp32 (model + optimizer master): biggest, safest. Optimizer state must be fp32 anyway (Adam m/v are variance-sensitive). Model weights in fp32 preserve full precision so you can resume in a different dtype later.

Save in bf16 (model weights only): half the size. Loses precision, but for inference-only checkpoints this is fine. Never save optimizer state in bf16 — Adam's variance term drifts if you round it.

Reasonable pattern:

# For resume-training checkpoints
payload["model"] = model.state_dict()                  # native (bf16 if you compile in bf16)
payload["optimizer"] = optimizer.state_dict()          # fp32 by default
 
# For a "final release" checkpoint that only needs inference
final_payload = {"model": {k: v.to(torch.bfloat16) for k, v in model.state_dict().items()}}
torch.save(final_payload, "final-bf16.pt")

The final "release" checkpoint is smaller and safe to publish (no optimizer moments, no RNG).

5 · The subtlest bug in checkpointing: resume drift

Symptom: you resume from step 3000, training continues, but the loss curve doesn't match the original run's continuation. Losses are similar but not identical, and the divergence grows over time.

Cause: something in the state didn't restore exactly. Common culprits:

1. Missing RNG restore. Your dataloader samples different batches on the resumed run.
2. Different torch.manual_seed in the boot code. Set it before loading the checkpoint; the checkpoint's torch_rng restore then overrides your seed's state to the checkpointed one.
3. torch.backends.cudnn.deterministic = False on one run, True on the other. This changes convolution kernel selection non-deterministically.
4. Optimizer state not fully restored. Adam's step counter (for bias correction) is inside optimizer state — if you saved from an older PyTorch and load in newer, the field name may have changed.

Test: after resuming, run ONE step and log the loss. Then compare to the original run's loss at the corresponding step. If different by more than 0.001, some state didn't survive. Fix before letting the run continue.

6 · Sharded checkpoints for FSDP/big models

Under FSDP, each rank owns only 1/N of the model. There are two ways to save:

Rank-0 gather + save: All ranks send their shards to rank 0, which assembles the full model and saves it. Simple, produces a "unified" checkpoint any downstream tool can load. Cost: rank 0 needs enough RAM/VRAM to hold the entire model temporarily — infeasible past ~30B.

Sharded save (each rank saves its shard): Every rank writes a shard-{rank}.pt file. Fast, memory-efficient. Cost: loading requires the same world_size to rejoin the shards, or a "resharding" step.

PyTorch's torch.distributed.checkpoint (DCP) handles sharded saves cleanly:

import torch.distributed.checkpoint as dcp
 
# Save
state = {"model": model.state_dict(), "optim": optim.state_dict()}
dcp.save(state, checkpoint_id="/workspace/ckpt/step-1000")
 
# Load (works even if world_size differs from save time)
dcp.load(state, checkpoint_id="/workspace/ckpt/step-1000")
model.load_state_dict(state["model"])

The DCP format is a directory of shards + a metadata file. Not human-portable, but resilient and reshardable.

For our 100M model (fits in one GPU), stick with the plain torch.save pattern. DCP only pays off past ~1B.

7 · The auto-resume launcher

Combine everything into a shell-loop wrapper that treats spot kills as normal:

#!/bin/bash
# train_forever.sh
CKPT_DIR=/workspace/checkpoints
 
while true; do
    python train.py --resume_from ${CKPT_DIR}/latest.pt
    EXIT=$?
    if [ -f ${CKPT_DIR}/DONE ]; then
        echo "Training complete at $(date)"
        break
    fi
    echo "Trainer exited with code ${EXIT} at $(date). Restarting in 10s."
    sleep 10
done

Your training script:

def main(args):
    start_step = load_checkpoint(args.resume_from, model, optimizer)
    for step in range(start_step, total_steps):
        # ... training loop
        pass
    open(f"{cfg.ckpt_dir}/DONE", "w").close()   # sentinel: don't restart

Combined with Runpod's "auto-restart on spot reclaim" setting, this makes your training run effectively uninterruptible. Longest gap between checkpoints on a spot kill: ckpt_every × step_time = 500 × 1.3s ≈ 11 minutes of lost compute. Acceptable.

8 · The complete file layout

After a completed 100M run, my /workspace/checkpoints/ looks like:

config.json is just json.dump(vars(cfg), open(...)) — human-readable snapshot so you don't have to torch.load a huge file just to check what LR you used.

9 · The "release" checkpoint — what to put on HuggingFace

When you're actually done, produce a single clean file for downstream users:

release = {
    "model": {k: v.to(torch.bfloat16) for k, v in model.state_dict().items()},
    "config": vars(cfg),
    "tokenizer_files": [...],  # or bundle in a separate file
    "training_metadata": {
        "total_tokens": cfg.total_tokens,
        "final_val_loss": 3.82,
        "hellaswag_10shot": 0.302,
        "trained_on": ["custom-cc-2024-30", "the-stack-v2", "wikipedia", "arxiv"],
        "training_hours_a100_40gb": 2.7,
    },
}
torch.save(release, "final-release.pt")

Include the training metadata as part of the checkpoint. When you (or anyone else) revisits in a year, you'll know exactly what this model is.

10 · War stories

11 · The 2025 checkpointing landscape

Checkpointing quietly became one of the biggest infra research areas of 2024–2025. The frontier labs discovered that at 10k+ GPU scale, checkpoint I/O was eating 5–10% of wall clock, and a hardware failure that took 10 min to detect + 30 min to restart from checkpoint cost more than the entire M09 pipeline budget. Here's what shipped.

11.1 · PyTorch DCP (Distributed Checkpoint) — the new standard

torch.distributed.checkpoint (DCP) stable since PyTorch 2.2. It's the parallel-aware, format-standardized replacement for hand-rolled FSDP checkpointing. Key features:

• Async save: dcp.async_save() kicks off save in background, training continues. Save-time cost drops from 30s to ~2s per checkpoint at 10B scale.
• Resharding: save on 8 GPUs, load on 16 or 4 — DCP handles it. This is a huge deal for elastic clusters.
• Storage backend abstraction: S3, GCS, local, all pluggable via dcp.FileSystemWriter.
• Works with FSDP2 and DTensor natively.

Typical usage: dcp.save({"model": model, "optim": optim}, storage_writer=dcp.FileSystemWriter(path)). Reference implementation in torchtitan.

11.2 · CheckFreq / Gemini — the "checkpoint every step" ambition

Microsoft's CheckFreq (Mohan et al., FAST 2021, arxiv 2011.14589) argued that with proper pipelining and delta-encoding, you can checkpoint every ~30 seconds with <3% overhead — vs the naive every-500-steps at 5–10% overhead. Gemini (Wang et al., SOSP 2023) pushed further: checkpoint every step by streaming deltas to remote memory (other GPUs' spare HBM), achieving <1% overhead even at 175B scale.

Why do you care? Because if a spot preemption gives you 60 seconds warning, and your checkpoint takes 2 minutes to write, you lose data. Modern checkpointing is designed around the reality that hardware fails constantly.

11.3 · The "model card" and full artifact standard

When you release your model, the community expectation as of 2025 is:

• model.safetensors (or sharded model-00001-of-N.safetensors + model.safetensors.index.json for >5GB models)
• config.json (HF-format architecture config, so AutoModel.from_pretrained Just Works)
• tokenizer.json (HF tokenizers format)
• tokenizer_config.json, special_tokens_map.json
• generation_config.json (temperature, top_p defaults)
• README.md — the model card. HuggingFace's spec: license, training data summary, evaluation results (with decontamination report per S054), intended use, known limitations, environmental cost (report CO2 per Cohere/HF joint template).
• Optional: training.log excerpts, wandb link.

Best reference: any recent Llama, Mistral, or Qwen release on the HF Hub. Copy their layout.

11.4 · The forgotten optim-state-only checkpoint

Here's a trick from the trenches. During pretrain, save two checkpoint variants:

1. Full (model + optim + sched + rng): every ~2500 steps. For resume. Big — ~4× model size.
2. Weights-only safetensors: every ~500 steps. For eval, sampling, sharing. Small — ~1× model size.

Why: weights-only checkpoints are useful all the time (kick off an eval on a spare box, share a snapshot with a collaborator), and they're cheap. Full checkpoints are only useful for resume, which is rare. Splitting them means your S3 bill drops 3× and your "give me the latest weights" workflow is trivial.

11.5 · Cloud vs local storage — the actual numbers

Practical pattern for a 100M run: active ckpts on the Network Volume, milestone ckpts rsync'd to R2 or HF Hub. Total storage cost for the whole M09 pipeline: under $2/month.

Further reading (11.x)

• PyTorch DCP docs — pytorch.org/docs/stable/distributed.checkpoint.html
• safetensors format — github.com/huggingface/safetensors
• CheckFreq paper — arxiv:2011.14589
• Gemini (SOSP 2023) — dl.acm.org/doi/10.1145/3600006.3613145
• HuggingFace model card guide — huggingface.co/docs/hub/model-cards
• torchtitan checkpoint reference — github.com/pytorch/torchtitan/blob/main/torchtitan/checkpoint.py

12 · 🧠 Retention scaffold

One-line summary (write it in your own words): _______________________________

Spaced review: re-read §1 (what's in a checkpoint) and §11.3 (model card) in 24 hours. Revisit the full session on day 7.

Next session (S056): the flip side — eight real pretraining bugs and their fingerprints on your wandb dashboard.

Sticky note (keep on your desk): "Weights + optim + sched + RNG. Safetensors for share, .pt for resume. Async DCP or you're wasting time."

Legacy recall drills (kept for spaced-review continuity)

Stretch

Extend save_checkpoint to also save wandb run_id, so wandb.init(resume="allow", id=...) continues logging to the same wandb run on resume. Bonus: verify by killing your training after 500 steps, restarting, and confirming the wandb dashboard shows one continuous curve, not two.

In your own words

"A resumed training run whose loss curve matches the original means ______ and ______ were correctly saved. If it drifts, one of those must be broken."

Spaced-review pointer

• S051 (DDP/FSDP) — sharded checkpoints only matter once you shard.
• S052 (renting GPU) — checkpoints saved to /workspace (persistent volume), NOT to ephemeral disk.

Next-session teaser

You now have a bulletproof training pipeline. Session 056 is the flip side: eight real pretraining bugs and their loss-curve fingerprints. Learn what a "bad init" looks like on wandb. What a "tokenizer/data mismatch" curve looks like. What a "dataloader stall" looks like. Pattern-recognize your way to faster debugging.

Bring back tomorrow

• Checklist: 4 things to save.
• Idiom: atomic write via torch.save(tmp) → os.rename(tmp, final).
• Command: the train_forever.sh while-loop.

Previous: ← DL S054 · Next: DL S056 →