Checkpointing Strategies

At 10K+ GPUs, hardware failures happen every few hours. Checkpointing is the insurance policy — saving training state so you can resume after failures without losing days of compute.

1. What to Save

A complete training checkpoint includes:

Component	Size (70B, TP=8, DP=32)	Per-rank	Total
Model parameters	140 GB	17.5 GB / TP rank	140 GB
Optimizer states (FP32)	840 GB	Sharded per DP rank	840 GB
Learning rate scheduler	~1 KB	Same on all ranks	~1 KB
Data loader state	~1 MB	Per DP rank	~32 MB
RNG states (CUDA + CPU)	~1 KB	Per rank	~256 KB
Total			~980 GB

2. Synchronous Checkpointing

All ranks save simultaneously, blocking training:

if step % save_interval == 0:
    # All ranks reach this barrier
    dist.barrier()
    
    # Each rank saves its shard
    torch.save({
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'scheduler': scheduler.state_dict(),
        'step': step,
        'rng_state': torch.cuda.get_rng_state(),
    }, f"checkpoint/rank_{rank}_step_{step}.pt")
    
    dist.barrier()  # ensure all ranks finished saving

Time cost: Writing 980 GB to a parallel filesystem (e.g., Lustre at 100 GB/s aggregate) takes ~10 seconds. With 256 GPUs writing their shards in parallel, throughput is good — but training is blocked for those 10 seconds.

3. Asynchronous Checkpointing

Copy state to CPU memory (pinned), then write to disk in background:

if step % save_interval == 0:
    # Stage 1: Copy GPU tensors to CPU (fast, ~1-2 seconds)
    cpu_state = {k: v.cpu() for k, v in model.state_dict().items()}
    cpu_optimizer = {k: v.cpu() for k, v in optimizer.state_dict().items()}
    
    # Stage 2: Resume training immediately
    # (background thread writes cpu_state to disk)
    save_thread = threading.Thread(
        target=save_to_disk, args=(cpu_state, cpu_optimizer, step)
    )
    save_thread.start()

Blocking time: Only the GPU→CPU copy (~1-2 seconds), not the full disk write. Training resumes while disk I/O happens in the background.

Requirement: Enough CPU RAM to hold one full checkpoint copy while training continues.

4. Distributed Checkpointing (PyTorch DCP)

PyTorch’s torch.distributed.checkpoint enables each rank to save its shard independently:

import torch.distributed.checkpoint as dcp

if step % save_interval == 0:
    dcp.save(
        state_dict={"model": model, "optimizer": optimizer},
        storage_writer=dcp.FileSystemWriter(f"checkpoints/step_{step}"),
    )

Benefits:

Each rank saves only its shard — parallel I/O
Resharding: can load a checkpoint saved with TP=8 into a different TP=4 configuration
No single rank needs to hold the full model

5. Optimal Save Interval

Too frequent: Wastes time on checkpointing (10s every 100 steps = 10% overhead). Too infrequent: Loses more work on failure (if MTBF is 4 hours and you save every 2 hours, you lose 1 hour on average).

Optimal interval formula

Given Mean Time Between Failures (MTBF) and checkpoint time $T_c$ :

T_{\text{optimal}} = \sqrt{2 \times \text{MTBF} \times T_c}

Cluster	MTBF	$T_c$ (async)	Optimal interval	Lost work per failure
256 GPUs	~8 hours	2 seconds	~6.7 minutes	~3.3 minutes
1024 GPUs	~2 hours	2 seconds	~3.4 minutes	~1.7 minutes
10K GPUs	~15 minutes	2 seconds	~1.4 minutes	~0.7 minutes

At 10K GPUs, you checkpoint every ~1.5 minutes!

6. Storage Architecture

Storage tier	Technology	Bandwidth	Latency	Cost	Use
Local NVMe	Samsung PM9A3	7 GB/s	~10 μs	$$	Fast first-stage save
Parallel FS	Lustre, GPFS	100+ GB/s	~ms	$$$	Shared checkpoints
Object store	S3, GCS	~50 GB/s	~10 ms	$	Long-term archival

Two-tier strategy: Save to local NVMe first (fast, unblocks training), then asynchronously replicate to parallel FS or object store for durability.

For detecting failures and recovering from them, see Fault Detection & Recovery.

Self-Check

List the components of a training checkpoint and their sizes for a 70B model
Compare synchronous vs asynchronous checkpointing blocking time
Calculate optimal checkpoint interval for a given MTBF and checkpoint time
Implement distributed checkpointing with PyTorch DCP
Design a two-tier storage strategy for checkpoints