Why is model checkpointing done asynchronously?

Because a synchronous save blocks the training step for as long as it takes to write the whole model state to storage, and a large checkpoint is terabytes. Accelerators — the most expensive hardware in the cluster — sit idle the entire time. An async checkpointer splits the save in two: a fast blocking device-to-host copy that snapshots the arrays into host memory, after which the training step resumes, and a background thread that does the slow storage write while training continues. The device stall shrinks from the full write time to just the copy.

What does it mean for a checkpoint to be committed atomically?

A crash halfway through writing a checkpoint must not leave a half-written directory that looks complete — a later restore would load garbage. So the write goes to a temporary location and is only made visible as one final, indivisible step. On a POSIX filesystem that step is an atomic directory rename. On an object store like GCS or S3, where there is no atomic rename of a directory, completion is instead signalled by writing a commit-success marker file after every byte is durable; a reader treats the checkpoint as valid only if that marker exists. Either way, a checkpoint is either entirely there or not there at all.

How does a sharded checkpoint decide which host writes which bytes?

Each array is already sharded across the devices of the cluster, so every host writes only the shards its own devices hold — no host holds or writes the whole array. When an array is replicated across several hosts, only one replica is chosen to write those bytes, so the same data is not written many times over. The result is that total write traffic scales with the size of the model, not with the size of the model times the number of hosts.

Can a checkpoint be restored onto a different device topology than it was saved on?

Yes, and this is the point of storing the checkpoint as sharded arrays rather than one host’s view of them. Restore is driven by the target sharding: each host reads exactly the byte ranges its new shards need, so a checkpoint saved on one mesh can be loaded onto a mesh with a different shape or a different number of hosts. Resharding happens at restore time, which decouples how the model was saved from how it is loaded — you can resume a run on a bigger or smaller cluster.

How are old checkpoints garbage-collected?

A checkpoint manager keeps a preservation policy — most commonly keep the latest N steps, but also keep one every N steps, keep the best by a metric, or keep specific steps — and deletes the rest in the background as new checkpoints land. Partial or temporary checkpoints left behind by a crashed save are also cleaned up, because they never received their commit marker and so are recognisable as incomplete. Garbage collection is what keeps a long run from filling storage with every step it ever saved.

The Checkpoint Lifecycle: What Happens Between save() and Durable

A training run that saves synchronously leaves a rack of accelerators idle for the seconds it takes to push a terabyte to storage. So the save goes async — a fast device-to-host copy, then a background thread that writes, coordinates every host, commits the directory as one atomic unit, and garbage-collects the old ones. Follow one save from the training step that triggers it to the moment it is durable and restorable onto a different mesh — drawn, computed, and animated.

System design · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Checkpoint Lifecycle: What Happens Between save() and Durable — built to be understood, not skimmed.

Questions

Why is model checkpointing done asynchronously?: Because a synchronous save blocks the training step for as long as it takes to write the whole model state to storage, and a large checkpoint is terabytes. Accelerators — the most expensive hardware in the cluster — sit idle the entire time. An async checkpointer splits the save in two: a fast blocking device-to-host copy that snapshots the arrays into host memory, after which the training step resumes, and a background thread that does the slow storage write while training continues. The device stall shrinks from the full write time to just the copy.
What does it mean for a checkpoint to be committed atomically?: A crash halfway through writing a checkpoint must not leave a half-written directory that looks complete — a later restore would load garbage. So the write goes to a temporary location and is only made visible as one final, indivisible step. On a POSIX filesystem that step is an atomic directory rename. On an object store like GCS or S3, where there is no atomic rename of a directory, completion is instead signalled by writing a commit-success marker file after every byte is durable; a reader treats the checkpoint as valid only if that marker exists. Either way, a checkpoint is either entirely there or not there at all.
How does a sharded checkpoint decide which host writes which bytes?: Each array is already sharded across the devices of the cluster, so every host writes only the shards its own devices hold — no host holds or writes the whole array. When an array is replicated across several hosts, only one replica is chosen to write those bytes, so the same data is not written many times over. The result is that total write traffic scales with the size of the model, not with the size of the model times the number of hosts.
Can a checkpoint be restored onto a different device topology than it was saved on?: Yes, and this is the point of storing the checkpoint as sharded arrays rather than one host’s view of them. Restore is driven by the target sharding: each host reads exactly the byte ranges its new shards need, so a checkpoint saved on one mesh can be loaded onto a mesh with a different shape or a different number of hosts. Resharding happens at restore time, which decouples how the model was saved from how it is loaded — you can resume a run on a bigger or smaller cluster.
How are old checkpoints garbage-collected?: A checkpoint manager keeps a preservation policy — most commonly keep the latest N steps, but also keep one every N steps, keep the best by a metric, or keep specific steps — and deletes the rest in the background as new checkpoints land. Partial or temporary checkpoints left behind by a crashed save are also cleaned up, because they never received their commit marker and so are recognisable as incomplete. Garbage collection is what keeps a long run from filling storage with every step it ever saved.

The Checkpoint Lifecycle: What Happens Between save() and Durable

Questions

Related explainers