TorchSnapshot (Beta发布)

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

安装

需要Python >= 3.7和PyTorch >= 1.12

从pip

# Stable
pip install 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 size更改时自动重新分片（弹性）（更多详情）。

安全性

• 无需pickle依赖项的安全张量序列化 [WIP]。

入门

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文件中找到。