跳转到主要内容
博客

PyTorch 中 x86 CPU 的 INT8 量化

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

概述

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 后端相比，x86 量化后端通过利用 FBGEMM 和 Intel® oneAPI 深度神经网络库 (oneDNN) 内核库的优势，提供了改进的 INT8 推理性能。

X86 后端的性能优势

为了衡量新 X86 后端的性能优势，我们使用 第四代 Intel® Xeon® 可扩展处理器 对 69 种流行的深度学习模型（如以下 图 1-3 所示）进行了 INT8 推理。结果显示，与 FP32 推理性能相比，几何平均性能加速了 2.97 倍，而使用 FBGEMM 后端时，加速为 1.43 倍。下表显示了比较 x86 后端和 FBGEMM 后端的每个模型性能加速。

Figure 1: Models with less than 2x performance boost with x86 backend1

图 1：使用 x86 后端性能提升低于 2 倍的模型1

Figure 2: Models with 2x-4x performance boost with x86 backend1

图 2：使用 x86 后端性能提升 2 倍到 4 倍的模型1

Figure 3: Models with larger than 4x performance boost with x86 backend1

图 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) 这一新量化路径正处于早期设计和概念验证阶段。这种新方法计划在未来取代 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 线程, Turbo Boost 开启, 超线程开启, 内存: 8x64GB, 存储: 192GB); 操作系统: Ubuntu 22.04.1 LTS; 内核: 5.15.0-1028-aws; 批量大小: 1; 每个实例的核心数: 4; PyTorch 2.0 RC3; TorchVision 0.15.0+cpu, 由英特尔于 2023 年 3 月 77 日测试。可能未反映所有公开可用的安全更新。