PyTorch 1.12：TorchArrow、模块功能 API 和 nvFuser 现已发布

我们非常高兴地宣布 PyTorch 1.12 版本发布（发布说明）！此版本包含超过 3124 次提交和 433 位贡献者。随 1.12 版本一同发布的还有 AWS S3 集成、CPU 上 Channels Last 格式的 PyTorch 视觉模型、通过 Bfloat16 和 FSDP API 为 Intel® Xeon® 可扩展处理器赋能 PyTorch 的测试版功能。我们衷心感谢社区成员的贡献。

总结

用于在给定参数集下以函数式方式应用模块计算的功能 API
PyTorch 中的 Complex32 和复数卷积
TorchData 中的 DataPipes 与 DataLoader 完全向后兼容
functorch 改进了 API 覆盖范围
nvFuser：用于 PyTorch 的深度学习编译器
Ampere 及更新 CUDA 硬件上 float32 矩阵乘法精度的变更
TorchArrow：一个新的机器学习批量数据预处理测试版库

前端 API

引入 TorchArrow

我们为您准备了一个新的测试版库：TorchArrow。这是一个用于机器学习批量数据预处理的库。它具有高性能、类似 Pandas 且易于使用的 API，旨在加快您的预处理工作流程和开发速度。

目前，它提供了一个 Python DataFrame 接口，具有以下功能：

高性能 CPU 后端，利用 Velox 实现向量化且可扩展的用户定义函数 (UDF)
与 PyTorch 或其他模型创作工具无缝对接，例如 Tensor 归集，并可轻松插入 PyTorch DataLoader 和 DataPipes
通过 Arrow 内存中列式格式实现外部读取器的零拷贝

有关更多详细信息，请参阅我们的 10 分钟教程、安装指南、API 文档，以及用于 TorchRec 数据预处理的原型。

（测试版）模块功能 API

PyTorch 1.12 引入了一个新的测试版功能，用于在给定参数集下以函数式方式应用模块计算。有时，内部维护一组静态参数的传统 PyTorch 模块使用模式限制较多。在实现元学习算法时，往往需要在优化器步骤中维护多组参数，此时该功能尤为有用。

新的 torch.nn.utils.stateless.functional_call() API 允许：

对所使用的参数集具有完全灵活性的模块计算
无需以函数式方式重新实现您的模块
模块中存在的任何参数或缓冲区都可以与外部定义的值进行交换以用于调用。引用参数/缓冲区的命名遵循模块 state_dict() 中的全限定形式。

例子

import torch
from torch import nn
from torch.nn.utils.stateless import functional_call

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3, 3)
        self.bn = nn.BatchNorm1d(3)
        self.fc2 = nn.Linear(3, 3)

    def forward(self, x):
        return self.fc2(self.bn(self.fc1(x)))

m = MyModule()

# Define parameter / buffer values to use during module computation.
my_weight = torch.randn(3, 3, requires_grad=True)
my_bias = torch.tensor([1., 2., 3.], requires_grad=True)
params_and_buffers = {
    'fc1.weight': my_weight,
    'fc1.bias': my_bias,
    # Custom buffer values can be used too.
    'bn.running_mean': torch.randn(3),
}

# Apply module computation to the input with the specified parameters / buffers.
inp = torch.randn(5, 3)
output = functional_call(m, params_and_buffers, inp)

（测试版）PyTorch 中的 Complex32 和复数卷积

目前，PyTorch 原生支持复数、复数自动微分、复数模块以及大量复数运算，包括线性代数和快速傅里叶变换 (FFT) 算子。许多库（包括 torchaudio 和 ESPNet）已经在使用 PyTorch 中的复数，而 PyTorch 1.12 进一步扩展了复数功能，增加了复数卷积和实验性的 complex32（“复数半精度”）数据类型，支持半精度 FFT 运算。由于 CUDA 11.3 包存在 Bug，如果您正在使用复数，建议使用 wheels 中的 CUDA 11.6 包。

（测试版）前向模式自动微分

前向模式自动微分 (AD) 允许在前向传递中即时计算方向导数（或等效的雅可比向量积）。PyTorch 1.12 显著改进了前向模式 AD 的算子覆盖范围。更多信息请参阅我们的教程。

TorchData

BC DataLoader + DataPipe

来自 TorchData 的 `DataPipe` 在多进程和分布式环境中的 shuffle 确定性和动态分片方面，实现了与现有 `DataLoader` 的完全向后兼容。

（测试版）AWS S3 集成

基于 AWSSDK 的 DataPipes 已集成到 TorchData 中。它提供了由原生 AWSSDK 支持的以下功能：

根据前缀从每个 S3 存储桶检索 URL 列表
- 支持超时设置，防止无限挂起
- 支持指定 S3 存储桶区域
从 S3 URL 加载数据
- 支持缓冲和分段下载
- 支持指定 S3 存储桶区域

AWS 原生 DataPipes 目前仍处于测试阶段。我们将持续调整以提升其性能。

（原型）DataLoader2

DataLoader2 现以原型模式可用。我们引入了 DataPipes、DataLoading API 和后端（即 ReadingServices）之间交互的新方式。该功能在 API 层面已趋于稳定，但在功能上尚不完善。我们欢迎早期采用者提供反馈，并期待潜在的贡献者。

更多详情，请查看链接。

functorch

受 Google JAX 的启发，functorch 是一个提供可组合 vmap（向量化）和自动微分变换的库。它实现了在 PyTorch 中难以表达的高级自动微分用例，例如：

我们很高兴发布 functorch 0.2.0，它带来了一些改进和新的实验性功能。

显著改进的覆盖范围

我们显著改进了对 functorch.jvp（我们的前向模式自动微分 API）及其他依赖于它的 API（functorch.{jacfwd, hessian}）的覆盖范围。

（原型）functorch.experimental.functionalize

给定一个函数 f，functionalize(f) 将返回一个无变异（有例外情况）的新函数。这对于构建不含原地操作的 PyTorch 函数追踪非常有用。例如，您可以使用 make_fx(functionalize(f)) 来构建 pytorch 函数的无变异追踪。要了解更多信息，请参阅文档。

更多详情，请参见我们的安装指南、文档、教程以及发布说明。

性能改进

引入 nvFuser：用于 PyTorch 的深度学习编译器

在 PyTorch 1.12 中，Torchscript 将其默认熔合器（针对 Volta 及以后的 CUDA 加速器）更新为 nvFuser。与之前 CUDA 设备的熔合器 NNC 相比，nvFuser 支持更广泛的算子且运行速度更快。即将发布的一篇博文将详细阐述 nvFuser 并展示它如何加速各种网络上的训练。

有关使用和调试的更多详细信息，请参阅 nvFuser 文档。

Ampere 及更新 CUDA 硬件上 float32 矩阵乘法精度的变更

PyTorch 支持多种“混合精度”技术，如 torch.amp (自动混合精度) 模块，以及在 Ampere 及更新的 CUDA 硬件上使用 TensorFloat32 数据类型进行 float32 矩阵乘法，以加速内部计算。在 PyTorch 1.12 中，我们将 float32 矩阵乘法的默认行为更改为始终使用完整的 IEEE fp32 精度，这比用于内部计算的 TensorFloat32 数据类型更精确但速度较慢。对于像 A100 这样 TensorFloat32 与 float32 吞吐量比率特别高的设备，默认值的更改可能会导致明显的性能下降。

如果您一直在使用 TensorFloat32 矩阵乘法，可以通过设置 torch.backends.cuda.matmul.allow_tf32 = True 继续使用，

该设置自 PyTorch 1.7 起即受支持。从 PyTorch 1.12 开始，还可以使用新的矩阵乘法精度 API：torch.set_float32_matmul_precision(“highest”|”high”|”medium”)

重申一下，PyTorch 对所有设备类型的新默认值均为“最高 (highest)”精度。我们认为这在不同设备类型之间为矩阵乘法提供了更好的一致性。新精度 API 的文档可以在此处找到。设置“高 (high)”或“中 (medium)”精度类型将在 Ampere 及更新的 CUDA 设备上启用 TensorFloat32。如果您正在升级到 PyTorch 1.12，为了保持矩阵乘法在 Ampere 设备上当前的性能，请将精度设置为“高 (high)”。

使用混合精度技术对于高效训练现代深度学习网络至关重要。如果您已经在用 torch.amp，那么这一更改不太可能对您产生影响。如果您不熟悉混合精度训练，请查看我们即将发布的博文《PyTorch 用户关于混合精度训练需要知道的一切》。

（测试版）CPU 上 Channels Last 格式的 PyTorch 视觉模型加速

内存格式对视觉模型运行性能有显著影响。通常，由于数据局部性更好，Channels Last 在性能方面更具优势。1.12 版本引入了内存格式的基本概念，并展示了在 Intel® Xeon® 可扩展处理器上流行的 PyTorch 视觉模型使用 Channels Last 所带来的性能优势。

支持 CPU 上 CV 领域常用算子的 Channels Last 内存格式，适用于推理和训练
针对 ATen 中 Channels Last 内核提供原生级优化，适用于 AVX2 和 AVX512
在 Intel® Xeon® Ice Lake（或更新）CPU 上，相较于 Channels First，TorchVision 模型的推理性能提升了 1.3 倍至 1.8 倍

（测试版）通过 Bfloat16 为 Intel® Xeon® 可扩展处理器上的 PyTorch 赋能

像 bfloat16 这样的低精度数值格式可提高多种深度学习训练工作负载的 PyTorch 性能。PyTorch 1.12 包含了 bfloat16 的最新软件增强功能，适用于更广泛的用户场景，并展示了更高的性能增益。主要改进包括：

使用 Intel® Xeon® Cooper Lake CPU 上引入的新 bfloat16 原生指令 VDPBF16PS，硬件计算吞吐量相比 float32 提高 2 倍
内存占用仅为 float32 的一半，提升了内存带宽密集型算子的速度
在 Intel® Xeon® Cooper Lake（或更新）CPU 上，TorchVision 模型的推理性能相较于 float32 提升 1.4 倍至 2.2 倍

（原型）引入 Mac 上的 PyTorch 加速训练

随着 PyTorch 1.12 的发布，开发者和研究人员现在可以利用 Apple Silicon GPU 实现显著更快的模型训练。这解锁了在 Mac 上本地进行原型设计和微调等机器学习工作流程的能力。加速的 GPU 训练使用 Apple 的 Metal Performance Shaders (MPS) 作为后端。其优势包括通过加速 GPU 训练带来的性能提升，以及能够在本地训练更大的网络或设置更大的批量大小。点击此处了解更多。

加速 GPU 训练和评估速度比仅 CPU 快（倍数）

除了新的 MPS 设备支持外，过去几个版本中提供的核心和领域库的 M1 二进制文件现已成为正式的原型功能。这些二进制文件可用于在 Apple Silicon 上原生运行 PyTorch。

（原型）BetterTransformer：Transformer 编码器推理的快速路径执行

PyTorch 现在支持 CPU 和 GPU 的快速路径实现（“BetterTransformer”），适用于包括 TransformerEncoder、TransformerEncoderLayer 和 MultiHeadAttention (MHA) 在内的多个 Transformer 编码器模块。BetterTransformer 快速路径架构在许多常见执行场景下性能持续提升 2 倍，具体取决于模型和输入特征。启用 BetterTransformer 的新模块与之前版本的 PyTorch Transformer API 保持 API 兼容，如果满足快速路径执行要求，它们将加速现有模型，并可读取用之前版本 PyTorch 训练的模型。PyTorch 1.12 包括：

为 Torchtext 的预训练 RoBERTa 和 XLM-R 模型提供 BetterTransformer 集成
基于 PyTorch Transformer API 构建的 Torchtext
通过使用将多个算子合并为单个内核的融合内核减少执行开销，从而实现快速路径执行以提升性能
在自然语言处理中处理填充标记时，利用数据稀疏性获得额外加速的选项（在创建 TransformerEncoder 时设置 enable_nested_tensor=True）
诊断工具，帮助用户了解为何未触发快速路径执行

分布式

（测试版）完全分片数据并行 (FSDP) API

FSDP API 有助于通过在数据并行工作节点之间对模型的参数、梯度和优化器状态进行分片，轻松扩展大模型训练，同时保持数据并行性的简洁性。原型版本于 PyTorch 1.11 发布，其最小功能集助力了 1 万亿参数模型规模的测试。

在此测试版中，FSDP API 增加了以下功能以支持各种生产工作负载。此测试版中新功能的亮点包括：

通用分片策略 API – 用户只需修改一行代码即可轻松切换分片策略，从而比较并使用 DDP（仅数据分片）、FSDP（全模型和数据分片）或 Zero2（仅优化器和梯度分片），以优化特定训练需求的内存和性能
细粒度混合精度策略 – 用户可以通过混合精度策略为模型参数、梯度通信和缓冲区指定半精度和全精度数据类型（bfloat16、fp16 或 fp32）的组合。模型会自动以 fp32 保存以实现最大程度的移植性
Transformer 自动包装策略 – 通过注册模型层类，允许对基于 Transformer 的模型进行最优包装，从而加速训练性能
使用 device_id 初始化实现更快的模型初始化 – 以流式方式执行初始化，以避免 OOM 问题并优化初始化性能（相对于 CPU 初始化）
针对大模型全模型保存的 Rank0 流式传输 – 可以由所有 GPU 将其分片流式传输到 Rank 0 GPU 来保存完全分片的模型，并在 Rank 0 CPU 上构建完整状态的模型以进行保存

更多详情和代码示例，请查看文档和教程。

感谢阅读，如果您对这些更新感兴趣并希望加入 PyTorch 社区，我们鼓励您加入讨论论坛并提交 GitHub issue。要获取 PyTorch 的最新消息，请在 Twitter、Medium、YouTube 和 LinkedIn 上关注我们。

干杯！

PyTorch 团队

PyTorch 1.12：TorchArrow、模块的函数式 API 和 nvFuser 现已可用