Skip to main content

A performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind.

Project description

TorchSnapshot (Beta Release)

build status pypi version pypi nightly version codecov bsd license

A performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind.

Install

Requires Python >= 3.7 and PyTorch >= 1.12

From pip:

# Stable
pip install 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.

Project details


Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

torchsnapshot-nightly-2022.11.24.tar.gz (52.3 kB view details)

Uploaded Source

Built Distribution

torchsnapshot_nightly-2022.11.24-py3-none-any.whl (68.6 kB view details)

Uploaded Python 3

File details

Details for the file torchsnapshot-nightly-2022.11.24.tar.gz.

File metadata

File hashes

Hashes for torchsnapshot-nightly-2022.11.24.tar.gz
Algorithm Hash digest
SHA256 c266a2b80b07c37ceba281fe6405bdd45c7a6665f13c85e244a75a6e6da3ee01
MD5 1403821c8840b8c330bd086ef0d29322
BLAKE2b-256 55ee55f9b4431c5cf4166e3b9ba8eee1c585963c3c8f3b451c8fa7e0604022cc

See more details on using hashes here.

Provenance

File details

Details for the file torchsnapshot_nightly-2022.11.24-py3-none-any.whl.

File metadata

File hashes

Hashes for torchsnapshot_nightly-2022.11.24-py3-none-any.whl
Algorithm Hash digest
SHA256 22cb8c572545e97c287745f6d2da51bcffc8070074192d8a7f6edc298b38148d
MD5 9872d91a7cd6f149d0bb9d89194af260
BLAKE2b-256 01802df5d89647484bd10fd1517d8a73faddd041da749db9020ad53324f3db44

See more details on using hashes here.

Provenance

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page