概览
INT8 量化是一种强大的技术，可用于加速 x86 CPU 平台上的深度学习推理。通过将模型权重和激活的精度从 32 位浮点数 (FP32) 降低到 8 位整数 (INT8)，INT8 量化可以在不牺牲准确性的情况下显著提高推理速度并减少内存需求。
在这篇博客中，我们将讨论 PyTorch 中面向 x86 CPU 的 INT8 量化最新进展，重点介绍新的 x86 量化后端。我们还将简要介绍 PyTorch 2.0 Export (PT2E) 和 TorchInductor 的新量化路径。
X86 量化后端
PyTorch 中当前推荐的量化方式是 FX。在 PyTorch 2.0 之前，x86 CPU 上的默认量化后端（也称为 QEngine）是 FBGEMM，它利用 FBGEMM 性能库实现性能加速。在 PyTorch 2.0 版本中，引入了一个新的量化后端 X86 来取代 FBGEMM。通过利用 FBGEMM 和 Intel® oneAPI 深度神经网络库 (oneDNN) 内核库的优势，x86 量化后端与原始 FBGEMM 后端相比，提供了改进的 INT8 推理性能。
X86 后端的性能优势
为了衡量新 X86 后端的性能优势，我们使用 第四代 Intel® Xeon® 可扩展处理器对 69 个流行的深度学习模型（如下所示的图 1-3）进行了 INT8 推理。结果显示，与 FP32 推理性能相比，几何平均性能提升了 2.97 倍，而使用 FBGEMM 后端时性能提升为 1.43 倍。下面的图表显示了按模型比较 x86 后端和 FBGEMM 后端的性能提升。
图 1：使用 x86 后端性能提升小于 2 倍的模型1
图 2：使用 x86 后端性能提升 2-4 倍的模型1
图 3：使用 x86 后端性能提升大于 4 倍的模型1
X86 后端的使用
默认情况下，在 2.0 版本中，x86 平台上的用户将使用 x86 量化后端，并且在使用默认后端时，其 PyTorch 程序将保持不变。此外，用户也可以显式指定 x86 作为量化后端。
下面是使用 x86 量化后端进行 PyTorch 静态训练后量化的代码示例片段。
import torch
from torch.ao.quantization import get_default_qconfig_mapping
from torch.quantization.quantize_fx import prepare_fx, convert_fx
qconfig_mapping = get_default_qconfig_mapping()
# Or explicity specify the qengine
# qengine = 'x86'
# torch.backends.quantized.engine = qengine
# qconfig_mapping = get_default_qconfig_mapping(qengine)
model_fp32 = MyModel().eval()
x = torch.randn((1, 3, 224, 224), dtype=torch.float)
x = x.to(memory_format=torch.channels_last)
# Insert observers according to qconfig and backend config
prepared_model = prepare_fx(model_fp32, qconfig_mapping, example_inputs=x)
# Calibration code not shown
# Convert to quantized model
quantized_model = convert_fx(prepared_model)
X86 后端的技术细节
根据我们基准测试模型的性能数据，我们设计了启发式分派规则，以决定调用 oneDNN 还是 FBGEMM 性能库来执行卷积或矩阵乘法运算。这些规则是操作类型、形状、CPU 架构信息等的组合。详细逻辑可在此处获取。如需更多设计和技术讨论，请参阅征求意见稿。
PyTorch 2.0 Export 新量化路径的下一步
虽然距离最终确定尚远，但一个新的量化路径，PyTorch 2.0 Export (PT2E)，正处于早期设计和概念验证 (PoC) 阶段。这种新方法未来将取代 FX 量化路径。它建立在 TorchDynamo Export 的能力之上，TorchDynamo Export 是 PyTorch 2.0 版本中引入的用于捕获 FX 图的特性。然后将此图进行量化并下放到不同的后端。TorchInductor 作为 PyTorch 的新 DL 编译器，在 x86 CPU 上的 FP32 推理速度提升方面已显示出令人鼓舞的结果。我们正在积极努力将其启用为 PT2E 的量化后端之一。我们相信，由于在不同级别融合的灵活性更高，新路径将进一步改善 INT8 推理性能。
结论
PyTorch 2.0 版本中引入的 x86 后端在 x86 CPU 平台上的 INT8 推理速度方面表现出显著提升。与原始 FBGEMM 后端相比，它提供了 1.43 倍的速度提升，同时保持了向后兼容性。这一增强功能可以使终端用户受益，且只需极少或无需修改其程序。此外，一个新的量化路径 PT2E 正在开发中，预计未来将提供更多可能性。
致谢
特别感谢 Nikita Shulga、Vasiliy Kuznetsov、Supriya Rao 和 Jongsoo Park。我们共同在改进 PyTorch CPU 生态系统的道路上又向前迈出了一步。
配置
1 AWS EC2 r7iz.metal-16xl 实例 (Intel(R) Xeon(R) Gold 6455B, 32 核/64 线程, 开启睿频加速, 开启超线程, 内存: 8x64GB, 存储: 192GB); 操作系统: Ubuntu 22.04.1 LTS; 内核: 5.15.0-1028-aws; 批处理大小: 1; 每实例核心数: 4; PyTorch 2.0 RC3; TorchVision 0.15.0+cpu, Intel 于 2023 年 3 月 7 日测试。可能未反映所有公开可用的安全更新。
