跳转到主要内容

一个针对PyTorch应用程序的高性能、内存高效的检查点库,专为大型、复杂的分布式工作负载设计。

项目描述

TorchSnapshot (Beta发行版)

build status pypi version conda version pypi nightly version codecov bsd license

一个针对PyTorch应用程序的高性能、内存高效的检查点库,专为大型、复杂的分布式工作负载设计。

安装

需要Python >= 3.8和PyTorch >= 2.0.0

从pip

# Stable
pip install torchsnapshot
# Or, using conda
conda install -c conda-forge torchsnapshot

# Nightly
pip install --pre torchsnapshot-nightly

从源

git clone https://github.com/pytorch/torchsnapshot
cd torchsnapshot
pip install -r requirements.txt
python setup.py install

为什么选择TorchSnapshot

性能

  • TorchSnapshot通过采用包括对大多数张量类型进行零拷贝序列化、重叠设备到主机复制和存储I/O、并行化存储I/O在内的各种优化,提供了一个快速的检查点实现。
  • TorchSnapshot通过在整个rank之间分配写负载,大大加快了DistributedDataParallel工作负载的检查点过程(基准测试)。
  • 当主机内存充足时,TorchSnapshot允许在所有存储I/O完成之前恢复训练,减少由检查点保存引起的阻塞时间。

内存使用

  • TorchSnapshot的内存使用量会根据主机可用资源进行调整,大大减少了保存和加载检查点时出现内存不足问题的可能性。
  • TorchSnapshot支持对快照中单个对象的高效随机访问,即使快照存储在云对象存储中也是如此。

可用性

  • 在分布式和非分布式工作负载之间保持一致的简单API。
  • 与常用云对象存储系统开箱即用的集成。
  • 对支持的工作负载在world大小变化时自动重新分片(弹性)(更多信息)。

安全

  • 无需pickle依赖的安全张量序列化 [进行中]。

入门指南

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)

有关更多详细信息,请参阅文档

许可证

torchsnapshot遵循BSD许可协议,如LICENSE文件所示。

项目详情


发行历史 发布通知 | RSS源

下载文件

下载适用于您平台的文件。如果您不确定选择哪个,请了解更多关于安装包的信息。

源分发

torchsnapshot_nightly-2024.7.26.tar.gz (93.0 kB 查看哈希值)

上传时间

构建分发

torchsnapshot_nightly-2024.7.26-py3-none-any.whl (84.7 kB 查看哈希值)

上传时间 Python 3

由以下支持

AWS AWS 云计算和安全赞助商 Datadog Datadog 监控 Fastly Fastly CDN Google Google 下载分析 Microsoft Microsoft PSF赞助商 Pingdom Pingdom 监控 Sentry Sentry 错误记录 StatusPage StatusPage 状态页面