Checkpointing and resuming: consistent state and durable publication

Checkpoint the transition, not just the weights

Goal: define restart state, publish a durable local generation, verify continuation, bound retention and explain distributed/remote variants. A checkpoint is a snapshot at a consistent optimizer boundary. It cannot guarantee surviving every failure, nor does a plausible file size prove recoverability.

Think of training as a state transition S(k+1)=U(S(k),batch(k)). To resume the same transition, S must include every input to U: parameters and buffers, optimizer moments and counters, schedule/scaler state, next update/token clocks, random states and data position. Configuration, tokenizer and immutable split/shard hashes identify what those values mean. Validation selection, best metric, EMA, early stopping and tracker run identity also matter when used.

Saving only weights is a warm start, not faithful continuation. An ordinary SGD optimizer without momentum may have little state; Adam stores moments. Their size depends on parameter/state dtype and implementation, not a universal four-times-weight rule. Autocast does not itself convert FP32 model weights to BF16, and torch.compile does not imply BF16 storage.

What must be in the snapshot

Python random, NumPy's legacy global RNG, each independent Generator, torch CPU and each CUDA device are separate sources. np.random.get_state() does not capture np.random.default_rng() used in a loader. Workers, prefetch queues, distributed sampler epochs and intra-epoch positions can require more than a seed. At a deterministic index-addressed boundary, store enough to reconstruct the next sample; otherwise document the permitted statistical rather than bitwise continuation.

A mid-accumulation checkpoint would additionally need partial gradients, their denominator and the microbatch cursor. Prefer committing after a successful optimizer update and zero_grad policy is understood. A process killed halfway through an optimizer update resumes the previous completed checkpoint, not a half-mutated live model.

Atomicity, durability and validation are different

1. Serialize to a unique temporary file in the destination directory, from a stable snapshot.
2. Flush Python buffers and fsync the file.
3. os.replace the temporary name over the final name on the same filesystem.
4. fsync the directory so the name update is persisted under the filesystem's guarantees.
5. Validate the completed generation before pruning older validated ones.

Rename gives atomic namespace replacement on the supported local filesystem; it does not magically make earlier data writes durable, work across filesystems, or create an atomic transaction across multiple object-store uploads. A failure after replacement but before directory sync is a durability uncertainty. Power loss, storage/controller faults and remote-filesystem behavior need their own fault tests.

A latest symlink updated by unlink-then-create has a missing-pointer window. Prefer replacing a small manifest/symlink atomically after a generation is validated. Keep the actual generation immutable and give the manifest its digest, size, format and next-update counter. A hash detects corruption against a trusted expected digest; a digest bundled by an adversary does not authenticate an untrusted checkpoint.

Complete local persistence utilities

Save as checkpointing.py. This is a synchronous trusted-local format using torch2.9 restricted loading; it is not an untrusted-file parser or a cloud/distributed checkpoint implementation. The failure injection is a local exception before replacement, not a simulated power failure.

import os, random, tempfile, hashlib, re
from pathlib import Path
import numpy as np
import torch
 
def capture_rng():
    a=np.random.get_state()
    return {'python':random.getstate(),
            'numpy_global':(a[0],a[1].tolist(),a[2],a[3],a[4]),
            'torch_cpu':torch.get_rng_state(),
            'torch_cuda':torch.cuda.get_rng_state_all() if torch.cuda.is_available() else []}
 
def restore_rng(s):
    random.setstate(s['python'])
    a=s['numpy_global']; np.random.set_state((a[0],np.array(a[1],dtype=np.uint32),a[2],a[3],a[4]))
    torch.set_rng_state(s['torch_cpu'])
    if s['torch_cuda']:
        if len(s['torch_cuda'])!=torch.cuda.device_count(): raise ValueError('CUDA topology changed')
        torch.cuda.set_rng_state_all(s['torch_cuda'])
 
def atomic_save(payload,path,fail_before_replace=False):
    path=Path(path); path.parent.mkdir(parents=True,exist_ok=True)
    fd,tmp=tempfile.mkstemp(prefix=path.name+'.',suffix='.tmp',dir=path.parent)
    try:
        with os.fdopen(fd,'wb') as f:
            torch.save(payload,f); f.flush(); os.fsync(f.fileno())
        if fail_before_replace: raise RuntimeError('injected before publication')
        os.replace(tmp,path)
        dfd=os.open(path.parent,os.O_RDONLY|os.O_DIRECTORY)
        try: os.fsync(dfd)
        finally: os.close(dfd)
    finally:
        if os.path.exists(tmp): os.unlink(tmp)
    return hashlib.sha256(path.read_bytes()).hexdigest()
 
def trusted_load(path,expected_sha256=None):
    path=Path(path)
    if expected_sha256 is not None and hashlib.sha256(path.read_bytes()).hexdigest()!=expected_sha256:
        raise ValueError('checkpoint digest mismatch')
    payload=torch.load(path,map_location='cpu',weights_only=True)
    if not isinstance(payload,dict) or payload.get('format')!=1:
        raise ValueError('checkpoint format')
    return payload
 
def latest_valid(paths,required=('model','optimizer','next_update','rng')):
    failures=[]
    # Caller supplies newest-first committed generations, never glob order.
    for p in paths:
        try:
            state=trusted_load(p)
            if not all(k in state for k in required): raise ValueError('missing restart state')
            if not isinstance(state['next_update'],int) or state['next_update']<0:
                raise ValueError('invalid update')
            return p,state,failures
        except Exception as e: failures.append({'path':str(p),'error':type(e).__name__})
    raise RuntimeError('no validated restart generation: '+repr(failures))
 
def retention_plan(paths,keep=3,milestone_every=2500,pins=()):
    if keep<1 or milestone_every<1: raise ValueError('retention config')
    steps=[]
    for p in map(Path,paths):
        m=re.fullmatch(r'step-(\d+)\.pt',p.name)
        if not m: raise ValueError('unexpected checkpoint filename')
        steps.append((int(m[1]),p))
    steps.sort(key=lambda x:x[0])
    retained={p for _,p in steps[-keep:]}|set(map(Path,pins))
    retained|={p for step,p in steps if step%milestone_every==0}
    return {'keep':[str(p) for _,p in steps if p in retained],
            'candidates_for_pruning':[str(p) for _,p in steps if p not in retained]}
 
if __name__=='__main__':
    import copy
    torch.set_num_threads(1); torch.manual_seed(9); np.random.seed(9); random.seed(9)
    model=torch.nn.Sequential(torch.nn.Linear(3,5),torch.nn.Dropout(.25),torch.nn.Linear(5,2))
    opt=torch.optim.AdamW(model.parameters(),lr=.01)
    def step():
        x=torch.randn(4,3); y=torch.randn(4,2)
        opt.zero_grad(set_to_none=True); loss=(model(x)-y).square().mean()
        loss.backward(); opt.step(); return float(loss.detach()),x
    for _ in range(5): step()
    payload={'format':1,'model':copy.deepcopy(model.state_dict()),
             'optimizer':copy.deepcopy(opt.state_dict()),'rng':capture_rng(),'next_update':5}
    with tempfile.TemporaryDirectory() as directory:
        p=Path(directory)/'step-5.pt'; digest=atomic_save(payload,p)
        expected_loss,expected_x=step(); expected=copy.deepcopy(model.state_dict())
        s=trusted_load(p,digest); model.load_state_dict(s['model']); opt.load_state_dict(s['optimizer'])
        restore_rng(s['rng']); loss,x=step()
        assert loss==expected_loss and torch.equal(x,expected_x)
        assert all(torch.equal(expected[k],v) for k,v in model.state_dict().items())
        model.load_state_dict(s['model']); opt=torch.optim.AdamW(model.parameters(),lr=.01)
        restore_rng(s['rng']); step()
        assert any(not torch.equal(expected[k],v) for k,v in model.state_dict().items())
        before=p.read_bytes()
        try: atomic_save({'format':1,'broken':True},p,fail_before_replace=True)
        except RuntimeError: pass
        else: raise AssertionError('failure injection did not fire')
        assert p.read_bytes()==before
        bad=Path(directory)/'step-6.pt'; bad.write_bytes(b'truncated')
        chosen,_,failures=latest_valid([bad,p]); assert chosen==p and len(failures)==1
        try: trusted_load(p,'0'*64)
        except ValueError: pass
        else: raise AssertionError('wrong checksum accepted')
    plan=retention_plan(['step-9.pt','step-10.pt','step-100.pt','step-2500.pt'],keep=2)
    assert plan['keep']==['step-100.pt','step-2500.pt']
    print('disk resume exact; missing-optimizer, corrupt-latest, hash and atomic-publication controls passed')

latest_valid performs structural screening, not a complete restart rehearsal: it checks the format, required keys and update counter. It does not verify tensor shapes, optimizer parameter pairing, RNG schema, experiment contract or a trusted per-generation digest. An empty model dictionary can pass screening and still fail load_state_dict. Before treating a generation as recoverable or pruning older ones, load it into freshly constructed matching objects, validate the data/config contract, restore every used state component and check the next batch/update. The caller can use trusted_load(path, expected_sha256) with a trusted manifest digest; the fallback helper does not supply that digest automatically. Keep trusted generations immutable during this check; these helpers are not a concurrent-writer transaction or an untrusted-file security boundary.

A synchronous save blocks updates while serialization reads state_dict; it therefore sees one consistent state. state_dict() contains live tensor references, so blindly handing it to a background writer while training continues can mix update generations. Async saving needs an owned copy/staging boundary, completion/failure handling and a limit on in-flight writes. Moving to CPU changes memory pressure; it is not free.

Retention, disk sizing and cadence derivation

Sort checkpoints by parsed numeric update or zero-pad names. Lexical order places step100 before step9. The planner above never deletes user files; pruning is a separate operation performed only after validating the newest generation and preserving pinned best/final snapshots. Last-K plus every-N milestones grows with total run duration; it is not bounded forever. For a finite run with M milestone generations, an upper bound is (K+M)*S, minus overlaps, plus temporary/staged/in-flight copies and logs. Disk-full before replacement should leave the old generation readable; retain headroom for writing the next one.

Let a completed save cost C seconds, checkpoint interval T seconds, and mean failure interval M seconds. For independent rare failures and roughly uniform failure phase, overhead fraction is approximately

Differentiate: -C/T²+1/(2M)=0, so T*=sqrt(2CM). With C=10s and M=14400s (four hours), T≈536.7s, about8.94minutes. This is an illustrative stochastic model, not a recovery SLA. Restore time, correlated faults, overlap, staging and failure warning change the optimum. Convert seconds to updates using measured step time; “every500” has no universal meaning.

Saving weights every500 but full state every2500 allows frequent evaluation, yet restart loses up to2500 updates since the latest full state. Do not claim500-step restart protection from weight-only files. Remote durability can lag local cadence further. Record the last completed remote generation, not merely an enqueued upload.

Resume drift: compare in causal order

First compare contract/hash, then loaded weights/buffers, next batch/mask, RNG draw, used LR, loss, gradients and updated weights. Seed fresh objects before restoring RNG, and restore RNG only after constructing objects whose initialization consumes randomness. Loading an optimizer or scheduler in a different order can alter group LR; test the actual library version rather than copying last_epoch blindly.

Exact equality is achievable in this bounded CPU environment and is the regression criterion here. It is not guaranteed across devices, kernels, PyTorch versions or world sizes. A failed equality test can reveal missing state or nondeterminism; a matching scalar loss can conceal different state. Compare tensors and batch identity, not a folklore0.001/0.05 threshold that supposedly names the missing component.

The training recipe stores next_update, tokens_seen, the sampler Generator state, a deterministic LR function's config, model/optimizer and global RNG. It rejects a changed contract. It restores the same fixed validation sampler after each evaluation, so monitoring does not consume training samples. Best-checkpoint selection or EMA would need explicit extra state if added.

Sharded checkpoints and asynchronous recovery

FSDP partitions parameter/state storage; gathering a full export can put the entire state on one rank's host memory. Whether that fits depends on dtype and RAM, not a10B/30B parameter threshold. A sharded checkpoint has multiple tensor shards plus metadata and requires participating ranks and a supported resharding protocol.

PyTorch2.9 DCP supports distributed save/load and load-time resharding. Plain model.state_dict() and optimizer IDs are not automatically a correct FSDP restart contract; use the documented distributed state-dict APIs for model/optimizer pairing. Loading writes into allocated state before applying it back to the model/optimizer. Full optimizer loading matters as much as model loading. DCP explicitly warns against assuming checkpoint compatibility across PyTorch versions.

async_save stages state and returns a completion object; wait/check errors and finish outstanding saves before shutdown. FileSystemWriter is a filesystem writer, not automatically an S3/GCS adapter. Backend-specific storage readers/writers and manifests supply remote behavior. CheckFreq studies runtime-adaptive checkpoint frequency/pipelining; Gemini studies checkpointing to host CPU memory with redundancy/placement and traffic scheduling for faster recovery, not streaming deltas into spare GPU HBM. Their experimental overheads do not prove “every step below1%” on your system, and volatile peer memory still needs a failure-domain plan.

No DCP/distributed/cloud run was executed. A supported production integration must test missing shard, changed world size, slow rank, partial write, backend failure and recovery into the intended model/optimizer topology.

Local, remote and release artifacts

Local NVMe can recover a process crash yet vanish with the instance; an attached persistent/network volume has a different lifecycle. On Runpod, stop, terminate, container disk, volume disk and network-volume semantics are distinct and must be read for the actual pod type. “Always terminate” is unsafe advice without verifying that required artifacts are durable. This lesson performs no provider actions and quotes no current storage price.

For each candidate storage tier—local, network volume, S3 Standard/intelligent tiering/archive, R2 or the Hub—estimate stored GB-months, write/read requests, transfer/egress, restore latency, minimum retention, quota and access controls. A cold archive can be inappropriate for a short recovery target; public model hosting is not a private training-state backup. Multiply actual bytes by dated rates, not a timeless under-USD2/month table.

Release weights are a separate product from restart state. Publish only after authorization: architecture/code version, tokenizer/template/special tokens, generation config, weight shards/index, model card, license/provenance, measured evaluation protocol/results, intended uses, limitations and measured compute/environmental information. A custom decoder .pt is not automatically loadable by AutoModelForCausalLM; compatible configuration and architecture implementation are required.

Safetensors stores tensors without pickle execution. Tied/shared tensors need its documented model-saving/sharing treatment rather than assuming any state_dict can be passed to save_file. Non-tensor optimizer metadata can be represented explicitly alongside tensors; pickle is not inherently required. Torch2.6+ defaults to restricted weights_only loading in the usual call path; restricted loading reduces code-execution risk but does not protect from all malformed-file/resource risks. Load only trusted files within this recipe. Do not cast integer buffers to BF16 when making an inference export. Removing optimizer/RNG does not itself establish privacy, license or safe-to-publish status.

Bounded restart policy, not an infinite launcher

A shell loop on a dead instance cannot provision its replacement. On a live host, distinguish intentional completion/cancellation from recoverable failure and permanent configuration/disk errors. Use a maximum retry count, backoff, operator stop signal, disk-space preflight and validated checkpoint fallback. Never delete latest just to force a fresh run; create a separate run directory and preserve the evidence. A DONE record must bind config, token count and final digest, not a stale empty file from yesterday.

A finite reference policy is: attempt at most3 starts; on zero exit verify DONE and checkpoint identity; on a permanent validation/config error stop; on an allowed transient failure log it, wait10/20/40seconds and retry from the newest validated full state. This control policy is specified, not executed against a cloud service.

Worked recall and stretch answers

1. Four state components? Weights/buffers, optimizer, schedule/scaler/counters and RNG are the mnemonic; data position and experiment contract are also essential. Include only state actually used, but account for every independent generator.
2. Stop versus terminate? Provider and storage type decide preservation and billing. Verify remote artifact readability before any destructive lifecycle action; no current fee or automatic termination is prescribed.
3. Why safetensors? Tensor-only serialization avoids pickle's arbitrary-object execution path. It is not an all-purpose training-state container or blanket guarantee against resource exhaustion/privacy problems.
4. DCP? Distributed state persistence, asynchronous save and supported resharding. No fixed model-size threshold, speedup or cross-version portability follows.
5. Resume drift? Compare batch, RNG, clocks, moments and weights in order; then consider deterministic-kernel/environment differences. Scalar closeness alone is insufficient.
6. 7B/32-GPU/four-hour-failure exercise: measure save/restore cost and failure-domain rate, compute an initial interval using the formula, retain local and completed remote generations, test all-rank recovery and then propose an SLA. A promised five-minute recovery without measurement is not an answer.

Legacy drills: temp/fsync/replace/directory-sync closes specified local crash windows; rank-zero export and sharded restart serve different uses; rolling last3 plus2500-step milestones requires temporary headroom and finite retention planning. To resume a tracker, persist its run ID and logging step, but keep it optional and avoid conflating remote log continuity with numerical continuation.

Sources and execution boundary

PyTorch2.9 serialization, reproducibility, DCP, Python3.12 filesystem operations, safetensors0.6.2, CheckFreq, Gemini, SOSP2023 author-hosted paper.

Executed: trusted local disk round-trip, same-runtime next-step equality, missing-optimizer negative control, truncated-latest fallback, digest rejection and injected pre-replace failure. Not executed: power-cut durability, GPU/distributed resharding, cloud storage, instance restarts or public release. Optional reading and next steps are listed below.

Inherited optional resources and navigation

These pre-read links are optional further reading; their timestamps, video claims and optional pacing are not accepted technical evidence.

• Save/Load Checkpoints in PyTorch — Complete Guide
• Saving and loading models
• Distributed Checkpoint (torch.distributed.checkpoint)
• nanoGPT checkpoint code

Previous: ← DL S054 · Next: DL S056 →