我们非常高兴地宣布 PyTorch 1.8 正式发布。自 1.7 版本以来,该版本共包含 3,000 多次提交。它在编译、代码优化、科学计算前端 API 方面进行了重大更新和功能增强,并通过 pytorch.org 提供的二进制文件支持 AMD ROCm。此外,它还改进了流水线并行、模型并行和梯度压缩等大规模训练功能。以下是部分亮点:
- 支持通过
torch.fx进行 Python 到 Python 的函数式转换; - 添加或稳定了对 FFT(
torch.fft)和线性代数函数(torch.linalg)的支持,增加了对复数张量的自动求导支持,并更新了以提高计算 Hessian 和 Jacobian 的性能;以及 - 对分布式训练进行了重大更新和改进,包括:提升了 NCCL 的可靠性;支持流水线并行;RPC 分析;以及支持通过通信钩子(communication hooks)进行梯度压缩。查看完整发布说明请点击此处。
随着 1.8 版本的发布,我们还对 PyTorch 库进行了重大更新,包括 TorchCSPRNG、TorchVision、TorchText 和 TorchAudio。有关库发布的更多信息,请参阅此文章。如前所述,PyTorch 发布版中的功能分为“稳定版”(Stable)、“测试版”(Beta)和“原型版”(Prototype)。您可以在此文章中了解更多定义。
全新与更新的 API
PyTorch 1.8 版本带来了一系列全新和更新的 API,涵盖了更多兼容 NumPy 的 API,以及旨在提高代码性能和扩展性(推理和训练阶段)的方法。以下是本次发布主要功能的简要总结:
[稳定版] Torch.fft 支持高性能 NumPy 风格 FFT
作为 PyTorch 支持科学计算目标的一部分,我们投入了大量精力改进 FFT 支持。在 PyTorch 1.8 中,我们发布了 torch.fft 模块。该模块实现了与 NumPy 的 np.fft 模块相同的函数,但支持硬件加速和自动求导。
[测试版] 通过 torch.linalg 支持 NumPy 风格的线性代数函数
torch.linalg 模块仿照 NumPy 的 np.linalg 模块,为常见的线性代数运算提供了 NumPy 风格的支持,包括 Cholesky 分解、行列式、特征值等。
[测试版] 使用 FX 进行 Python 代码转换
FX 允许您编写形式为 transform(input_module : nn.Module) -> nn.Module 的转换,您可以输入一个 Module 实例并获得一个转换后的 Module 实例。
此类功能适用于多种场景。例如,基于 FX 的图模式量化(Graph Mode Quantization)产品与 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 中使用通信钩子实现的梯度压缩。详细信息如下:
[测试版] 流水线并行(Pipeline Parallelism)
随着机器学习模型规模不断增长,传统的分布式数据并行(DDP)训练不再具备扩展性,因为这些模型无法放入单个 GPU 设备中。新的流水线并行功能提供了一个易于使用的 PyTorch API,以便在您的训练循环中利用流水线并行。
[测试版] DDP 通信钩子(DDP Communication Hook)
DDP 通信钩子是一个通用接口,通过重写 DistributedDataParallel 中的原生 allreduce,从而控制工作节点间梯度通信的方式。提供了一些内置的通信钩子,包括 PowerSGD,用户可以轻松应用这些钩子来优化通信。此外,通信钩子接口还可以支持用户定义的通信策略,以满足更高级的用例。
分布式训练的其他原型功能
除了此版本中主要的稳定和测试版分布式训练功能外,我们的夜间构建版本中还提供了一些原型功能供大家试用并提供反馈。我们已在下方的草案文档中提供了参考链接。
- (原型) ZeroRedundancyOptimizer – 基于并与微软 DeepSpeed 团队合作,此功能通过在
ProcessGroup组的所有参与进程中分片优化器状态,帮助减少每个进程的内存占用。有关详细信息,请参考此文档。 - (原型) 进程组 NCCL 发送/接收(Process Group NCCL Send/Recv) – NCCL 发送/接收 API 在 v2.7 中引入,此功能增加了在 NCCL 进程组中的支持。此功能将为用户提供一种在 Python 层而不是 C++ 层实现集合操作的选择。请参考此文档和代码示例了解更多。
- (原型) 使用 TensorPipe 进行 RPC 中的 CUDA 支持 – 对于在多 GPU 机器上使用 PyTorch RPC 的用户,此功能应能带来显著的速度提升。TensorPipe 将在可用时自动利用 NVLink,并在进程间交换 GPU 张量时避免昂贵的主机内存拷贝。如果不在同一台机器上,TensorPipe 将回退到将张量复制到主机内存并将其作为常规 CPU 张量发送。这也将改善用户体验,因为用户能够在代码中像对待常规 CPU 张量一样对待 GPU 张量。有关详细信息,请参考此文档。
- (原型) 远程模块(Remote Module) – 此功能允许用户像使用本地模块一样操作远程工作节点上的模块,其中 RPC 对用户是透明的。过去,此功能以特殊方式实现,总体而言,此功能将提高 PyTorch 上模型并行的可用性。有关详细信息,请参考此文档。
PyTorch Mobile
PyTorch Mobile 的支持正在通过一套新的教程不断扩展,以帮助新用户更快地在设备上启动模型,并为现有用户提供充分利用我们框架的工具。这些包括:
我们的新演示应用程序还包括图像分割、目标检测、神经机器翻译、问答和视觉 Transformer 的示例。它们均可在 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 版本中,我们发布了基准测试工具(benchmark utils)的支持,使用户能够更好地监控性能。我们还开放了一个新的自动化量化 API。详见下文:
(测试版) 基准测试工具(Benchmark utils)
基准测试工具允许用户进行精确的性能测量,并提供可组合的工具来帮助进行基准测试制定和后期处理。预计这将有助于 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 Graph Mode Quantization)
FX 图模式量化是 PyTorch 中新的自动化量化 API。它通过增加对函数式的支持和自动化量化过程,改进了 Eager 模式量化。尽管用户可能需要重构模型以使其与 FX 图模式量化兼容(可以使用 torch.fx 进行符号追踪)。
硬件支持
[测试版] 在 C++ 中扩展 PyTorch 分发器(Dispatcher)以支持新后端的可能性
在 PyTorch 1.8 中,您现在可以创建在 pytorch/pytorch 仓库之外的树外(out-of-tree)设备。下方链接的教程展示了如何注册您的设备并使其与原生 PyTorch 设备保持同步。
[测试版] AMD GPU 二进制文件现已推出
从 PyTorch 1.8 开始,我们增加了对 ROCm wheels 的支持,为使用 AMD GPU 提供了便捷的入门途径。您可以直接访问标准的 PyTorch 安装选择器,选择 ROCm 作为安装选项,然后执行提供的命令。
感谢阅读!如果您对这些更新感到兴奋并希望参与 PyTorch 的未来,我们鼓励您加入讨论论坛并在 GitHub 上提交问题。
干杯!
PyTorch 团队