透過 Inductor 使用 X86 後端的 PyTorch 2 匯出量化

建立於：2024 年 1 月 25 日 | 最後更新：2024 年 4 月 19 日 | 最後驗證：2024 年 11 月 05 日

作者：Leslie Fang、Weiwen Xia、Jiong Gong、Jerry Zhang

先決條件

• PyTorch 2 Export 後訓練量化

• PyTorch 2 Export 量化感知訓練

• PyTorch 中的 TorchInductor 和 torch.compile 概念

• Inductor C++ Wrapper 概念

簡介

本教學介紹使用 PyTorch 2 Export 量化流程產生針對 x86 inductor 後端自訂的量化模型，以及如何將量化模型降低到 inductor 中的步驟。

pytorch 2 export 量化流程使用 torch.export 將模型擷取到圖表中，並在 ATen 圖表上執行量化轉換。預計這種方法具有顯著更高的模型覆蓋率、更好的可程式性以及簡化的 UX。TorchInductor 是一種新的編譯器後端，可將 TorchDynamo 產生的 FX Graphs 編譯為最佳化的 C++/Triton 核心。

具有 Inductor 的量化 2 流程支援靜態和動態量化。靜態量化最適合 CNN 模型，例如 ResNet-50。動態量化更適合 NLP 模型，例如 RNN 和 BERT。有關兩種量化類型之間的差異，請參閱以下頁面。

量化流程主要包括三個步驟

• 步驟 1：根據 torch export 機制從 eager 模型中擷取 FX Graph。

• 步驟 2：根據擷取的 FX Graph 應用量化流程，包括定義後端特定的量化器、產生具有觀察器的準備模型、執行準備模型的校準或量化感知訓練，以及將準備模型轉換為量化模型。

• 步驟 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 中的量化，我們可以使用新的量化前端獲得彈性和生產力，並使用編譯器後端獲得出色的開箱即用效能。尤其是在 Intel 第四代 (SPR) Xeon 處理器上，它可以透過利用進階矩陣擴充功能功能進一步提升模型的效能。

後訓練量化

現在，我們將逐步引導您如何將其與torchvision resnet18 模型一起用於後訓練量化。

1. 擷取 FX Graph

我們將從執行必要的匯入開始，從 eager 模組中擷取 FX Graph。

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-mixed-bf16 量化開始發揮作用。在這種情況下，如果沒有後續的量化節點，Convolution 或 GEMM 運算子會產生 BFloat16 輸出資料類型，而不是 Float32。隨後，BFloat16 張量會順暢地透過後續的 pointwise 運算子傳播，有效地最小化記憶體使用量，並可能提高效能。此功能的使用方式與常規 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 Export 量化感知訓練 (QAT)，隨後將量化模型降低到 Inductor。為了更深入地了解 PT2 Export 量化感知訓練，我們建議參考專門的PyTorch 2 Export 量化感知訓練。

PyTorch 2 Export 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 量化中使用 Inductor 與 X86 CPU。使用者可以學習如何使用 X86InductorQuantizer 量化模型，並使用 X86 CPU 裝置將其降低到 inductor。