PyTorch 2 通过 Inductor 和 X86 后端导出量化¶

创建于：2024 年 1 月 25 日 | 最后更新：2024 年 4 月 19 日 | 最后验证：2024 年 11 月 05 日

作者: Leslie Fang, Weiwen Xia, Jiong Gong, Jerry Zhang

先决条件¶

简介¶

本教程介绍了使用 PyTorch 2 导出量化流程生成为 x86 inductor 后端定制的量化模型的步骤，并解释了如何将量化模型降低到 inductor 中。

pytorch 2 导出量化流程使用 torch.export 将模型捕获到图中，并在 ATen 图之上执行量化转换。这种方法有望实现更高的模型覆盖率、更好的可编程性和简化的用户体验。 TorchInductor 是新的编译器后端，可将 TorchDynamo 生成的 FX 图编译为优化的 C++/Triton 内核。

Inductor 的量化 2 流程支持静态和动态量化。静态量化最适合 CNN 模型，如 ResNet-50。动态量化更适合 NLP 模型，如 RNN 和 BERT。有关两种量化类型之间的差异，请参阅以下页面。

量化流程主要包括三个步骤

步骤 1：基于 torch export 机制从 eager 模型捕获 FX 图。
步骤 2：应用基于捕获的 FX 图的量化流程，包括定义后端特定的量化器，生成带有观察者的准备好的模型，执行准备好的模型的校准或量化感知训练，并将准备好的模型转换为量化模型。
步骤 3：使用 API torch.compile 将量化模型降低到 inductor 中。

此流程的高级架构可能如下所示

float_model(Python)                          Example Input
    \                                              /
     \                                            /
—--------------------------------------------------------
|                         export                       |
—--------------------------------------------------------
                            |
                    FX Graph in ATen
                            |            X86InductorQuantizer
                            |                 /
—--------------------------------------------------------
|                      prepare_pt2e                     |
|                           |                           |
|                     Calibrate/Train                   |
|                           |                           |
|                      convert_pt2e                     |
—--------------------------------------------------------
                            |
                     Quantized Model
                            |
—--------------------------------------------------------
|                    Lower into Inductor                |
—--------------------------------------------------------
                            |
                         Inductor

结合 PyTorch 2 Export 和 TorchInductor 中的量化，我们通过新的量化前端获得了灵活性和生产力，并通过编译器后端获得了出色的开箱即用性能。特别是在英特尔第四代 (SPR) 至强处理器上，它们可以通过利用高级矩阵扩展功能进一步提升模型性能。

训练后量化¶

现在，我们将逐步引导您完成如何将此方法与torchvision resnet18 模型一起用于训练后量化的教程。

1. 捕获 FX 图¶

我们将从执行必要的导入开始，从 eager 模块捕获 FX 图。

import torch
import torchvision.models as models
import copy
from torch.ao.quantization.quantize_pt2e import prepare_pt2e, convert_pt2e
import torch.ao.quantization.quantizer.x86_inductor_quantizer as xiq
from torch.ao.quantization.quantizer.x86_inductor_quantizer import X86InductorQuantizer
from torch._export import capture_pre_autograd_graph

# Create the Eager Model
model_name = "resnet18"
model = models.__dict__[model_name](pretrained=True)

# Set the model to eval mode
model = model.eval()

# Create the data, using the dummy data here as an example
traced_bs = 50
x = torch.randn(traced_bs, 3, 224, 224).contiguous(memory_format=torch.channels_last)
example_inputs = (x,)

# Capture the FX Graph to be quantized
with torch.no_grad():
     # if you are using the PyTorch nightlies or building from source with the pytorch master,
    # use the API of `capture_pre_autograd_graph`
    # Note 1: `capture_pre_autograd_graph` is also a short-term API, it will be updated to use the official `torch.export` API when that is ready.
    exported_model = capture_pre_autograd_graph(
        model,
        example_inputs
    )
    # Note 2: if you are using the PyTorch 2.1 release binary or building from source with the PyTorch 2.1 release branch,
    # please use the API of `torch._dynamo.export` to capture the FX Graph.
    # exported_model, guards = torch._dynamo.export(
    #     model,
    #     *copy.deepcopy(example_inputs),
    #     aten_graph=True,
    # )

接下来，我们将拥有要量化的 FX 模块。

2. 应用量化¶

在我们捕获要量化的 FX 模块之后，我们将导入 X86 CPU 的后端量化器，并配置如何量化模型。

quantizer = X86InductorQuantizer()
quantizer.set_global(xiq.get_default_x86_inductor_quantization_config())

注意

X86InductorQuantizer 中的默认量化配置对激活和权重都使用 8 位。

当向量神经网络指令不可用时，oneDNN 后端会静默选择假定乘法为 7 位 x 8 位的内核。换句话说，在没有向量神经网络指令的 CPU 上运行时，可能会发生潜在的数值饱和和精度问题。

量化配置默认用于静态量化。要应用动态量化，在获取配置时添加参数 is_dynamic=True。

quantizer = X86InductorQuantizer()
quantizer.set_global(xiq.get_default_x86_inductor_quantization_config(is_dynamic=True))

在我们导入后端特定的量化器之后，我们将为训练后量化准备模型。 prepare_pt2e 将 BatchNorm 运算符折叠到前面的 Conv2d 运算符中，并在模型中的适当位置插入观察者。

prepared_model = prepare_pt2e(exported_model, quantizer)

现在，我们将在模型中插入观察者后校准 prepared_model。此步骤仅静态量化需要。

# We use the dummy data as an example here
prepared_model(*example_inputs)

# Alternatively: user can define the dataset to calibrate
# def calibrate(model, data_loader):
#     model.eval()
#     with torch.no_grad():
#         for image, target in data_loader:
#             model(image)
# calibrate(prepared_model, data_loader_test)  # run calibration on sample data

最后，我们将校准后的模型转换为量化模型。 convert_pt2e 接受校准后的模型并生成量化模型。

converted_model = convert_pt2e(prepared_model)

完成这些步骤后，我们完成了运行量化流程，并将获得量化模型。

3. 降低到 Inductor¶

在我们获得量化模型后，我们将进一步将其降低到 inductor 后端。默认的 Inductor 封装器生成 Python 代码以调用生成的内核和外部内核。此外，Inductor 支持生成纯 C++ 代码的 C++ 封装器。这允许生成的内核和外部内核的无缝集成，有效减少 Python 开销。在未来，利用 C++ 封装器，我们可以扩展功能以实现纯 C++ 部署。有关 C++ 封装器的一般全面详细信息，请参阅关于 Inductor C++ 封装器教程的专门教程。

# Optional: using the C++ wrapper instead of default Python wrapper
import torch._inductor.config as config
config.cpp_wrapper = True

with torch.no_grad():
    optimized_model = torch.compile(converted_model)

    # Running some benchmark
    optimized_model(*example_inputs)

在更高级的场景中，int8 混合 bf16 量化开始发挥作用。在这种情况下，Convolution 或 GEMM 运算符会生成 BFloat16 输出数据类型，而不是在没有后续量化节点的情况下生成 Float32。随后，BFloat16 张量无缝地通过后续的逐点运算符传播，从而有效地最大限度地减少内存使用并可能提高性能。此功能的利用类似于常规 BFloat16 Autocast，就像将脚本包装在 BFloat16 Autocast 上下文中一样简单。

with torch.autocast(device_type="cpu", dtype=torch.bfloat16, enabled=True), torch.no_grad():
    # Turn on Autocast to use int8-mixed-bf16 quantization. After lowering into Inductor CPP Backend,
    # For operators such as QConvolution and QLinear:
    # * The input data type is consistently defined as int8, attributable to the presence of a pair
        of quantization and dequantization nodes inserted at the input.
    # * The computation precision remains at int8.
    # * The output data type may vary, being either int8 or BFloat16, contingent on the presence
    #   of a pair of quantization and dequantization nodes at the output.
    # For non-quantizable pointwise operators, the data type will be inherited from the previous node,
    # potentially resulting in a data type of BFloat16 in this scenario.
    # For quantizable pointwise operators such as QMaxpool2D, it continues to operate with the int8
    # data type for both input and output.
    optimized_model = torch.compile(converted_model)

    # Running some benchmark
    optimized_model(*example_inputs)

将所有这些代码放在一起，我们将得到玩具示例代码。请注意，由于 Inductor freeze 功能默认情况下尚未启用，请使用 TORCHINDUCTOR_FREEZING=1 运行示例代码。

例如

TORCHINDUCTOR_FREEZING=1 python example_x86inductorquantizer_pytorch_2_1.py

在 PyTorch 2.1 版本中，TorchBench 测试套件中的所有 CNN 模型都已测量并证明与 Inductor FP32 推理路径相比有效。有关详细的基准测试数字，请参阅此文档。

量化感知训练¶

现在 X86 CPU 上支持使用 X86InductorQuantizer 的 PyTorch 2 导出量化感知训练 (QAT)，然后将量化模型后续降低到 Inductor 中。为了更深入地了解 PT2 导出量化感知训练，我们建议参考专门的 PyTorch 2 导出量化感知训练。

PyTorch 2 导出 QAT 流程与 PTQ 流程大致相似

import torch
from torch._export import capture_pre_autograd_graph
from torch.ao.quantization.quantize_pt2e import (
  prepare_qat_pt2e,
  convert_pt2e,
)
import torch.ao.quantization.quantizer.x86_inductor_quantizer as xiq
from torch.ao.quantization.quantizer.x86_inductor_quantizer import X86InductorQuantizer

class M(torch.nn.Module):
   def __init__(self):
      super().__init__()
      self.linear = torch.nn.Linear(1024, 1000)

   def forward(self, x):
      return self.linear(x)

example_inputs = (torch.randn(1, 1024),)
m = M()

# Step 1. program capture
# NOTE: this API will be updated to torch.export API in the future, but the captured
# result shoud mostly stay the same
exported_model = capture_pre_autograd_graph(m, example_inputs)
# we get a model with aten ops

# Step 2. quantization-aware training
# Use Backend Quantizer for X86 CPU
# To apply dynamic quantization, add an argument ``is_dynamic=True`` when getting the config.
quantizer = X86InductorQuantizer()
quantizer.set_global(xiq.get_default_x86_inductor_quantization_config(is_qat=True))
prepared_model = prepare_qat_pt2e(exported_model, quantizer)

# train omitted

converted_model = convert_pt2e(prepared_model)
# we have a model with aten ops doing integer computations when possible

# move the quantized model to eval mode, equivalent to `m.eval()`
torch.ao.quantization.move_exported_model_to_eval(converted_model)

# Lower the model into Inductor
with torch.no_grad():
  optimized_model = torch.compile(converted_model)
  _ = optimized_model(*example_inputs)

请注意，Inductor freeze 功能默认情况下未启用。要使用此功能，您需要使用 TORCHINDUCTOR_FREEZING=1 运行示例代码。

例如

TORCHINDUCTOR_FREEZING=1 python example_x86inductorquantizer_qat.py

结论¶

在本教程中，我们介绍了如何在 PyTorch 2 量化中使用带有 X86 CPU 的 Inductor。用户可以了解如何使用 X86InductorQuantizer 量化模型并使用 X86 CPU 设备将其降低到 inductor 中。