asynchronous checkpointing

Asynchronous checkpointing is common at scale. In this scheme, the GPU is only blocked when copying checkpoint data from GPU memory to host CPU memory. The GPU then proceeds with computing while the CPU asynchronously flushes the checkpoint data to nonvolatile storage.

There may be multiple levels of nonvolatile storage that are used at different intervals to further reduce bandwidth requirements and increase durability of the checkpoints.

Supported frameworks

Microsoft’s Nebula framework does this by

1. Synchronously checkpointing to CPU memory
2. Asynchronously copying that to an adjacent node’s local SSD to protect against a single-node failure
3. Asynchronously copying that to object storage to protect against multi-node failures, allow rollback, and store checkpoints long-term

Nebula was a project out of Microsoft Research which has been abandoned, but I think DeepSpeed has adopted async checkpointing.

NVIDIA’s NeMo Framework also supports asynchronous checkpointing.

PyTorch also supports asynchronous checkpointing. It sounds very similar to what NeMo has implemented.

Ray offers ways to perform asynchronous checkpointing and recommends writing to local disk.¹ It also offers distributed checkpointing

Google offers transparent “multi-tier checkpointing” using Replicator, which automatically replications node-local checkpoints to peer nodes and automatically drains to GCS.

Companies who use it

Meta uses asynchronous checkpointing to train its Muse models.²

ByteDance’s MegaScale does this.³

Footnotes

1. Saving and Loading Checkpoints — Ray 2.59.0 ↩

2. Inside Prometheus: Meta’s Journey to a Gigawatt-Scale Supercluster - Engineering at Meta ↩

3. [2402.15627] MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs (arxiv.org) ↩