Checkpointing and Recovery › Asynchronous Checkpointing

The key insight behind asynchronous checkpointing is that writing a checkpoint to disk does not require the GPU to be involved. Once the checkpoint data has been copied from GPU memory to CPU memory, the GPU can continue training while the CPU handles the I/O.

To understand why this separation is possible, consider what a checkpoint really is at the moment it is created: a snapshot of tensor values that were valid at step $N$. Once those values are safely in CPU memory, the GPU is free to advance to steps $N+1$, $N+2$, and so on. The CPU copy remains a valid representation of step-$N$ state regardless of what the GPU does afterward. The I/O operation is therefore completely independent of GPU computation, and the two can proceed in parallel.

This observation motivates a clean architectural separation: checkpointing logic consists of a fast "snapshot" phase that briefly blocks training (the GPU-to-CPU copy) and a slow "persist" phase that runs in the background (the CPU-to-disk write). Minimizing the snapshot phase minimizes the training stall, while the persist phase can take as long as necessary without affecting training progress.

In synchronous checkpointing, the workflow at each checkpoint step is:

1. Training pauses
2. GPU tensors are transferred to CPU memory
3. CPU writes tensors to disk
4. Training resumes

Both steps 2 and 3 contribute to the stall. The GPU-to-CPU transfer alone can take several seconds for a large model (at roughly 60 GB/s PCIe 4.0 bandwidth, a 300 GB model takes about 5 seconds), and the disk write takes tens of seconds (at 10 GB/s NFS bandwidth, the same 300 GB takes 30 seconds). During all 35 of these seconds, hundreds of GPUs sit completely idle.

For a checkpoint interval of 10 minutes (600 seconds), this represents about 5.8% overhead from checkpointing alone. Over a 20-day run, you lose more than a full day of compute to synchronous checkpointing. The overhead scales linearly with model size and inversely with checkpoint interval, creating strong pressure to checkpoint less frequently, which in turn increases the expected work lost per failure.

Asynchronous checkpointing breaks this uncomfortable tradeoff. By running the disk write in the background, you pay only the GPU-to-CPU copy cost (the fast serialization phase) while the GPU continues training. The disk write proceeds concurrently and the GPU never waits for it.

Asynchronous checkpointing breaks the write into two stages that overlap with ongoing training:

First, at the checkpoint step, the training loop snapshots the GPU tensors to CPU memory. This pinned-memory copy is fast, typically 2-5 seconds for a 300 GB model at roughly 60-100 GB/s PCIe bandwidth. Training then resumes immediately after this copy, without waiting for anything to be written to disk.

Second, a background thread or process continues transferring the CPU buffer to disk while training proceeds on the GPU. Since the CPU has a complete copy of the checkpoint state at the moment of the snapshot, it does not need the GPU again for I/O. The background writer can proceed at whatever rate the storage system supports.

The result is that the effective checkpoint overhead is reduced to the GPU-to-CPU copy time (a few seconds) rather than the full GPU-to-CPU plus CPU-to-disk time (tens of seconds). For the 300 GB model example above, asynchronous checkpointing reduces the training stall from 35 seconds to 5 seconds, an 85% reduction.

There is one important safety condition: the background write must complete before the next checkpoint is triggered. If the background writer is still writing the step-1000 checkpoint when it is time to checkpoint step-1100, you have two choices: wait for the previous write to finish before starting the new one, or start the new write anyway and accept that you briefly need two full checkpoint buffers in CPU memory. In practice, most systems wait, which means the effective checkpoint interval is bounded below by the write time. If your checkpoint write takes 30 seconds and you want to checkpoint every 10 minutes, asynchronous I/O works fine. If you want to checkpoint every 20 seconds, it does not help.

PyTorch provides torch.distributed.checkpoint with async support through AsyncCheckpointManager. The manager accepts a saving plan and executes it in a background thread, allowing training to proceed concurrently:

The future object can be waited on at a later step if needed, or simply allowed to run to completion in the background. If the process crashes before the write completes, the partial checkpoint is typically left on disk but is not used for recovery (the previous complete checkpoint is used instead, which is why the "latest checkpoint" pointer is updated only after all ranks have confirmed their writes are complete).

Asynchronous checkpointing requires holding two full copies of the model state simultaneously: one on the GPU (continuing training) and one in CPU pinned memory (being written to disk). For a 70B model with full optimizer state at ten bytes per parameter, this means roughly 700 GB of CPU RAM in addition to the GPU memory already in use.

Many clusters have sufficient CPU RAM. Modern servers with 8 GPUs often have 1-2 TB of system RAM. For a 70B model running on 64 H100 GPUs, each rank's shard is about 11 GB, and each host needs two copies of roughly 88 GB of tensors (for 8 GPUs per host), well within typical host memory. But for very large models or memory-constrained systems, the CPU RAM requirement can become a binding constraint.

In such cases, partial asynchronous checkpointing is used: save the optimizer state (the largest component) asynchronously while saving the model parameters synchronously, or vice versa. This reduces the peak CPU memory requirement proportionally, at the cost of a longer training stall than fully asynchronous checkpointing. The system designer must choose the right balance based on the available CPU RAM and the desired training stall duration.

Another memory-conscious approach is staged asynchronous checkpointing: copy one layer's parameters from GPU to CPU, then allow the background writer to drain that layer to disk, then proceed to the next layer. This keeps the CPU buffer small but serializes the I/O, losing some parallelism. It is most useful for single-GPU training where CPU RAM is the dominant constraint.

In distributed training, each rank must participate in checkpointing to ensure the complete model state is saved. The ranks must coordinate to ensure they all snapshot at the same step and that all shards are committed before a checkpoint is considered complete.

Coordination happens at two points. First, all ranks must reach the checkpoint step before any of them begins snapshotting. A global barrier at the checkpoint step ensures this. Second, after all ranks have launched their background writers, a second barrier (or a coordinator-rank acknowledgment protocol) ensures that the "latest checkpoint" pointer is updated only after all shards have been successfully written to durable storage.

If the second barrier is omitted and rank 0 updates the "latest checkpoint" pointer immediately after launching its background write, a process crash on rank 3 partway through its write would leave the checkpoint in an incomplete state with the pointer pointing to it. Recovery would attempt to load from this checkpoint, find that rank 3's shard is corrupt or missing, and fail. The coordinator-rank acknowledgment protocol prevents this by requiring all ranks to confirm success before the pointer is moved.

A typical coordination pattern uses a two-phase approach. In the first phase, all ranks write their shard data asynchronously. In the second phase, each rank reports completion to a coordinator (usually rank 0 or a dedicated checkpoint daemon), and the coordinator updates the metadata only after all ranks have reported. If any rank fails to report within a timeout, the checkpoint is considered incomplete and the previous pointer is retained.