通过 Inductor 使用 X86 后端的 PyTorch 2 导出量化¶
Created On: Jan 25, 2024 | Last Updated: Apr 19, 2024 | Last Verified: Nov 05, 2024
作者: Leslie Fang, Weiwen Xia, Jiong Gong, Jerry Zhang
前置条件¶
简介¶
本教程介绍了利用 PyTorch 2 导出量化流程生成定制化的量化模型以适配 x86 Inductor 后端的步骤,并说明如何将量化模型降低到 Inductor。
PyTorch 2 导出量化流程使用 torch.export 捕获模型到图中,并在 ATen 图的基础上执行量化转换。这种方法预计具有显著更高的模型覆盖率、更好的可编程性,以及简化的用户体验。TorchInductor 是新的编译器后端,可将 TorchDynamo 生成的 FX 图编译为优化的 C++/Triton 内核。
Inductor 的这种量化流程支持静态和动态量化。静态量化对 CNN 模型(如 ResNet-50)效果最佳,动态量化更适合于 NLP 模型(如 RNN 和 BERT)。有关两种量化类型之间的区别,请参阅 相关页面。
量化流程主要包括三个步骤:
步骤1: 基于 torch 导出机制 从动态模型捕获 FX 图。
步骤 2:基于捕获的 FX 图应用量化流程,包括定义特定后端量化器、生成带观察者的准备模型、执行准备模型的校准或量化感知训练,以及将准备模型转换为量化模型。
步骤 3:使用
torch.compile
API 将量化模型降低到 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 导出的量化与 TorchInductor 相结合,我们通过新的量化前端获得灵活性和生产力,通过编译器后端获得突出的开箱即用性能。特别是在 Intel 第四代 (SPR) Xeon 处理器上,通过利用 高级矩阵扩展 功能可以进一步提升模型性能。
后训练量化¶
现在,我们将逐步向您介绍如何在 torchvision resnet18 模型 上执行后训练量化。
1. 捕获 FX 图¶
我们将从必要的导入开始,从动态模块中捕获 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 位。
当 Vector Neural Network 指令不可用时,oneDNN 后端会默默选择假设 乘法是 7 位与 8 位的内核。换句话说,在没有 Vector Neural Network 指令的 CPU 上运行时,可能会发生潜在的数值饱和和准确性问题。
默认情况下,量化配置是用于静态量化的。要应用动态量化,请在获取配置时添加参数 is_dynamic=True
。
quantizer = X86InductorQuantizer()
quantizer.set_global(xiq.get_default_x86_inductor_quantization_config(is_dynamic=True))
导入特定后端的 Quantizer 之后,我们将为后训练量化准备模型。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 量化派上用场。在这种情况下,当缺少后续量化节点时,卷积或 GEMM 操作会生成 BFloat16 输出数据类型而不是 Float32。随后,BFloat16 张量可以无缝地通过后续的点操作传播,有效地减少内存使用,并可能提高性能。使用此功能的方式与常规 BFloat16 自动转型相同,只需将脚本包裹在 BFloat16 自动转型上下文中即可。
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 推理路径相比效果良好。有关详细的基准指标,请参阅 本文档。
量化感知训练¶
PyTorch 2 导出量化感知训练 (QAT) 现已支持使用 X86InductorQuantizer 在 X86 CPU 上运行,随后将量化模型降低到 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。