TorchSnapshot (Beta Release)
A performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind.
Install
Requires Python >= 3.8 and PyTorch >= 2.0.0
From pip:
# Stable
pip install torchsnapshot
# Or, using conda
conda install -c conda-forge torchsnapshot
# Nightly
pip install --pre torchsnapshot-nightly
From source:
git clone https://github.com/pytorch/torchsnapshot
cd torchsnapshot
pip install -r requirements.txt
python setup.py install
Why TorchSnapshot
Performance
- TorchSnapshot provides a fast checkpointing implementation employing various optimizations, including zero-copy serialization for most tensor types, overlapped device-to-host copy and storage I/O, parallelized storage I/O.
- TorchSnapshot greatly speeds up checkpointing for DistributedDataParallel workloads by distributing the write load across all ranks (benchmark).
- When host memory is abundant, TorchSnapshot allows training to resume before all storage I/O completes, reducing the time blocked by checkpoint saving.
Memory Usage
- TorchSnapshot's memory usage adapts to the host's available resources, greatly reducing the chance of out-of-memory issues when saving and loading checkpoints.
- TorchSnapshot supports efficient random access to individual objects within a snapshot, even when the snapshot is stored in a cloud object storage.
Usability
- Simple APIs that are consistent between distributed and non-distributed workloads.
- Out of the box integration with commonly used cloud object storage systems.
- Automatic resharding (elasticity) on world size change for supported workloads (more details).
Security
- Secure tensor serialization without pickle dependency [WIP].
Getting Started
from torchsnapshot import Snapshot
# Taking a snapshot
app_state = {"model": model, "optimizer": optimizer}
snapshot = Snapshot.take(path="/path/to/snapshot", app_state=app_state)
# Restoring from a snapshot
snapshot.restore(app_state=app_state)
See the documentation for more details.
License
torchsnapshot is BSD licensed, as found in the LICENSE file.
Metadata
Release files for torchsnapshot-nightly 2024.7.26
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | Actions |
|---|---|---|---|---|
| torchsnapshot_nightly-2024.7.26.tar.gz | 93.0 kB | Without provenance | Details for torchsnapshot_nightly-2024.7.26.tar.gz |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Actions Reset |
|---|---|---|---|---|
| torchsnapshot_nightly-2024.7.26-py3-none-any.whl | Python 3 | none | any | Details for torchsnapshot_nightly-2024.7.26-py3-none-any.whl |
Total release size: 177.6 kB
Release files / torchsnapshot_nightly-2024.7.26.tar.gz
| Download URL | torchsnapshot_nightly-2024.7.26.tar.gz |
|---|---|
| Size | 93.0 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
48ef217adb7f2f46b54391e056276fda9f8ae9415254b313b8ccad0f00393928
|
|
BLAKE2b-256 checksum How to use checksums |
e9b12d239671cd5622ab47975d48b97e290e51b34b8f8a7e588fab9bc8e22d78
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.1.1 CPython/3.8.19
|
Release files / torchsnapshot_nightly-2024.7.26-py3-none-any.whl
| Download URL | torchsnapshot_nightly-2024.7.26-py3-none-any.whl |
|---|---|
| Size | 84.7 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
8fa53e4eda4fc334bafc523f3245ddab767f8666913164036ff960983b55051b
|
|
BLAKE2b-256 checksum How to use checksums |
f205f2f19c1838811757e80a5e64ad1401154d09e90bd2caebd0b30ce7f557da
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/5.1.1 CPython/3.8.19
|
Security for torchsnapshot-nightly 2024.7.26
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
To learn more about protecting yourself from malicious actors, visit our package security documentation.
