如何为 PyTorch 2 export 量化编写 Quantizer

如何为 PyTorch 2 export 量化编写 Quantizer#

参考：pt2e_quantizer

（原型）PyTorch 2 Export 训练后量化引入了 PyTorch 2 Export 量化的整体 API，与 fx 图模式量化的主要区别在于 API 上明确表明量化是针对特定后端的。因此要使用新流程，后端需要实现 Quantizer 类，该类编码：

  1. 后端支持哪些量化算子或模式

  2. 如何让用户表达他们希望其浮点模型如何被量化，例如将整个模型量化为 int8 对称量化，或仅量化线性层等。

为 XNNPACK 定义的现有量化器对象是 XNNPackQuantizer

注解 API#

Quantizer 使用注解 API 来传达不同算子/模式下的量化意图。注解 API 主要由 QuantizationSpec 和 QuantizationAnnotation 组成。

QuantizationSpec 用于传达张量将如何量化的意图，例如 dtype、位宽、最小值、最大值、对称与非对称等。此外， QuantizationSpec 还允许量化器指定如何观测张量值，例如 MinMaxObserver ，或 HistogramObserver ，或一些自定义的观测者。

由 QuantizationSpec 对象组成的 QuantizationAnnotation 用于注解模式的输入张量和输出张量。注解输入张量相当于注解输入边，而注解输出张量相当于注解节点。 QuantizationAnnotation 是一个 dataclass ，包含多个字段：

  • input_qspec_map 字段是 Dict 类型的，用于将每个输入张量（作为输入边）映射到 QuantizationSpec 上。

  • output_qspec 字段表示用于标注输出张量的 QuantizationSpec 类型；

  • _annotated 字段表示该节点是否已经被量化器标注。

总而言之，标注 API 要求量化器标注图的边（输入张量）或节点（输出张量）。现在，将介绍如何使用标注 API 与不同类型的 QuantizationSpec 。

标注常见算子模式#

为了使用量化模式/算子，例如 quantized add ，后端开发者将有意对模式（如 QuantizationSpec 所表达）的输入和输出进行量化。以下示例流程（以 add 算子为例），说明如何在量化工作流中通过标注 API 传达这种意图。

  • 步骤 1：在 FX 图中识别原始浮点模式。识别此模式有几种方法：量化器可以使用模式匹配器来匹配算子模式；量化器可以从头到尾遍历节点，并将节点的目标类型与算子模式进行比较。在这个示例中，可以使用 get_source_partitions 来匹配这个模式。原始浮点 add 模式只包含一个 add 节点。

add_partitions = get_source_partitions(gm.graph, [operator.add, torch.add])
add_partitions = list(itertools.chain(*add_partitions.values()))
for add_partition in add_partitions:
    add_node = add_partition.output_nodes[0]

  • 步骤 2：定义模式的输入和输出的 QuantizationSpec 。 QuantizationSpec 定义了 data type 、 qscheme 以及其他关于用户如何观测或模拟量化的张量意图的量化参数。

act_quantization_spec = QuantizationSpec(
    dtype=torch.int8,
    quant_min=-128,
    quant_max=127,
    qscheme=torch.per_tensor_affine,
    is_dynamic=False,
    observer_or_fake_quant_ctr=HistogramObserver.with_args(eps=2**-12),
)

input_act_qspec = act_quantization_spec
output_act_qspec = act_quantization_spec

  • 步骤 3：使用 QuantizationAnnotation 标注模式的输入和输出。在这个例子中，将使用步骤 2 中创建的 QuantizationAnnotation 对象，为 add 节点的两个输入和一个输出创建 QuantizationSpec 。

input_qspec_map = {}
input_act0 = add_node.args[0]
input_qspec_map[input_act0] = input_act_qspec

input_act1 = add_node.args[1]
input_qspec_map[input_act1] = input_act_qspec

add_node.meta["quantization_annotation"] = QuantizationAnnotation(
    input_qspec_map=input_qspec_map,
    output_qspec=output_act_qspec,
    _annotated=True,
)

在像这样标注 add 节点后，在接下来的量化流程中， HistogramObserver 将在准备阶段插入到其两个输入节点和一个输出节点。在转换阶段， HistogramObserver 将被替换为 quantize 节点和 dequantize 节点。

标注共享量化参数的算子#

用户自然希望标注量化模型，其中量化参数可以明确地在一些张量之间共享。两种典型用例是：

  • 示例 1：一个例子是针对 add ，其中两个输入共享量化参数使得算子实现更加简单。如果不使用 SharedQuantizationSpec，必须在上面的第 1 节中标注 add 作为示例，其中 add 的两个输入具有不同的量化参数。

  • 示例 2：另一个例子是输入和输出之间共享量化参数。这通常是由 maxpool 、 average_pool 、 concat 等算子引起的。

SharedQuantizationSpec 是为这个用例设计的，用于标注其量化参数与其他张量共享的张量。 SharedQuantizationSpec 的输入是 EdgeOrNode 对象，该对象可以是输入边或输出值。

备注

共享是传递的： - 有些张量可能有效地使用了共享量化规范，原因如下： - 两个节点/边被配置为使用 SharedQuantizationSpec 。 - 存在一些节点的现有共享。

例如，假设有两个 conv 节点 conv1 和 conv2 ，它们都被输入到一个 cat 节点： cat([conv1_out, conv2_out], ...) 。假设 conv1 的输出、 conv2 以及 cat 的第一个输入被配置为使用与 QuantizationSpec 相同的配置。 cat 的第二个输入被配置为使用与第一个输入相同的 SharedQuantizationSpec 。

conv1_out: qspec1(dtype=torch.int8, ...)
conv2_out: qspec1(dtype=torch.int8, ...)
cat_input0: qspec1(dtype=torch.int8, ...)
cat_input1: SharedQuantizationSpec((conv1, cat))  # conv1 node is the first input of cat

首先， conv1 的输出隐式地与 cat 的第一个输入共享量化参数（以及观察者对象），同样， conv2 的输出也与 cat 的第二个输入共享。因此，由于用户配置了 cat 的两个输入共享量化参数，根据传递性， conv2_out 和 conv1_out 也将共享量化参数。在观察到的图中，你会看到以下内容：

conv1 -> obs -> cat
conv2 -> obs   /

并且 obs 将是同一个观察者实例。

  • 输入边是输入节点和消耗输入的节点之间的连接，所以它是一个 Tuple[Node, Node] 。

  • 输出值是一个 FX Node 。

现在，如果想用 SharedQuantizationSpec 重新编写 add 注释示例，以指示两个输入张量共享量化参数。可以将其 QuantizationAnnotation 定义为：

  • 步骤 1：在 FX 图中识别原始浮点模式。可以使用 QuantizationSpec 示例中介绍的方法来识别 add 模式。

  • 第二步：用 QuantizationSpec 标注 add 的 input_act0。

  • 第三步：创建 SharedQuantizationSpec 对象，其输入边定义为 (input_act0, add_node) ，这意味着共享用于该边的观察者。然后，用户可以用这个 SharedQuantizationSpec 对象标注 input_act1。

input_qspec_map = {}
share_qparams_with_input_act0_qspec = SharedQuantizationSpec((input_act0, add_node))
input_qspec_map = {input_act0: act_quantization_spec, input_act1: share_qparams_with_input_act0_qspec}

add_node.meta["quantization_annotation"] = QuantizationAnnotation(
    input_qspec_map=input_qspec_map,
    output_qspec=act_quantization_spec,
    _annotated=True,
)

用固定量化参数标注算子#

标注量化模型的另一个典型用例是对于量化参数事先已知的张量。例如，像 sigmoid 这样的算子，其输入和输出张量具有预定义且固定的 scale/zero_point。FixedQParamsQuantizationSpec 就是为了这个用例设计的。要使用 FixedQParamsQuantizationSpec ，用户需要显式地传入 scale 和 zero_point 的参数。

  • 步骤 1：在 FX 图中识别原始浮点模式。我们可以使用 QuantizationSpec 示例中介绍的方法来识别 sigmoid 模式。

  • 步骤 2：创建具有固定 scale 、 zero_point 值的 FixedQParamsQuantizationSpec 对象。这些值将在转换阶段用于创建 quantize 节点和 dequantize 节点。

  • 步骤 3：标注输入和输出以使用这个 FixedQParamsQuantizationSpec 对象。

import torch
from torchao.quantization.pt2e.quantizer.quantizer import FixedQParamsQuantizationSpec, QuantizationAnnotation

act_qspec = FixedQParamsQuantizationSpec(
    dtype=torch.uint8,
    quant_min=0,
    quant_max=255,
    qscheme=torch.per_tensor_affine,
    scale=1.0 / 256.0,
    zero_point=0,
)
sigmoid_node.meta["quantization_annotation"] = QuantizationAnnotation(
    input_qspec_map={input_act: act_qspec},
    output_qspec=act_qspec,
    _annotated=True,
)

使用派生量化参数标注张量#

另一个用例是定义量化参数由其他张量派生而来的张量的约束条件。例如，如果我们想标注一个卷积节点，并定义其偏置输入张量的 scale 为激活张量的 scale 与权重张量的 scale 的乘积，我们可以使用 DerivedQuantizationSpec 来标注这个卷积节点。

  • 步骤 1：在 FX 图中识别原始浮点模式。我们可以使用 QuantizationSpec 示例中介绍的方法来识别 convolution 模式。

  • 第二步：定义 derive_qparams_fn 函数，它接受 ObserverOrFakeQuantize （ObserverBase 或 FakeQuantizeBase）的列表作为输入。从每个 ObserverOrFakeQuantize 对象中，用户可以获取 scale 、 zero point 值。用户可以定义其启发式方法，基于从观察者或假量化实例计算出的量化参数来推导新的 scale 、 zero point 值。

  • 步骤 3：定义 DerivedQuantizationSpec 对象，它接受 EdgeOrNode 对象的列表作为输入。每个 EdgeOrNode 对象对应的观察者将被传递到 derive_qparams_fn 函数； derive_qparams_fn 函数；以及其他一些量化参数，例如 dtype 、 qscheme 。

  • 步骤 4：用 QuantizationAnnotation 标注这个卷积节点的输入和输出。

def derive_qparams_fn(obs_or_fqs: list[ObserverOrFakeQuantize]) -> tuple[Tensor, Tensor]:
    assert len(obs_or_fqs) == 2, \
        "Expecting two obs/fqs, one for activation and one for weight, got: {}".format(len(obs_or_fq))
    act_obs_or_fq = obs_or_fqs[0]
    weight_obs_or_fq = obs_or_fqs[1]
    act_scale, act_zp = act_obs_or_fq.calculate_qparams()
    weight_scale, weight_zp = weight_obs_or_fq.calculate_qparams()
    return torch.tensor([act_scale * weight_scale]).to(torch.float32), torch.tensor([0]).to(torch.int32)

bias_qspec = DerivedQuantizationSpec(
    derived_from=[(input_act, node), (weight, node)],
    derive_qparams_fn=derive_qparams_fn,
    dtype=torch.int32,
    quant_min=-2**31,
    quant_max=2**31 - 1,
    qscheme=torch.per_tensor_symmetric,
)
input_qspec_map = {input_act: act_quantization_spec, weight: weight_quantization_spec, bias: bias_qspec}
node.meta["quantization_annotation"] = QuantizationAnnotation(
    input_qspec_map=input_qspec_map,
    output_qspec=act_quantization_spec,
    _annotated=True,
)

Resnet18 的玩具示例#

在用 QuantizationAnnotation API 注释方法定义之后，现在可以将它们组合起来构建 BackendQuantizer 并用 Torchvision Resnet18 运行示例。为了更好地理解最终的示例，这里列出了示例中使用的类和工具函数：

  • QuantizationConfig 包含分别针对激活值、权重和偏置的 QuantizationSpec 。

  • 在标注模型时，可以使用 get_input_act_qspec、get_output_act_qspec、get_weight_qspec 和 get_bias_qspec 来为特定模式获取 QuantizationSpec 中的 QuantizationConfig 。

目录