我们很高兴看到 PyTorch 2.1 的发布。在这篇博客中，我们将讨论英特尔对 PyTorch 2.1 做出重大贡献的五项特性。
- TorchInductor-CPU 优化，包括针对 torch.compile 的 Bfloat16 推理路径
- 针对 torch.compile 的 CPU 动态形状推理路径
- C++ 封装器（原型）
- 基于 Flash-attention 的 CPU 缩放点积算法
- 通过 Inductor 导出 PyTorch 2 训练后量化，采用 x86 后端
在英特尔，我们很高兴能成为 PyTorch 社区的一员，并感谢与 Meta* 同事在共同开发这些特性方面的合作和反馈。
让我们开始吧。
TorchInductor-CPU 优化
此特性优化了 TorchInductor 的 bfloat16 推理性能。第三代和第四代英特尔® 至强® 可扩展处理器内置硬件加速器，可加速 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 几何平均加速比（单插槽多线程） | |||
| 编译器 | TorchBench | HuggingFace | TIMM_Models |
| Inductor | 1.81 倍 | 1.25x | 2.35 倍 |
| Bfloat16 几何平均加速比（单核单线程） | |||
| 编译器 | TorchBench | HuggingFace | TIMM_Models |
| Inductor | 1.74 倍 | 1.28 倍 | 1.29 倍 |
开发人员可以将其模型完全部署在第四代英特尔至强处理器上，以利用英特尔® 高级矩阵扩展（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 所示。我们看到图模式的性能比即时模式高出 1.15 倍到 1.79 倍不等。
表 2. 动态形状几何平均加速比，与即时模式相比
| 动态形状几何平均加速比（单插槽多线程） | |||
| 编译器 | TorchBench | HuggingFace | TIMM_Models |
| Inductor | 1.35 倍 | 1.15x | 1.79 倍 |
| 动态形状几何平均加速比（单核单线程） | |||
| 编译器 | TorchBench | HuggingFace | TIMM_Models |
| Inductor | 1.48 倍 | 1.15x | 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) |
| Inductor | 1.06 倍 | 1.01 倍 | 1.00 倍 |
| FP32 静态形状模式几何平均加速比（单核单线程） | |||
| 编译器 | 小型 (t <= 0.04s) | 中型 (0.04s < t <= 1.5s) | 大型 (t > 1.5s) |
| Inductor | 1.13 倍 | 1.02 倍 | 1.01 倍 |
| FP32 动态形状模式几何平均加速比（单插槽多线程） | |||
| 编译器 | 小型 (t <= 0.04s) | 中型 (0.04s < t <= 1.5s) | 大型 (t > 1.5s) |
| Inductor | 1.05 倍 | 1.01 倍 | 1.00 倍 |
| FP32 动态形状模式几何平均加速比（单核单线程） | |||
| 编译器 | 小型 (t <= 0.04s) | 中型 (0.04s < t <= 1.5s) | 大型 (t > 1.5s) |
| Inductor | 1.14 倍 | 1.02 倍 | 1.01 倍 |
| BF16 静态形状模式几何平均加速比（单插槽多线程） | |||
| 编译器 | 小型 (t <= 0.04s) | 中型 (0.04s < t <= 1.5s) | 大型 (t > 1.5s) |
| Inductor | 1.09 倍 | 1.03 倍 | 1.04 倍 |
| BF16 静态形状模式几何平均加速比（单核单线程） | |||
| 编译器 | 小型 (t <= 0.04s) | 中型 (0.04s < t <= 1.5s) | 大型 (t > 1.5s) |
| Inductor | 1.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 几何加速比 |
| Inductor | 1.15 倍，20/20 | 1.07 倍，20/20 |
| SDPA 几何平均加速比（单核单线程） | ||
| 编译器 | FP32 几何加速比 | BF16 几何加速比 |
| Inductor | 1.02 倍，20/20 | 1.04 倍，20/20 |
通过 Inductor 导出 PyTorch 2 训练后量化，采用 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 导出训练后量化后端的代码片段
与 Inductor FP32 推理路径相比，TorchBench 测试套件中的所有卷积神经网络 (CNN) 模型都已测量并证明有效。性能指标如表 5 所示。
| 编译器 | 几何加速比 | 几何相关精度损失 |
| Inductor | 3.25 倍，12/12 | 0.44%, 12/12 |
下一步
获取软件
试用 PyTorch 2.1，并亲自体验英特尔贡献的这些特性带来的性能优势。
我们鼓励您查看英特尔的其他 AI 工具 和 框架 优化，并了解开放的、基于标准的 oneAPI 多架构、多供应商编程模型，它构成了英特尔 AI 软件组合的基础。
有关第四代英特尔至强可扩展处理器的更多详细信息，请访问 AI 平台，您可以在其中了解英特尔如何赋能开发人员运行高性能、高效的端到端 AI 流水线。
PyTorch 资源
产品和性能信息
1 Amazon EC2* m7i.16xlarge：1 节点，英特尔至强铂金 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 2017 年 10 月 16 日；软件：PyTorch 2.1.0_rc4，英特尔® oneAPI 深度神经网络库 (oneDNN) 3.1.1 版，TorchBench，TorchVision，TorchText，TorchAudio，TorchData，TorchDynamo 基准测试，英特尔于 2023 年 9 月 12 日测试。
2 Amazon EC2 c6i.16xlarge：1 节点，英特尔至强铂金 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.0 2017 年 10 月 16 日；软件：PyTorch 2.1.0_rc4，oneDNN 3.1.1 版，TorchBench，TorchVision，TorchText，TorchAudio，TorchData，TorchDynamo 基准测试，TorchBench cpu 用户基准测试，英特尔于 2023 年 9 月 12 日测试。