我们很高兴地宣布 PyTorch 1.8 版本现已推出。此版本自 1.7 以来包含了 3,000 多次提交。它包含了编译、代码优化、科学计算前端 API 和 AMD ROCm 支持(通过 pytorch.org 提供的二进制文件)的重大更新和新功能。它还为管道并行和模型并行的大规模训练以及梯度压缩提供了改进的功能。其中一些亮点包括:
- 通过
torch.fx
支持 Python 到 Python 的功能转换; - 添加或稳定了 API,以支持 FFT (
torch.fft
)、线性代数函数 (torch.linalg
),并添加了对复杂张量自动微分的支持,以及改进了计算 Hessian 和 Jacobian 的性能;以及 - 分布式训练的重大更新和改进,包括:NCCL 可靠性改进;管道并行支持;RPC 性能分析;以及支持添加梯度压缩的通信钩子。完整的发布说明请参阅此处。
与 1.8 版本一起,我们还发布了 PyTorch 库的重大更新,包括 TorchCSPRNG、TorchVision、TorchText 和 TorchAudio。有关库发布的更多信息,请参阅此处的帖子。如前所述,PyTorch 版本中的功能分为稳定版、Beta 版和原型版。您可以在此处的帖子中了解有关定义的更多信息。
新增和更新的 API
PyTorch 1.8 版本带来了一系列新的和更新的 API 表面,包括用于 NumPy 兼容性的附加 API,以及改进和扩展代码以在推理和训练时提高性能的方法。以下是此版本中主要功能的简要摘要:
[稳定版] Torch.fft
支持高性能 NumPy 风格的 FFT
作为 PyTorch 支持科学计算目标的一部分,我们投入了大量精力来改进我们的 FFT 支持,并在 PyTorch 1.8 中发布了 torch.fft
模块。此模块实现了与 NumPy 的 np.fft
模块相同的功能,但支持硬件加速和自动微分。
[Beta 版] 通过 torch.linalg
支持 NumPy 风格的线性代数函数
torch.linalg
模块模仿 NumPy 的 np.linalg 模块,为常见的线性代数运算(包括 Cholesky 分解、行列式、特征值等)带来了 NumPy 风格的支持。
[Beta 版] 使用 FX 进行 Python 代码转换
FX 允许您编写 transform(input_module : nn.Module)
-> nn.Module
形式的转换,您可以将 Module
实例输入其中并从中获得转换后的 Module
实例。
这种功能适用于许多场景。例如,基于 FX 的图模式量化产品与 FX 同时作为原型发布。图模式量化自动化了神经网络的量化过程,并利用了 FX 的程序捕获、分析和转换功能。我们还在使用 FX 开发许多其他转换产品,我们很高兴与社区分享这个强大的工具包。
由于 FX 转换消耗和生成 nn.Module 实例,它们可以在许多现有的 PyTorch 工作流程中使用。这包括例如在 Python 中训练然后通过 TorchScript 部署的工作流程。
您可以在官方文档中阅读更多关于 FX 的信息。您还可以在此处找到几个使用 torch.fx
实现的程序转换示例。我们正在不断改进 FX,并邀请您在论坛或问题跟踪器上分享您对工具包的任何反馈。
我们谨此鸣谢 TorchScript 追踪、Apache MXNet hybridize,以及最近的 JAX 作为通过追踪进行程序获取的影响者。我们还谨此鸣谢 Caffe2、JAX 和 TensorFlow 作为简单、有向数据流图程序表示及其转换价值的灵感来源。
分布式训练
PyTorch 1.8 版本新增了许多功能,并改进了可靠性和可用性。具体来说,增加了对 稳定级别的异步错误/超时处理 的支持,以提高 NCCL 的可靠性;以及对 基于 RPC 的性能分析 的稳定支持。此外,我们还通过在 DDP 中使用通信钩子增加了对管道并行和梯度压缩的支持。详情如下:
[Beta 版] 管道并行
随着机器学习模型规模的不断增长,传统的分布式数据并行 (DDP) 训练不再适用,因为这些模型无法在单个 GPU 设备上运行。新的管道并行功能提供了一个易于使用的 PyTorch API,可将管道并行作为训练循环的一部分。
[Beta 版] DDP 通信钩子
DDP 通信钩子是一个通用接口,通过覆盖 DistributedDataParallel 中的普通 allreduce 来控制如何跨 workers 通信梯度。提供了几个内置的通信钩子,包括 PowerSGD,用户可以轻松应用这些钩子中的任何一个来优化通信。此外,通信钩子接口还可以支持用户定义的通信策略,以应对更高级的用例。
分布式训练的其他原型功能
除了此版本中主要的稳定版和 Beta 版分布式训练功能外,我们的夜间构建中还提供了一些原型功能,供您试用并提供反馈。我们已在下方链接了草稿文档以供参考:
- (原型) ZeroRedundancyOptimizer – 基于 Microsoft DeepSpeed 团队并与其合作,此功能通过在
ProcessGroup
组中的所有参与进程之间分片优化器状态,帮助减少每个进程的内存占用。有关更多详细信息,请参阅此文档。 - (原型) 进程组 NCCL 发送/接收 – NCCL 发送/接收 API 在 v2.7 中引入,此功能为其在 NCCL 进程组中添加了支持。此功能将为用户提供一个选项,以便在 Python 层而不是 C++ 层实现集体操作。请参阅此文档和代码示例以了解更多信息。
- (原型) 使用 TensorPipe 在 RPC 中支持 CUDA – 此功能将为使用 PyTorch RPC 的多 GPU 机器用户带来显著的速度改进,因为 TensorPipe 会在可用时自动利用 NVLink,并避免在进程之间交换 GPU 张量时进行昂贵的主机内存复制。当不在同一台机器上时,TensorPipe 将回退到将张量复制到主机内存并将其作为常规 CPU 张量发送。这还将改善用户体验,因为用户将能够在代码中将 GPU 张量视为常规 CPU 张量。有关更多详细信息,请参阅此文档。
- (原型) 远程模块 – 此功能允许用户在远程工作器上操作模块,就像使用本地模块一样,其中 RPC 对用户是透明的。过去,此功能以临时方式实现,总的来说,此功能将提高 PyTorch 上模型并行的可用性。有关更多详细信息,请参阅此文档。
PyTorch Mobile
PyTorch Mobile 的支持正在扩展,新增了一系列教程,以帮助新用户更快地在设备上启动模型,并为现有用户提供一个工具,以从我们的框架中获得更多收益。其中包括:
我们的新演示应用程序还包括图像分割、目标检测、神经机器翻译、问答和视觉转换器的示例。它们在 iOS 和 Android 上均可用。
除了 MobileNetV3 和其他模型的 CPU 性能改进之外,我们还改进了 Android GPU 后端原型,以实现更广泛的模型覆盖和更快的推理速度。
最后,我们在此版本中作为原型功能推出了 PyTorch Mobile Lite Interpreter。Lite Interpreter 允许用户减小运行时二进制文件的大小。请试用这些功能,并在 PyTorch 论坛上向我们发送您的反馈。所有最新更新都可以在 PyTorch Mobile 页面上找到。
[原型] PyTorch Mobile Lite Interpreter
PyTorch Lite Interpreter 是 PyTorch 运行时的一个精简版本,它可以在资源受限的设备上执行 PyTorch 程序,并减小二进制文件大小。与当前版本中的设备上运行时相比,此原型功能将二进制文件大小减少了多达 70%。
性能优化
在 1.8 版本中,我们发布了对基准测试工具的支持,以使用户能够更好地监控性能。我们还开放了一个新的自动化量化 API。详情如下:
(Beta 版) 基准测试工具
基准测试工具允许用户进行准确的性能测量,并提供可组合的工具来帮助基准测试制定和后处理。这预计将有助于 PyTorch 的贡献者快速了解他们的贡献如何影响 PyTorch 性能。
例子
from torch.utils.benchmark import Timer
results = []
for num_threads in [1, 2, 4]:
timer = Timer(
stmt="torch.add(x, y, out=out)",
setup="""
n = 1024
x = torch.ones((n, n))
y = torch.ones((n, 1))
out = torch.empty((n, n))
""",
num_threads=num_threads,
)
results.append(timer.blocked_autorange(min_run_time=5))
print(
f"{num_threads} thread{'s' if num_threads > 1 else ' ':<4}"
f"{results[-1].median * 1e6:>4.0f} us " +
(f"({results[0].median / results[-1].median:.1f}x)" if num_threads > 1 else '')
)
1 thread 376 us
2 threads 189 us (2.0x)
4 threads 99 us (3.8x)
(原型) FX 图模式量化
FX 图模式量化是 PyTorch 中新的自动化量化 API。它通过增加对函数式的支持和自动化量化过程来改进 Eager 模式量化,尽管人们可能需要重构模型以使其与 FX 图模式量化兼容(可通过 torch.fx
进行符号追踪)。
硬件支持
[Beta 版] 在 C++ 中为新后端扩展 PyTorch Dispatcher 的能力
在 PyTorch 1.8 中,您现在可以创建存在于 pytorch/pytorch
存储库之外的新树外设备。下面链接的教程展示了如何注册您的设备并使其与原生 PyTorch 设备保持同步。
[Beta 版] AMD GPU 二进制文件现已可用
从 PyTorch 1.8 开始,我们添加了对 ROCm wheels 的支持,为使用 AMD GPU 提供了简便的入门。您只需访问标准的 PyTorch 安装选择器,选择 ROCm 作为安装选项并执行提供的命令。
感谢您的阅读,如果您对这些更新感到兴奋并希望参与 PyTorch 的未来,我们鼓励您加入讨论论坛并提出 GitHub 问题。
干杯!
PyTorch 团队