Checkpoint Storage Design

Checkpoint writes are the most intense I/O event in training — all GPUs dump their state simultaneously. Without proper design, checkpoint I/O can block training for minutes.

---

1. The Checkpoint I/O Storm

When 256 GPUs checkpoint simultaneously:

Component	Size per GPU	Total (256 GPUs)	Write time @ 10 GB/s
Model params (sharded)	0.5-2 GB	128-512 GB	13-51s
Optimizer states (sharded)	2-8 GB	512-2048 GB	51-205s
RNG + metadata	~1 MB	~256 MB	<1s
Total	2.5-10 GB	640-2560 GB	64-256s

2-4 minutes of training blocked just for checkpointing on a parallel FS. At a checkpoint interval of 10 minutes, that’s 20-40% overhead!

---

2. Two-Stage Checkpoint (NVMe Staging)

Stage 1: GPU → Local NVMe (fast, ~1-2 seconds)
  Each GPU writes its shard to local NVMe SSD
  Training resumes IMMEDIATELY

Stage 2: NVMe → Shared Storage (background, async)
  Background thread copies from NVMe to Lustre/S3
  Takes 1-5 minutes but doesn't block training

Training blocked time: 1-2 seconds (GPU → NVMe only). The 1-5 minute upload happens while training continues.

Implementation

import threading

def async_checkpoint(model, optimizer, step, local_path, remote_path):
    # Stage 1: fast local save (blocks training ~1-2s)
    local_file = f"{local_path}/ckpt_step_{step}.pt"
    torch.save({
        'model': model.state_dict(),
        'optimizer': optimizer.state_dict(),
        'step': step,
    }, local_file)
    
    # Stage 2: async upload (doesn't block training)
    def upload():
        shutil.copy(local_file, f"{remote_path}/ckpt_step_{step}.pt")
        # Or: boto3 upload to S3
    
    thread = threading.Thread(target=upload)
    thread.start()
    return thread  # caller can join() if needed

---

3. Retention Policies

Not all checkpoints need to be kept forever:

Tier	Retention	Storage	Cost
Hot (last 3 checkpoints)	Always available	Local NVMe	Included
Warm (last 24 hours)	Quick restore	Parallel FS / S3 Standard	$25/TB/mo
Cold (weekly snapshots)	Disaster recovery	S3 Glacier	$4/TB/mo
Archived (milestone models)	Permanent	S3 Deep Archive	$1/TB/mo

# Lifecycle policy
retention = {
    "keep_last_n": 3,           # always keep 3 most recent on NVMe
    "keep_hourly_for": "24h",   # hourly checkpoints for 24 hours
    "keep_daily_for": "7d",     # daily checkpoints for 1 week
    "archive_milestones": True,  # keep 25%, 50%, 75%, 100% checkpoints forever
}

---

4. Distributed Checkpointing

With FSDP/ZeRO-3, each rank saves only its shard — enabling parallel I/O:

Rank 0: save shard_0.pt (2 GB) → NVMe 0
Rank 1: save shard_1.pt (2 GB) → NVMe 1
...
Rank 255: save shard_255.pt (2 GB) → NVMe 255

Total: 512 GB written in parallel across 256 NVMe drives
Time: ~0.3 seconds (512 GB / 256 drives / 7 GB/s per drive)

Resharding: PyTorch DCP (torch.distributed.checkpoint) supports loading checkpoints saved with different TP/PP/DP configurations — essential for changing parallelism after a checkpoint.

For the training-side checkpoint strategies, see Checkpointing Strategies.

---

Self-Check

• Calculate checkpoint I/O time for a given model and GPU count
• Design a two-stage checkpoint (NVMe staging + async upload)
• Create a retention policy balancing cost and recovery speed
• Explain distributed checkpointing with FSDP shard-per-rank

---

References

• PyTorch Distributed Checkpoint
• DeepSpeed Checkpointing