跳转到主要内容
博客

PyTorch 1.8 发布,包括编译器和分布式训练更新,以及新的移动教程

作者: 2021年3月4日2024年11月16日暂无评论

我们很高兴地宣布 PyTorch 1.8 已发布。自 1.7 版本以来,此版本包含 3,000 多次提交。它包括编译、代码优化、科学计算前端 API 以及通过 pytorch.org 提供的二进制文件支持 AMD ROCm 的重大更新和新功能。它还为管道和模型并行性以及梯度压缩提供了改进的大规模训练功能。其中一些亮点包括:

  1. 通过 torch.fx 支持 Python 到 Python 的函数式转换;
  2. 添加或稳定了支持 FFT(torch.fft)和线性代数函数(torch.linalg)的 API,增加了对复数张量自动求导的支持,并更新以提高赫斯矩阵和雅可比矩阵的计算性能;以及
  3. 分布式训练的重大更新和改进,包括:提高 NCCL 可靠性;支持管道并行性;RPC 性能分析;以及支持添加梯度压缩的通信钩子。请在此处查看完整的发布说明:此处

除了 1.8 版本,我们还发布了 PyTorch 库的重大更新,包括 TorchCSPRNGTorchVisionTorchTextTorchAudio。有关库发布的更多信息,请参阅 此处 的文章。如前所述,PyTorch 发布中的功能分为稳定版、测试版和原型版。您可以在 此处 的文章中了解更多关于这些定义的详细信息。

新增和更新的 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 作为通过跟踪进行程序获取的影响者。我们还要感谢 Caffe2JAXTensorFlow 作为简单、有向数据流图程序表示及其转换的价值的灵感来源。

分布式训练

PyTorch 1.8 版本增加了一些新功能,并改进了可靠性和可用性。具体来说,支持:稳定级别的异步错误/超时处理已添加到以提高 NCCL 可靠性;并稳定支持 基于 RPC 的性能分析。此外,我们还通过在 DDP 中使用通信钩子增加了管道并行性支持和梯度压缩。详细信息如下:

[Beta] 管道并行性

随着机器学习模型规模的不断增长,传统的分布式数据并行(DDP)训练不再适用,因为这些模型无法在单个 GPU 设备上容纳。新的管道并行性功能提供了一个易于使用的 PyTorch API,可将管道并行性作为训练循环的一部分。

[Beta] DDP 通信钩子

DDP 通信钩子是一个通用接口,通过覆盖 DistributedDataParallel 中的普通 allreduce 来控制如何跨工作进程通信梯度。提供了几个内置通信钩子,包括 PowerSGD,用户可以轻松应用任何这些钩子来优化通信。此外,通信钩子接口还可以支持用户自定义的通信策略,以应对更高级的用例。

分布式训练的其他原型功能

除了此版本中主要的稳定版和 Beta 版分布式训练功能之外,我们的每夜构建版本中还有一些原型功能可供试用和提供反馈。我们已在下面链接了草案文档供参考:

  • (原型)ZeroRedundancyOptimizer – 基于微软 DeepSpeed 团队并与其合作,此功能通过将优化器状态分片到 ProcessGroup 组中的所有参与进程来帮助减少每个进程的内存占用。有关更多详细信息,请参阅此文档
  • (原型)进程组 NCCL 发送/接收 – NCCL 发送/接收 API 在 v2.7 中引入,此功能为其在 NCCL 进程组中添加了支持。此功能将为用户提供一个选项,以在 Python 层而不是 C++ 层实现集合操作。有关更多信息,请参阅此文档代码示例
  • (原型)使用 TensorPipe 在 RPC 中支持 CUDA – 此功能应为使用多 GPU 机器的 PyTorch RPC 用户带来显著的速度提升,因为 TensorPipe 在可用时会自动利用 NVLink,并避免在进程之间交换 GPU 张量时进行昂贵的宿主内存拷贝。当不在同一台机器上时,TensorPipe 将回退到将张量拷贝到宿主内存并将其作为常规 CPU 张量发送。这还将改善用户体验,因为用户将能够在代码中将 GPU 张量视为常规 CPU 张量。有关更多详细信息,请参阅此文档
  • (原型)远程模块 – 此功能允许用户像使用本地模块一样在远程工作器上操作模块,其中 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 版本中,我们发布了对基准测试工具的支持,以使用户能够更好地监控性能。我们还开放了一个新的自动化量化 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 团队