0

Design review: checkpointing a large training run

by · about 4 hours ago· ⑂ 0 remixes

Question: Design review: checkpointing a large training run

Companies: Anthropic

Difficulty: Hard

Problem statement

You are reviewing the checkpointing design for a month-long training run. Critique it: what fails, what it costs, what you would change, and what the doc does not cover. You have 45 minutes.

Checkpointing for the 70B run (draft)

Setup. A 70B-parameter model trained on 1,024 GPUs (128 nodes) for 30 days with mixed precision (bf16 compute) and the Adam optimizer, sharded across ranks.

Size. A checkpoint is the model weights: 70B × 2 bytes ≈ 140 GB.

Schedule. Every 2 hours, all ranks stop training and write their shards to one shared NFS server (about 5 GB/s of write bandwidth).

Commit. Rank 0 first updates the latest file to point at the new checkpoint directory, then every rank writes its shard there.

Restart. After a failure the job restarts from latest. The data loader restarts from the beginning of the current epoch (one epoch is about 6 days).

Retention. All checkpoints are kept for the life of the project.

Topology. Restoring requires the same 1,024-GPU layout, because each shard maps to one rank.

Validation. None; NFS writes are reliable.

How to run the review

Checkpointing exists for one reason: when (not if) the job fails, lose as little work as possible and resume exactly where it was. Judge every line of the doc against two questions: does a restart produce the same training run that would have happened without the failure, and what fraction of the cluster's time does checkpointing cost?

  1. Fix the size first. Every cost in the doc follows from it, and it is wrong.
  2. Model failures. At 1,024 GPUs, interruptions (a GPU, a node, a switch, a host process) are routine. Plan for about one a day and measure your own rate; the math below uses that.
  3. Walk a failure at every point: mid-training, mid-write, mid-commit, on restore with a missing node.
  4. Rank by what corrupts or silently changes the run first, then by GPU-hours wasted.
  5. Credit what is good: sharded writes by rank and a single latest pointer are the right building blocks.

The numbers

Size. Mixed-precision Adam training keeps more than the bf16 weights: fp32 master weights (4 bytes), Adam's two moments (4 + 4 bytes) and the bf16 copy (2 bytes), 14 bytes per parameter. 70B × 14 bytes ≈ 980 GB, about 7 times the doc's 140 GB. Saving only the bf16 weights would not let training resume at all: the optimizer state would be lost.

Write time. 980 GB at 5 GB/s on one NFS server is 196 seconds, with all 1,024 GPUs idle.

Cost of the doc's plan. With a checkpoint every 2 hours and about one failure a day:

  • Stall: 196 s every 7,200 s = 2.7% of the run.
  • Lost work: on average half an interval (1 hour) per failure, so 1 hour a day = 4.2%. That is 1,024 GPU-hours thrown away per failure.
  • Total ≈ 6.9% of 737,280 GPU-hours for the month, about 51,000 GPU-hours, before counting restart time.

The best interval for a given checkpoint cost C and mean time between failures M is about √(2·C·M) (Young's approximation). With the doc's 196 s stall, that is about 97 minutes: the 2-hour interval is already near its best, so the interval is not the problem. The cost per checkpoint is. If the training pause drops to about 5 seconds (copy each rank's shard to host memory and write it in the background), the best interval becomes √(2 × 5 × 86,400) ≈ 15 minutes and the total overhead drops to about 1.1%, saving roughly 43,000 GPU-hours over the month.

Retention. Keeping all 360 checkpoints (12 a day for 30 days) is about 350 TB. Keeping the last 3 plus one a day is about 32 TB.

Planted flaws

# Flaw Why it matters Severity Fix
1 latest is updated before the shards are written A failure during the write (the most likely moment, since writing takes minutes and stresses storage) leaves latest pointing at a partial checkpoint. The restart crashes, or worse, loads a mix of old and missing shards Critical Write all shards to a new directory, write a manifest with every shard's size and checksum, verify that every rank finished, then update the pointer last with an atomic rename
2 Size counts only bf16 weights 980 GB, not 140 GB. Without optimizer state the run cannot resume; with it, every time and storage estimate in the doc is 7× too small Critical Save weights, master weights, optimizer moments, step counter, learning-rate scheduler, gradient scaler, and RNG states
3 Data loader restarts at the start of the epoch An epoch is 6 days. After a failure on day 5 the model sees 5 days of data again: about one failure a day means the early part of each epoch is trained on many times and the end barely at all. The run is no longer the run you designed Critical Save the data loader position per rank (shard, offset, shuffle seed) and resume exactly; test that a resumed run sees the same next batch as an uninterrupted one
4 All 1,024 GPUs stall for 196 s on one NFS server 2.7% of the run idle, and one NFS server is a single point of failure for the whole job High Two stages: snapshot to host memory (seconds), then write in the background to local NVMe (about 4 s per node at 2 GB/s for its 7.7 GB) and upload to object storage with aggregate bandwidth across 128 nodes
5 No validation A truncated file, a bit flip, or NaNs in the weights surface only when you need the checkpoint, after the failure High Checksums in the manifest, checked on load; a NaN and norm check before commit; regularly restore a checkpoint on a separate small job and compare loss on a fixed batch
6 Restore needs the same 1,024-GPU layout If a node dies and there is no spare, the job cannot restart until hardware is repaired High Hot spare nodes, and a checkpoint format that stores tensors with their global shapes so it can be resharded to a different layout
7 Every checkpoint kept forever About 350 TB for one month Medium Keep the last 3, one a day, and named milestones; delete the rest after they are verified superseded
8 Nothing about monitoring Nobody knows when a checkpoint fails, how long writes take, or how much work was lost Medium Alerts on a missed or failed checkpoint, write duration, time since the last good checkpoint

Revised design

{"type":"excalidraw","elements":[
 {"type":"rectangle","id":"ranks","x":0,"y":0,"width":180,"height":70,"label":{"text":"1,024 ranks"}},
 {"type":"rectangle","id":"snap","x":260,"y":0,"width":200,"height":70,"label":{"text":"Host-memory\nsnapshot"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"ellipse","id":"nvme","x":540,"y":-5,"width":200,"height":80,"label":{"text":"Local NVMe"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"ellipse","id":"obj","x":820,"y":-5,"width":200,"height":80,"label":{"text":"Object store"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"rectangle","id":"commit","x":1100,"y":0,"width":220,"height":70,"label":{"text":"Manifest, then\npointer (last)"},"backgroundColor":"#d0ebff","fillStyle":"solid"},
 {"type":"arrow","id":"a1","x":180,"y":35,"width":80,"height":0,"start":{"id":"ranks"},"end":{"id":"snap"}},
 {"type":"arrow","id":"a2","x":460,"y":35,"width":80,"height":0,"start":{"id":"snap"},"end":{"id":"nvme"}},
 {"type":"arrow","id":"a3","x":740,"y":35,"width":80,"height":0,"start":{"id":"nvme"},"end":{"id":"obj"}},
 {"type":"arrow","id":"a4","x":1020,"y":35,"width":80,"height":0,"start":{"id":"obj"},"end":{"id":"commit"}},
 {"type":"text","id":"legend","x":0,"y":130,"text":"Shaded: what the review adds. Training pauses only for the snapshot (seconds).","fontSize":16}
]}
  1. Every 15 minutes, each rank copies its full state (shard of weights, master weights, optimizer moments, data loader position, RNG) to pinned host memory: a few seconds of pause.
  2. A background thread writes the snapshot to local NVMe and then uploads it to object storage, while training continues. If a new snapshot is due before the upload finishes, skip it rather than queue it.
  3. When every rank has uploaded, a coordinator writes the manifest (shard list, sizes, checksums, step, data position, layout) and only then moves the latest pointer. A checkpoint without a manifest does not exist.
  4. Restore reads the manifest, verifies checksums, and reshards if the layout changed. A small job periodically restores the latest checkpoint and checks the loss on a fixed batch.
  5. Retention keeps the last 3, one a day, and milestones.

Questions to ask the author

  • What interruption rate have we measured on this cluster, and how long does a restart take end to end (scheduling, loading, warm-up)?
  • What does the parallelism look like (tensor, pipeline, data parallel), and does anything outside the ranks need saving (data loader workers, a separate parameter server)?
  • Do we need to evaluate or fine-tune from intermediate checkpoints? That decides which milestones to keep.
  • Is there a spare-node pool, and who replaces failed nodes?

Interviewer expectations

Mid-level (L4)

  • Notices the pointer is updated before the data is written and proposes writing the pointer last.
  • Points out that optimizer state is missing and that the data loader should resume where it stopped.
  • Suggests validating checkpoints and limiting retention.

Senior (L5)

  • Recomputes the size (about 1 TB) and the stall, and quantifies the cost of the plan in GPU-hours.
  • Proposes asynchronous two-stage checkpointing with a manifest-based atomic commit and exact data loader resume.
  • Raises the topology lock-in and proposes spares or a reshardable format.

Staff (L6+)

  • Uses the interval formula to show that the interval is near optimal for the doc's cost, so the real lever is a cheaper checkpoint, and estimates the savings (about 43,000 GPU-hours).
  • Treats correctness of the resumed run as the core requirement: deterministic resume, a test that proves it, and periodic restore drills.
  • Covers operations: monitoring, alerting, storage cost, and how the design changes at 10× the scale.

Sign in to comment, vote or remix.