Metadata-Version: 2.4
Name: torch_checkpointing
Version: 0.1.0
Summary: A high-performance, distributed-training-ready checkpointing library for PyTorch models
Author: Meta Platforms, Inc.
License: BSD 3-Clause License
        
        Copyright (c) Meta Platforms, Inc. and affiliates.
        All rights reserved.
        
        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions are met:
        
        * Redistributions of source code must retain the above copyright notice, this
          list of conditions and the following disclaimer.
        
        * Redistributions in binary form must reproduce the above copyright notice,
          this list of conditions and the following disclaimer in the documentation
          and/or other materials provided with the distribution.
        
        * Neither the name of the copyright holder nor the names of its
          contributors may be used to endorse or promote products derived from
          this software without specific prior written permission.
        
        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
        AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
        IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
        DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
        FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
        DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
        SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
        CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
        OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
        OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
        
Project-URL: Homepage, https://github.com/meta-pytorch/torch_checkpointing
Project-URL: Repository, https://github.com/meta-pytorch/torch_checkpointing
Project-URL: Issues, https://github.com/meta-pytorch/torch_checkpointing/issues
Keywords: pytorch,checkpointing,distributed-training,distributed-checkpointing
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: BSD License
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.6.0
Requires-Dist: numpy
Requires-Dist: safetensors>=0.4.3
Requires-Dist: typing_extensions>=4.6.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.0.0; extra == "dev"
Requires-Dist: expecttest; extra == "dev"
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: isort>=5.12.0; extra == "dev"
Requires-Dist: mypy>=1.5.0; extra == "dev"
Provides-Extra: test
Requires-Dist: pytest>=7.0.0; extra == "test"
Requires-Dist: pytest-cov>=4.0.0; extra == "test"
Requires-Dist: expecttest; extra == "test"
Dynamic: license-file

# `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](./docs/extensibility.md).

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

## Installation

```bash
pip install torch_checkpointing
```

Requires Python >= 3.10 and torch >= 2.6.

Saving is asynchronous by default. The optimized async staging defaults
currently require CUDA; CPU-only users should use the explicit configuration in
[Troubleshooting](./docs/troubleshooting.md).

## 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](./docs/tutorials.md) — checkpoint and resume a complete training loop.
- [Overview](./docs/index.md) — what the library does and how the pieces fit together.
- [Key concepts](./docs/key_concepts.md) — the `CheckpointManager`, the payload/`into=` model, and how async save and load work.
- [Configuring checkpoints](./docs/configuring_checkpoints.md) — per-item `layout`, `requires_copy`, and `resharder` via `ItemSpec`.
- [Troubleshooting & FAQ](./docs/troubleshooting.md) — common errors and how to fix them.
- [API reference](./docs/api_reference.md) — the public symbols at a glance.

**Building bespoke components (power users)**

- [Extensibility](./docs/extensibility.md) — the extension points, and how to plug in your own infrastructure.
- [Storage](./docs/storage.md) — the `Storage` / `StorageConfig` interface and writing a custom backend.
- [Distributed and resharding](./docs/distributed_and_resharding.md) — multi-rank saves and custom resharding across mesh/placement changes.
- [Design & internals](./docs/design.md) — the async staging and background-write architecture.

**Contributing**

- [Contributing](./CONTRIBUTING.md) — development setup, testing, and pull-request guidance.

## License

BSD 3-Clause License. See [LICENSE](LICENSE).
