PyTorch 2 导出量化与 X86 后端通过 Inductor¶
作者: Leslie Fang,Weiwen Xia,Jiong Gong,Jerry Zhang
先决条件¶
简介¶
本教程介绍了使用 PyTorch 2 导出量化流程生成针对 x86 inductor 后端定制的量化模型的步骤,并解释了如何将量化模型降低到 inductor 中。
pytorch 2 导出量化流程使用 torch.export 将模型捕获到图中,并在 ATen 图之上执行量化变换。此方法预计将具有显著更高的模型覆盖率、更好的可编程性和简化的 UX。TorchInductor 是新的编译器后端,它将 TorchDynamo 生成的 FX 图编译成优化的 C++/Triton 内核。
此量化 2 与 Inductor 的流程支持静态和动态量化。静态量化最适合 CNN 模型,例如 ResNet-50。动态量化更适合 NLP 模型,例如 RNN 和 BERT。有关两种量化类型之间的差异,请参阅以下页面。
量化流程主要包括三个步骤
步骤 1:根据torch 导出机制从 eager 模型中捕获 FX 图。
步骤 2:根据捕获的 FX 图应用量化流程,包括定义后端特定量化器、生成带有观察器的准备模型、执行准备模型的校准或量化感知训练,并将准备模型转换为量化模型。
步骤 3:使用 API
torch.compile
将量化模型降低到归纳器。
此流程的高级架构可能如下所示
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 相结合,我们在新的量化前端具有灵活性和生产力,并在编译器后端具有出色的开箱即用性能。尤其是在英特尔第四代 (SPR) 至强处理器上,通过利用 高级矩阵扩展 功能可以进一步提升模型的性能。
训练后量化¶
现在,我们将逐步指导您如何使用 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 位。
当向量神经网络指令不可用时,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 包装器生成 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 量化开始发挥作用。在这种情况下,卷积或 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 量化中将 Inductor 与 X86 CPU 结合使用。用户可以了解如何使用 X86InductorQuantizer
量化模型并将其降低到具有 X86 CPU 设备的归纳器。