PyTorch 2.1 包含面向 AI 开发人员的新性能特性

作者：Intel

我们很高兴看到 PyTorch 2.1 发布。在这篇博客中，我们将讨论 Intel 为 PyTorch 2.1 做出重要贡献的五大特性

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

在 Intel，我们很高兴能成为 PyTorch 社区的一份子，并感谢我们在共同开发这些特性时与 Meta* 同事的协作和反馈。

让我们开始吧。

TorchInductor-CPU 优化

此特性优化了 TorchInductor 的 bfloat16 推理性能。第三代和第四代 Intel® 至强® 可扩展处理器内置硬件加速器，可加快 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 性能在图模式下的几何平均加速比，与 eager 模式相比

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

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

动态形状是 PyTorch 2.0 的关键特性之一。PyTorch 2.0 默认假定所有内容都是静态的。如果由于大小改变而重新编译，我们将尝试将该大小重新编译为动态（已更改的大小未来也可能更改）。动态形状支持是大型语言模型 (LLM) 等流行模型所必需的。支持广泛模型的动态形状可以帮助用户从 torch.compile 中获得更多收益。对于动态形状，我们为 conv/gemm 算子提供后操作融合，并为非 conv/gemm 算子提供向量化代码生成。

CUDA\* 的 inductor 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 所示。在此，我们看到图模式的性能优于 eager 模式，加速因子范围从 1.15 倍到 1.79 倍。

表 2. 动态形状几何平均加速比，与 Eager 模式相比

C++ Wrapper（原型）

此特性生成 C++ 代码而不是 Python* 代码，用于调用 TorchInductor 中生成的内核和外部内核，以减少 Python 开销。它也是在没有 Python 的环境中支持部署的中间步骤。

要启用此特性，请使用以下配置：

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

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

表 3. C++ Wrapper 几何平均加速比，与 Eager 模式相比

基于 Flash-Attention 的 CPU 缩放点积算法

缩放点积注意力 (SDPA) 是 PyTorch 2.0 的旗舰特性之一，有助于加速 Transformer 模型。它通过最优化的 CUDA 内核加速，但仍缺乏优化的 CPU 内核。这种 Flash-Attention 实现的目标是训练和推理，支持 FP32 和 Bfloat16 数据类型。用户无需改变前端用法即可利用此 SDPA 优化。调用 SDPA 时，将自动选择特定实现，包括此新实现。

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

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

PyTorch 2 导出后训练量化，通过 Inductor 支持 x86 后端

PyTorch 在 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. 显示将 Inductor 用作 PyTorch 2 导出后训练量化后端的代码片段

TorchBench 测试套件中的所有卷积神经网络 (CNN) 模型都已测量并证明有效，与 Inductor FP32 推理路径相比。性能指标如表 5 所示。

下一步

获取软件

试用 PyTorch 2.1，亲自体验 Intel 贡献的这些特性带来的性能优势。

我们鼓励您查看 Intel 的其他 AI 工具 和 框架 优化，并了解基于开放标准的多架构、多供应商编程模型 oneAPI，它是 Intel AI 软件产品组合的基础。

有关第四代 Intel 至强可扩展处理器的更多详情，请访问 AI 平台，在此您可以了解 Intel 如何赋能开发者运行高性能、高效的端到端 AI 管线。

PyTorch 资源

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

产品与性能信息

1 Amazon EC2* m7i.16xlarge: 1 节点，配备 256 GB 内存（1 x 256 GB DDR5 4800 MT/s）的 Intel 至强铂金 8488C 处理器，微码 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，Intel 于 2023 年 9 月 12 日测试。

2 Amazon EC2 c6i.16xlarge: 1 节点，配备 128 GB 内存（1 x 128 GB DDR4 3200 MT/s）的 Intel 至强铂金 8375C 处理器，微码 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.0 10/16/2017；软件：PyTorch 2.1.0_rc4，oneDNN 版本 3.1.1，TorchBench，TorchVision，TorchText，TorchAudio，TorchData，TorchDynamo Benchmarks，TorchBench cpu userbenchmark，Intel 于 2023 年 9 月 12 日测试。