torch_checkpointing

High-performance asynchronous checkpointing for PyTorch. It takes checkpoint saving off your training loop's critical path:

• Zero-overhead saves — save() returns immediately; model state is staged off the training device and written by a background process while your training step keeps running.
• Save and load through one API — a single CheckpointManager drives both; you pass plain {name: value} dicts and decide when to block on a save (for example, before exit).
• Single-rank to distributed — the same API scales from one process to large distributed jobs, and reshards across different parallelism layouts on load.

You interact with one object, CheckpointManager: save(checkpoint_id, {...}) and load(checkpoint_id, into={...}) over a pluggable storage backend. A checkpoint_id is a string interpreted by that backend; the default local filesystem backend treats it as the path to a checkpoint directory. Rank, storage, sharding metadata, and per-item copy/reshard behavior are configured for you. Power users can still swap in bespoke components — storage backends, resharders, cross-rank coordination — through the extension points.

Experimental and pre-1.0. The public API may still change.

Installation

pip install torch_checkpointing

Requires Python >= 3.10 and torch >= 2.11.

Saving is asynchronous by default. The optimized async staging defaults currently require CUDA; CPU-only users should use the explicit configuration in Troubleshooting.

Key features

• Non-blocking async saves overlapped with training (host-side staging + a background-process write).
• One high-level CheckpointManager for both save and load, with auto-detected rank, storage, and metadata.
• Plain-dict payloads: save(id, {...}) / load(id, into={...}) — tensors restored in place (identity preserved), scalars and JSON/bytes are first-class top-level items.
• Resharding on load across different distributed layouts (mesh / placement changes), wired automatically when an item declares a resharder.
• Pluggable storage behind the Storage / StorageConfig interface; a local filesystem backend ships in the package.

Documentation

Getting started

• Tutorial — checkpoint and resume a complete training loop.
• Overview — what the library does and how the pieces fit together.
• Key concepts — the CheckpointManager, the payload/into= model, and how async save and load work.
• Configuring checkpoints — per-item layout, requires_copy, and resharder via ItemSpec.
• Troubleshooting & FAQ — common errors and how to fix them.
• API reference — the public symbols at a glance.

Building bespoke components (power users)

• Extensibility — the extension points, and how to plug in your own infrastructure.
• Storage — the Storage / StorageConfig interface and writing a custom backend.
• Distributed and resharding — multi-rank saves and custom resharding across mesh/placement changes.
• Design & internals — the async staging and background-write architecture.

Contributing

• Contributing — development setup, testing, and pull-request guidance.

License

BSD 3-Clause License. See LICENSE.