一个针对PyTorch应用程序的高性能、内存高效的检查点库,专为大型、复杂的分布式工作负载设计。
项目描述
TorchSnapshot (Beta发行版)
一个针对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文件所示。
项目详情
下载文件
下载适用于您平台的文件。如果您不确定选择哪个,请了解更多关于安装包的信息。
源分发
torchsnapshot_nightly-2024.7.26.tar.gz (93.0 kB 查看哈希值)
构建分发
关闭
哈希值 for torchsnapshot_nightly-2024.7.26-py3-none-any.whl
算法 | 哈希摘要 | |
---|---|---|
SHA256 | 8fa53e4eda4fc334bafc523f3245ddab767f8666913164036ff960983b55051b |
|
MD5 | 1c505b2d00201e658000ceeadf1de946 |
|
BLAKE2b-256 | f205f2f19c1838811757e80a5e64ad1401154d09e90bd2caebd0b30ce7f557da |