(原型) PyTorch 2 Export 量化简介

(原型) PyTorch 2 Export 量化简介#

参考：pt2e_quant_ptq

定义简单的 PyTorch 模型，包含单个线性层：

import torch


class M(torch.nn.Module):
   def __init__(self):
      super().__init__()
      # 创建一个输入维度为5，输出维度为10的线性层
      self.linear = torch.nn.Linear(5, 10)

   def forward(self, x):
      # 前向传播仅包含线性变换
      return self.linear(x)

准备示例输入：形状为 (1, 5) 的随机张量

example_inputs = (torch.randn(1, 5),)

初始化模型并设置为评估模式：

m = M().eval()

  1. Step 1. 程序捕获 (Program Capture)：此功能适用于 PyTorch 2.6 及以上版本。对于较低版本的 PyTorch，请查看 使用 torch.export 导出模型获取更多详细信息

m = torch.export.export(m, example_inputs).module()

此时得到包含 aten 算子的模型。

  1. Step 2. 量化 (Quantization)：从 torchao 库导入 PT2E 量化所需的函数

from torchao.quantization.pt2e.quantize_pt2e import (
  prepare_pt2e,  # 准备模型进行量化
  convert_pt2e,  # 将准备好的模型转换为量化模型
)

  1. 从 executorch 的 XNNPACK 后端导入量化器：

    • 安装 executorch: pip install executorch

from executorch.backends.xnnpack.quantizer.xnnpack_quantizer import (
  get_symmetric_quantization_config,  # 获取对称量化配置
  XNNPACKQuantizer,  # XNNPACK量化器
)

后端开发人员会编写自己的量化器并公开方法，允许用户表达他们希望如何量化模型。

创建量化器并设置全局对称量化配置：

quantizer = XNNPACKQuantizer().set_global(get_symmetric_quantization_config())

准备模型进行量化

m = prepare_pt2e(m, quantizer)

# 校准步骤省略 (通常需要使用校准数据集来确定量化参数)
# 在校准过程中，模型会学习数据分布特征，以便更准确地进行量化
# 将准备好的模型转换为量化模型
m = convert_pt2e(m)
# 现在得到在可能的情况下使用整数计算的aten算子的模型
# 量化后的模型通常具有更小的体积和更快的推理速度，但可能会有轻微的精度损失

/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/torchao/quantization/pt2e/utils.py:145: UserWarning: must run observer before calling calculate_qparams. Returning default values.
  warnings.warn(
/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/torchao/quantization/pt2e/observer.py:1350: UserWarning: must run observer before calling calculate_qparams.                                    Returning default scale and zero point 
  warnings.warn(

PyTorch 2 export 量化的动机#

在 PyTorch 版本 2 之前，有 FX 图模式量化，它使用 QConfigMapping 和 BackendConfig 进行自定义。QConfigMapping 允许建模用户指定他们希望如何量化模型，BackendConfig 允许后端开发者指定他们后端支持量化方式。虽然这个 API 相对较好地涵盖了大多数用例，但它并非完全可扩展。当前 API 有两个主要限制：

  1. 在现有对象（ QConfig 和 QConfigMapping ）中表达复杂算子模式量化意图（即算子模式应如何被观测/量化）方面的限制。

  2. 用户表达其模型量化意图的支持有限。例如，如果用户希望对模型中的每隔一个线性层进行量化，或者量化行为依赖于张量的实际形状（例如，仅在具有 3D 输入的线性层时观察/量化输入和输出），后端开发者或建模用户需要更改核心量化 API/流程。

  3. 使用 QConfigMapping 和 BackendConfig 作为独立对象， QConfigMapping 描述用户希望其模型如何量化的意图， BackendConfig 描述后端支持何种量化方式。 BackendConfig 是后端特定的，但 QConfigMapping 不是，用户可以提供与特定 BackendConfig 不兼容的 QConfigMapping ，这不是很好的用户体验。理想情况下，可以通过使配置（ QConfigMapping ）和量化能力（ BackendConfig ）都成为后端特定的，来更好地组织这些内容，从而减少关于不兼容性的困惑。

  4. 在 QConfig 中，将观测者/ fake_quant 观测者类作为对象暴露给用户以配置量化，这增加了用户可能需要关心的事项。例如，不仅包括 dtype ，还包括观测应该如何进行，这些都可以对用户隐藏，从而简化用户流程。

以下是新 API 的优缺点总结：

  1. 可编程性（解决 1. 和 2.）：当用户的量化需求未被现有量化器覆盖时，用户可以构建自己的量化器，并像上面提到的那样将其与其他量化器组合。

  2. 简化用户体验（解决 3.）：提供单一实例，供后端和用户交互。因此，您不再需要用户面对量化配置映射来映射用户意图，以及后端交互的独立量化配置来配置后端支持的内容。仍然会为用户提供查询量化器支持内容的方法。通过单一实例，组合不同的量化能力也比以前更自然。

  3. 例如 XNNPACK 不支持 embedding_byte ，而 ExecuTorch 原生支持这个功能。因此，如果有 ExecuTorchQuantizer 仅量化 embedding_byte ，那么它就可以与 XNNPACKQuantizer 组合。（以前，这需要将两个 BackendConfig 连接在一起，并且由于 QConfigMapping 中的选项不是后端特定的，用户也需要自行决定如何指定与组合后端量化能力匹配的配置。通过单个量化器实例，可以组合两个量化器，并查询组合量化器的功能，这使得它更不容易出错且更清晰，例如 composed_quantizer.quantization_capabilities() 。

  4. 关注点分离（解决 4.）：在设计量化器 API 时，将量化的规格（以 dtype 表示，包括最小/最大位数、对称性等）与观测者概念解耦。目前，观测者既捕获量化规格，也捕获如何观测（直方图与最小最大观测者）。建模用户通过这一改变，可以免于与观测者和伪量化对象交互。

目录