跳转到主要内容
博客

PyTorch 2.1 包含适用于 AI 开发人员的新性能功能

作者： 2023 年 11 月 29 日2024 年 11 月 14 日暂无评论

我们很高兴看到 PyTorch 2.1 的发布。在本博客中，我们将讨论英特尔为 PyTorch 2.1 做出重大贡献的五项功能。

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

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

让我们开始吧。

TorchInductor-CPU 优化

此功能优化了 TorchInductor 的 bfloat16 推理性能。第三代和第四代 Intel® Xeon® 可扩展处理器内置了硬件加速器，可加快 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）的性能优于急切模式 1.25 倍到 2.35 倍不等。

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

Bfloat16 几何平均加速比（单插槽多线程）
编译器TorchBenchHugging FaceTIMM 模型
Inductor1.81 倍1.25 倍2.35 倍
Bfloat16 几何平均加速比（单核单线程）
编译器TorchBenchHugging FaceTIMM 模型
Inductor1.74 倍1.28 倍1.29 倍

开发人员可以将其模型完全部署在第四代 Intel Xeon 处理器上，以利用 Intel® Advanced Matrix Extensions (Intel® AMX) 功能，从而获得 torch.compile 的峰值性能。Intel AMX 包含两个主要组件：磁贴和磁贴矩阵乘法 (TMUL)。磁贴将大量数据存储在八个二维寄存器中，每个寄存器的大小为一千字节。TMUL 是一个连接到磁贴的加速器引擎，其中包含用于在单个操作中计算更大矩阵的指令。

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

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

动态形状受 Inductor Triton 后端（用于 CUDA*）和 C++ 后端（用于 CPU）支持。范围涵盖了功能（通过模型通过率衡量）和性能（通过推理延迟/吞吐量衡量）的改进。图 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. 动态形状几何平均加速比与急切模式相比

动态形状几何平均加速比（单插槽多线程）
编译器TorchBenchHugging FaceTIMM 模型
Inductor1.35 倍1.15 倍1.79 倍
动态形状几何平均加速比（单核单线程）
编译器TorchBenchHugging FaceTIMM 模型
Inductor1.48 倍1.15 倍1.48 倍

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 显示了 C++ 包装器相对于默认 Python 包装器实现的几何平均加速比。

表 3. C++ 包装器几何平均加速比与急切模式相比

FP32 静态形状模式几何平均加速比（单插槽多线程）
编译器小（t <= 0.04s）中（0.04s < t <= 1.5s）大（t > 1.5s）
Inductor1.06 倍1.01 倍1.00 倍
FP32 静态形状模式几何平均加速比（单核单线程）
编译器小（t <= 0.04s）中（0.04s < t <= 1.5s）大（t > 1.5s）
Inductor1.13 倍1.02 倍1.01 倍
FP32 动态形状模式几何平均加速比（单插槽多线程）
编译器小（t <= 0.04s）中（0.04s < t <= 1.5s）大（t > 1.5s）
Inductor1.05 倍1.01 倍1.00 倍
FP32 动态形状模式几何平均加速比（单核单线程）
编译器小（t <= 0.04s）中（0.04s < t <= 1.5s）大（t > 1.5s）
Inductor1.14 倍1.02 倍1.01 倍
BF16 静态形状模式几何平均加速比（单插槽多线程）
编译器小（t <= 0.04s）中（0.04s < t <= 1.5s）大（t > 1.5s）
Inductor1.09 倍1.03 倍1.04 倍
BF16 静态形状模式几何平均加速比（单核单线程）
编译器小（t <= 0.04s）中（0.04s < t <= 1.5s）大（t > 1.5s）
Inductor1.17 倍1.04 倍1.03 倍

基于 Flash-Attention 的 CPU 标量点积算法

标量点积注意力 (SDPA) 是 PyTorch 2.0 的标志性功能之一，有助于加速 Transformer 模型。它通过优化的 CUDA 内核进行加速，但仍缺乏优化的 CPU 内核。此 Flash-attention 实现针对训练和推理，支持 FP32 和 Bfloat16 数据类型。用户无需前端更改即可利用此 SDPA 优化。调用 SDPA 时，将自动选择特定的实现，包括此新实现。

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

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

SDPA 几何平均加速比（单插槽多线程）
编译器几何加速比 FP32几何加速比 BF16
Inductor1.15 倍，20/201.07 倍，20/20
SDPA 几何平均加速比（单核单线程）
编译器几何加速比 FP32几何加速比 BF16
Inductor1.02 倍，20/201.04 倍，20/20

通过 Inductor 使用 x86 后端进行 PyTorch 2 导出后训练量化

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 所示。

编译器几何加速比几何相关精度损失
Inductor3.25 倍，12/120.44%, 12/12

后续步骤

获取软件

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

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

有关第四代 Intel Xeon 可扩展处理器的更多详细信息，请访问 AI 平台，您可以在其中了解英特尔如何助力开发人员运行高性能、高效的端到端 AI 管道。

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 基准测试，英特尔于 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 基准测试，TorchBench CPU 用户基准测试，英特尔于 2023 年 9 月 12 日测试。