PyTorch 2 透過 Inductor 使用 Intel GPU 後端進行匯出量化¶
作者: Yan Zhiwei, Wang Eikan, Zhang Liangang, Liu River, Cui Yifeng
先決條件¶
PyTorch 2.7 或更高版本
介紹¶
本教程介紹了 XPUInductorQuantizer，它旨在為 Intel GPU 上的量化模型提供推理服務。 XPUInductorQuantizer 使用 PyTorch 匯出量化流程，並將量化模型降低到 inductor。
Pytorch 2 匯出量化流程使用 torch.export 將模型捕獲到圖中，並在 ATen 圖之上執行量化轉換。該方法有望實現更高的模型覆蓋率、更好的可程式設計性以及簡化的使用者體驗。TorchInductor 是一個編譯器後端，它將 TorchDynamo 生成的 FX 圖轉換為最佳化的 C++/Triton 核心。
量化流程包含三個步驟
步驟 1：基於 torch 匯出機制，從 eager 模型中捕獲 FX 圖。
步驟 2：基於捕獲的 FX 圖應用量化流程，包括定義後端特定的量化器、生成帶有觀察器的準備模型、執行準備模型的校準，以及將準備模型轉換為量化模型。
步驟 3：使用 API
torch.compile將量化模型降低到 inductor，這將呼叫 Triton 核心或 oneDNN GEMM/Convolution 核心。
這個流程的高階架構可能如下所示
float_model(Python) Example Input
\ /
\ /
—--------------------------------------------------------
| export |
—--------------------------------------------------------
|
FX Graph in ATen
| X86InductorQuantizer
| /
—--------------------------------------------------------
| prepare_pt2e |
| | |
| Calibrate/Train |
| | |
| convert_pt2e |
—--------------------------------------------------------
|
Quantized Model
|
—--------------------------------------------------------
| Lower into Inductor |
—--------------------------------------------------------
|
OneDNN kernels Triton Kernels
訓練後量化¶
靜態量化是我們目前唯一支援的方法。
建議透過 Intel GPU 渠道安裝以下依賴項
pip3 install torch torchvision torchaudio pytorch-triton-xpu --index-url https://download.pytorch.org/whl/xpu
請注意，由於 inductor 的 freeze 功能尚未預設開啟，您必須使用 TORCHINDUCTOR_FREEZING=1 來執行您的示例程式碼。
例如
TORCHINDUCTOR_FREEZING=1 python xpu_inductor_quantizer_example.py
1. 捕獲 FX 圖¶
我們將首先執行必要的匯入，從即時模式（eager）模組中捕獲 FX 圖。
import torch
import torchvision.models as models
from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e, convert_pt2e
import torchao.quantization.pt2e.quantizer.xpu_inductor_quantizer as xpuiq
from torchao.quantization.pt2e.quantizer.xpu_inductor_quantizer import XPUInductorQuantizer
from torch.export import export
# Create the Eager Model
model_name = "resnet18"
model = models.__dict__[model_name](weights=models.ResNet18_Weights.DEFAULT)
# Set the model to eval mode
model = model.eval().to("xpu")
# Create the data, using the dummy data here as an example
traced_bs = 50
x = torch.randn(traced_bs, 3, 224, 224, device="xpu").contiguous(memory_format=torch.channels_last)
example_inputs = (x,)
# Capture the FX Graph to be quantized
with torch.no_grad():
exported_model = export(
model,
example_inputs,
).module()
接下來，我們將量化 FX Module。
2. 應用量化¶
捕獲 FX Module 後，我們將匯入 Intel GPU 的後端量化器並進行配置以量化模型。
quantizer = XPUInductorQuantizer()
quantizer.set_global(xpuiq.get_default_xpu_inductor_quantization_config())
XPUInductorQuantizer 的預設量化配置對啟用值和權重都使用有符號 8 位。張量是逐張量量化的，而權重是逐通道有符號 8 位量化的。
可選地，除了使用非對稱量化啟用值的預設量化配置外，還支援有符號 8 位對稱量化啟用值，這有可能提供更好的效能。
from torchao.quantization.pt2e.observer import HistogramObserver, PerChannelMinMaxObserver
from torchao.quantization.pt2e.quantizer.quantizer import QuantizationSpec
from torchao.quantization.pt2e.quantizer import QuantizationConfig
from typing import Any, Optional, TYPE_CHECKING
if TYPE_CHECKING:
from torchao.quantization.pt2e import ObserverOrFakeQuantizeConstructor
def get_xpu_inductor_symm_quantization_config():
extra_args: dict[str, Any] = {"eps": 2**-12}
act_observer_or_fake_quant_ctr = HistogramObserver
act_quantization_spec = QuantizationSpec(
dtype=torch.int8,
quant_min=-128,
quant_max=127,
qscheme=torch.per_tensor_symmetric, # Change the activation quant config to symmetric
is_dynamic=False,
observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args(
**extra_args
),
)
weight_observer_or_fake_quant_ctr: ObserverOrFakeQuantizeConstructor = (
PerChannelMinMaxObserver
)
weight_quantization_spec = QuantizationSpec(
dtype=torch.int8,
quant_min=-128,
quant_max=127,
qscheme=torch.per_channel_symmetric, # Same as the default config, the only supported option for weight
ch_axis=0, # 0 corresponding to weight shape = (oc, ic, kh, kw) of conv
is_dynamic=False,
observer_or_fake_quant_ctr=weight_observer_or_fake_quant_ctr.with_args(
**extra_args
),
)
bias_quantization_spec = None # will use placeholder observer by default
quantization_config = QuantizationConfig(
act_quantization_spec,
act_quantization_spec,
weight_quantization_spec,
bias_quantization_spec,
False,
)
return quantization_config
# Then, set the quantization configuration to the quantizer.
quantizer = XPUInductorQuantizer()
quantizer.set_global(get_xpu_inductor_symm_quantization_config())
匯入後端特定的量化器後，為訓練後量化準備模型。 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 後端。
with torch.no_grad():
optimized_model = torch.compile(converted_model)
# Running some benchmark
optimized_model(*example_inputs)
在更高階的場景中，int8-mixed-bf16 量化發揮作用。在這種情況下，卷積或 GEMM 運算元會以 BFloat16 而非 Float32 的形式輸出，前提是後面沒有量化節點。隨後，BFloat16 張量會無縫地傳播到後續的逐點運算元，從而有效地減少記憶體使用並可能提高效能。使用此功能與常規 BFloat16 Autocast 的使用方式相同，只需將指令碼包裝在 BFloat16 Autocast 上下文中即可。
with torch.amp.autocast(device_type="xpu", dtype=torch.bfloat16), torch.no_grad():
# Turn on Autocast to use int8-mixed-bf16 quantization. After lowering into indcutor 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)
結論¶
在本教程中，我們學習瞭如何利用 XPUInductorQuantizer，透過 PyTorch 2 的匯出量化流程，對用於 Intel GPU 推理的模型執行訓練後量化。我們涵蓋了捕獲 FX 圖、應用量化以及使用 torch.compile 將量化模型降低到 inductor 後端的步驟。此外，我們還探討了使用 int8-mixed-bf16 量化來提高記憶體效率和潛在效能優勢的好處，尤其是在使用 BFloat16 自動轉換時。
