PyTorch 2.1 包含针对 AI 开发者的新性能特性

由英特尔提供

我们非常期待 PyTorch 2.1 的发布。在这篇博客中，我们讨论了英特尔为 PyTorch 2.1 做出的五个重要贡献：

1. TorchInductor-CPU 优化，包括 torch.compile 的 Bfloat16 推理路径
2. torch.compile 的 CPU 动态形状推理路径
3. C++包装（原型）
4. 基于 Flash-attention 的 CPU 缩放点积算法
5. 通过电感器使用 x86 后端进行 PyTorch 2 导出后的训练后量化

在英特尔，我们很高兴成为 PyTorch 社区的一员，并感谢 Meta*的同事们在共同开发这些功能过程中的合作与反馈。

让我们开始吧。

TorchInductor-CPU 优化

此功能优化了 TorchInductor 的 bfloat16 推理性能。第 3 代和第 4 代英特尔® 至强® 可扩展处理器内置硬件加速器，用于加快 bfloat16 数据类型的点积计算。图 1 展示了如何指定 BF16 推理路径的代码片段。

user_model = ...

user_model.eval()
with torch.no_grad(), torch.autocast("cpu"):
	compiled_model = torch.compile(user_model)
	y = compiled_model(x)

图 1. 使用 TorchInductor 的 BF16 推理的代码片段

我们在三个 TorchInductor 基准测试套件——TorchBench、Hugging Face 和 TIMM——上进行了性能测试，结果如下表 1 所示。在这里，我们看到图模式（TorchInductor）的性能比 eager 模式高 1.25 倍到 2.35 倍。

表 1. 图模式下的 Bfloat16 性能几何平均加速，与急切模式相比

开发者可以将他们的模型完全部署在第四代英特尔至强处理器上，以充分利用英特尔® 高级矩阵扩展（英特尔® AMX）功能，以获得 torch.compile 的最高性能。英特尔 AMX 有两个主要组件：瓦片和瓦片矩阵乘法（TMUL）。瓦片将大量数据存储在八个二维寄存器中，每个寄存器大小为一千字节。TMUL 是一个连接到瓦片的加速引擎，其中包含用于单次操作计算较大矩阵的指令。

torch.compile 的 CPU 动态形状推理路径

动态形状是 PyTorch 2.0 中的关键特性之一。PyTorch 2.0 默认假设所有内容都是静态的。如果我们因为大小变化而重新编译，我们将尝试将那个大小作为动态的重新编译（大小变化的可能性很大，未来可能会发生变化）。动态形状支持对于大型语言模型（LLM）等流行模型是必需的。支持广泛模型范围的动态形状可以帮助用户从 torch.compile 中获得更多好处。对于动态形状，我们为 conv/gemm 运算符提供后操作融合，为非 conv/gemm 运算符提供向量化代码生成。

动态形状功能由 CUDA*后端的电感器 Triton 和 CPU 后端的 C++后端同时支持。范围包括对功能（通过模型通过率衡量）和性能（通过推理延迟/吞吐量衡量）的改进。图 2 显示了使用 TorchInductor 进行动态形状推理的代码片段。

user_model = ...

# Training example
compiled_model = torch.compile(user_model)
y = compiled_model(x_size1)
# Here trigger the recompile because the input size changed
y = compiled_model(x_size2)


# Inference example
user_model.eval()
compiled_model = torch.compile(user_model)
with torch.no_grad():
	y = compiled_model(x_size1)
 # Here trigger the recompile because the input size changed
 y = compiled_model(x_size2)

图 2. 显示使用 TorchInductor 进行动态形状推理的代码片段

我们再次在三个 TorchInductor 基准测试套件——TorchBench、Hugging Face 和 TIMM——上测量了性能，结果见表 2。在这里，我们看到图模式下的性能比急切模式高出 1.15 倍到 1.79 倍。

表 2. 与急切模式相比的动态形状几何平均速度提升

C++包装器（原型）

该功能生成 C++代码而不是 Python*代码来调用生成的内核和 TorchInductor 中的外部内核，以减少 Python 开销。这还是支持在无 Python 的环境中部署的中间步骤。

要启用此功能，请使用以下配置：

import torch
import torch._inductor.config as config
config.cpp_wrapper = True

对于轻量级工作负载，Python 包装器的开销更为明显，C++包装器则表现出更高的性能提升比率。我们根据单次迭代的平均推理时间，将 TorchBench、Hugging Face 和 TIMM 中的模型分组，并按小型、中型和大型分类。表 3 显示了与默认 Python 包装器相比，C++包装器实现的几何平均加速比。

表 3. 与 Eager 模式相比，C++包装器的几何平均加速比

基于 Flash-Attention 的缩放点积算法（SDPA）用于 CPU

缩放点积注意力（SDPA）是 PyTorch 2.0 的旗舰特性之一，有助于加速 Transformer 模型。它通过最优的 CUDA 内核进行加速，但仍然缺乏优化的 CPU 内核。这种闪速注意力实现针对训练和推理，支持 FP32 和 Bfloat16 数据类型。用户无需进行前端使用更改即可利用这种 SDPA 优化。调用 SDPA 时，将自动选择特定的实现，包括这种新的实现。

我们已经测量了 Hugging Face 中的 SDPA 相关模型，与未融合的 SDPA 相比，它们已被证明是有效的。表 4 展示了 SDPA 优化的几何平均加速比。

表 4. SDPA 优化性能的几何平均加速

PyTorch 2.0 导出后训练量化通过电感器使用 x86 后端

PyTorch 2.0 导出中提供了新的量化流程。该功能使用 TorchInductor，并以 x86 CPU 设备作为后端，实现了该新量化流程的后训练静态量化。示例代码片段如图 3 所示。

import torch
import torch._dynamo as torchdynamo
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
import torch.ao.quantization.quantizer.x86_inductor_quantizer as xiq

model = ... 

model.eval()
with torch.no_grad():
 # Step 1: Trace the model into an FX graph of flattened ATen operators
 exported_graph_module, guards = torchdynamo.export(
	 model,
	 *copy.deepcopy(example_inputs),
	 aten_graph=True,
 )

 # Step 2: Insert observers or fake quantize modules
 quantizer = xiq.X86InductorQuantizer()
 operator_config = xiq.get_default_x86_inductor_quantization_config()
 quantizer.set_global(operator_config)
 prepared_graph_module = prepare_pt2e(exported_graph_module, quantizer)

 # Doing calibration here.

 # Step 3: Quantize the model
 convert_graph_module = convert_pt2e(prepared_graph_module)

 # Step 4: Lower Quantized Model into the backend
 compile_model = torch.compile(convert_graph_module)

图 3. 展示使用电感器作为 PyTorch 2.0 导出后训练量化后端代码片段

所有来自 TorchBench 测试套件的卷积神经网络（CNN）模型在与 Inductor FP32 推理路径进行比较时均经过测量并证明有效。性能指标显示在表 5 中。

下一步

获取软件

尝试使用 PyTorch 2.1，并亲自体验由英特尔贡献的这些特性带来的性能优势。

我们鼓励您查看英特尔的其他 AI 工具和框架优化，并了解作为英特尔 AI 软件组合基础的开放、基于标准的 oneAPI 多架构、多厂商编程模型。

想了解更多关于第四代英特尔至强可扩展处理器的详细信息，请访问 AI 平台，了解英特尔如何赋能开发者运行高性能、高效的端到端 AI 管道。

PyTorch 资源

• PyTorch 入门
• 开发讨论
• 文档

产品和性能信息

1 Amazon EC2* m7i.16xlarge：1 节点，Intel Xeon Platinum 8488C 处理器，256 GB 内存（1 x 256 GB DDR5 4800 MT/s），微代码 0x2b000461，超线程开启，睿频开启，Ubuntu* 22.04.3 LTS，内核 6.2.0-1011-aws，GCC* 11.3.0，Amazon Elastic Block Store 200 GB，BIOS Amazon EC2 1.0 10/16/2017；软件：PyTorch 2.1.0_rc4，Intel® oneAPI 深度神经网络库（oneDNN）版本 3.1.1，TorchBench，TorchVision，TorchText，TorchAudio，TorchData，TorchDynamo Benchmarks，由英特尔于 2023 年 9 月 12 日测试。

2 Amazon EC2 c6i.16xlarge：1 节点，Intel Xeon Platinum 8375C 处理器，128 GB 内存（1 x 128 GB DDR4 3200 MT/s），微代码 0xd0003a5，超线程开启，睿频开启，Ubuntu 22.04.2 LTS，内核 6.2.0-1011-aws，gcc 11.3.0，Amazon Elastic Block Store 200 GB，BIOS Amazon EC2 1.010/16/2017；软件：PyTorch 2.1.0_rc4，oneDNN 版本 3.1.1，TorchBench，TorchVision，TorchText，TorchAudio，TorchData，TorchDynamo Benchmarks，TorchBench cpu userbenchmark，由英特尔于 2023 年 9 月 12 日测试。