PyTorch 量化 BackendConfig 完全指南:为多后端定制算子级量化行为

发布时间:2026/9/10 21:18:15
PyTorch 量化 BackendConfig 完全指南:为多后端定制算子级量化行为
PyTorch 量化 BackendConfig 完全指南为多后端定制算子级量化行为【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorchBackendConfig 是 PyTorchFX 图模式量化的可配置化基石它把哪些算子模式可以量化、支持什么数据类型、如何插入观察器、如何做模块融合与参考量化模块替换全部抽象为声明式配置从而让同一套量化流程适配 FBGEMM、QNNPACK含 xnnpack、TensorRT、Executorch 以及任意第三方后端。本文以 torch/ao/quantization/backend_config/README.md 为主体结合仓库内backend_config.py、各后端配置实现与 FX 量化 pass 源码系统讲解 BackendConfig 的四大核心能力——模式规范、观察器插入、参考量化模块映射、融合与数据类型约束并给出可运行的自定义后端配置示例帮助你掌握为任意推理后端定制量化规则的方法。BackendConfig 解决什么问题不同的量化后端或 kernel 库对量化算子的支持千差万别。PyTorch 量化尤其 FX 图模式需要一种机制把某个后端支持哪些量化算子模式以及同一算子模式在不同后端上的差异化处理完全参数化而不是硬编码在量化 pass 里。BackendConfig 正是为此而生它目前只服务于 FX 图模式量化与 FX 量化流程的集成细节可参考 FX 量化 README。BackendConfig 按算子模式operator pattern配置量化行为。对每个算子模式需要说明输入/输出激活、权重、偏置分别支持的数据类型float32、float16、int8、uint8、bfloat16 等量化方案对称 vs 非对称、per-channel vs per-tensor的适用范围量化及融合映射部分量化算子相对朴素的dequant - float_op - quant参考实现在数值上可能有差异对有权重算子如 conv、linear需要指定自定义参考模块以及从 float 模块到参考模块的映射QAT 映射对有权重算子需要替换为在权重上插入伪量化的量化感知训练QAT版本。以 fbgemm 为例其能力画像如下维度fbgemm量化方案激活 per-tensor权重 per-tensor 或 per-channel数据类型激活 quint8带 qmin/qmax 范围限制权重 qint8量化/融合算子及映射如torch.nn.Conv2d - torch.ao.nn.quantized.reference.Conv2dQAT 模块映射如torch.nn.Conv2d - torch.ao.nn.qat.Conv2d代码库中不再硬编码融合映射、float 到参考量化模块的映射、融合模式等而是统一从 BackendConfig 推导。这一设计让 PyTorch 量化既能服务第一方后端fbgemm、qnnpack也能服务差异较大的第三方后端TensorRT、Executorch 等近期集成到 qnnpack 后端中的 xnnpack正是依靠 BackendConfig 来定义 xnnpack 量化算子所需的额外约束见 qnnpack.py。模式规范Pattern SpecificationBackendConfig 中的算子模式可以是 float 模块、functional 算子、PyTorch 算子或上述元素的元组组合例如torch.nn.Lineartorch.nn.functional.lineartorch.addoperator.add(torch.nn.functional.linear, torch.nn.functional.relu)(torch.nn.Conv2d, torch.nn.BatchNorm2d, torch.nn.ReLU)元组模式被视为顺序模式sequential patterns当前仅支持 2 元组或 3 元组。用户面 API 中2 元组(a, b)与 3 元组(a, b, c)均按前向顺序书写。高级模式反向嵌套元组格式上述格式覆盖绝大多数场景但无法表达图DAG模式。为此 BackendConfig 提供替代的反向嵌套元组格式通过BackendPatternConfig()._set_pattern_complex_format(...)启用。注意该格式已废弃将在未来版本被替换。operator module_type | functional | torch op | native op | MatchAllNode Pattern (operator, Pattern, Pattern, ...) | operator其中每个 Pattern 的第一项是算子其余是该算子的各个参数的子模式。例如模式(nn.ReLU, (operator.add, MatchAllNode, (nn.BatchNorm2d, nn.Conv2d)))匹配如下计算图tensor_1 tensor_2 | | *(MatchAllNode) nn.Conv2d | | | nn.BatchNorm2d \ / -- operator.add -- | nn.ReLU在 prepare 和 convert 阶段匹配发生在最后一个节点即匹配的锚点然后从该节点向前回溯即可还原整张子图。上例中匹配到nn.ReLU节点node.args[0]就是operator.add节点。内部实现上用户面的正向元组会被转换成反向嵌套元组供模式匹配使用。转换规则见 utils.py 的_get_pattern_in_reversed_nested_tuple_format2 元组(a, b)转为(b, a)3 元组(a, b, c)转为(c, (b, a))。例如(nn.Linear, nn.ReLU)内部表示为(nn.ReLU, nn.Linear)(nn.Conv2d, nn.BatchNorm2d, nn.ReLU)内部表示为(nn.ReLU, (nn.BatchNorm2d, nn.Conv2d))。BackendConfig 内部正是以该格式为键存储配置见 backend_config.py。未来计划用 torch.fx 的 subgraph rewriter 取代这套自维护的模式匹配代码。BackendConfig 实现与完整示例BackendConfig 由一组 BackendPatternConfig 组成每个 BackendPatternConfig 定义单个算子模式的规格与要求。下面是 README 与 backend_config.py 文档字符串中给出的完整示例import torch from torch.ao.quantization.backend_config import ( BackendConfig, BackendPatternConfig, DTypeConfig, ObservationType, ) weighted_int8_dtype_config DTypeConfig( input_dtypetorch.quint8, output_dtypetorch.quint8, weight_dtypetorch.qint8, bias_dtypetorch.float) def fuse_conv2d_relu(is_qat, conv, relu): Return a fused ConvReLU2d from individual conv and relu modules. return torch.ao.nn.intrinsic.ConvReLU2d(conv, relu) # For quantizing Linear linear_config BackendPatternConfig(torch.nn.Linear) \ .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \ .add_dtype_config(weighted_int8_dtype_config) \ .set_root_module(torch.nn.Linear) \ .set_qat_module(torch.ao.nn.qat.Linear) \ .set_reference_quantized_module(torch.ao.nn.quantized.reference.Linear) # For fusing Conv2d ReLU into ConvReLU2d conv_relu_config BackendPatternConfig((torch.nn.Conv2d, torch.nn.ReLU)) \ .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \ .add_dtype_config(weighted_int8_dtype_config) \ .set_fused_module(torch.ao.nn.intrinsic.ConvReLU2d) \ .set_fuser_method(fuse_conv2d_relu) # For quantizing ConvReLU2d fused_conv_relu_config BackendPatternConfig(torch.ao.nn.intrinsic.ConvReLU2d) \ .set_observation_type(ObservationType.OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT) \ .add_dtype_config(weighted_int8_dtype_config) \ .set_root_module(torch.nn.Conv2d) \ .set_qat_module(torch.ao.nn.intrinsic.qat.ConvReLU2d) \ .set_reference_quantized_module(torch.ao.nn.quantized.reference.Conv2d) backend_config BackendConfig(my_backend) \ .set_backend_pattern_config(linear_config) \ .set_backend_pattern_config(conv_relu_config) \ .set_backend_pattern_config(fused_conv_relu_config)要点拆解BackendConfig(my_backend)以后端名称为参数后续通过set_backend_pattern_config逐个注册也可用set_backend_pattern_configs批量注册列表重复模式会覆盖旧配置见 backend_config.pyBackendPatternConfig(pattern)既接受单个算子也接受顺序元组模式set_pattern与_set_pattern_complex_format互斥只能二选一add_dtype_config追加一个支持的 dtype 组合set_dtype_configs则整体覆盖一个融合模式需要两组BackendPatternConfig 协作一组描述(Conv2d, ReLU)模式的融合规则set_fuser_methodset_fused_module另一组描述融合产物ConvReLU2d的量化规则root/QAT/reference 映射。这与仓库内_get_conv_configs、_get_linear_configs的编排方式一致见 _common_operator_config_utils.py。BackendConfig 还支持与字典形式互转to_dict/from_dict见 backend_config.pyget_native_backend_config_dict()即返回字典形式便于序列化或向后兼容。观察器插入Observer Insertion相关 APIset_observation_typeprepare 阶段按观察类型向图中插入观察器未来将改为插入带观察器/FakeQuantize 的 QuantDeQuantStub。ObservationType枚举定义在 backend_config.py共三档取值含义典型算子OUTPUT_USE_DIFFERENT_OBSERVER_AS_INPUT默认输入与输出使用不同观察器实例各自依据qconfig.activationconv、linear、softmaxOUTPUT_SHARE_OBSERVER_WITH_INPUT输出复用输入的观察器实例torch.cat、maxpool、reshape 等共享量化参数算子INPUT_OUTPUT_NOT_OBSERVED输入输出都不观察x.shape、x.size仓库中的使用佐证torch.cat配置为OUTPUT_SHARE_OBSERVER_WITH_INPUT_common_operator_config_utils.pyshape、size这类张量信息算子配置为INPUT_OUTPUT_NOT_OBSERVED同文件 L769-L783。此外二元算子add/mul还通过_set_num_tensor_args_to_observation_type按张量参数个数0/1/2分别指定观察类型以区分标量参与与双张量参与的情形。参考量化模式Reference Quantized Patterns相关 APIset_root_module、set_reference_quantized_moduleconvert 阶段构造参考量化模型时root 模块如nni.LinearReLU或nniqat.LinearReLU背后的torch.nn.Linear会被替换为对应的参考量化模块如torch.ao.nn.quantized.reference.Linear。这一对一映射允许自定义后端提供数值与其下移算子匹配的参考实现——root 模块与参考量化模块必须配置在同一条 BackendPatternConfig 中替换才会发生见 backend_config.py。utils.py 的get_root_module_to_quantized_reference_module会收集全部满足root 与 reference 同时非空的映射FX 的 convert.py 在 convert 阶段使用该映射完成模块替换。同理get_module_to_qat_module收集 pattern 到 QAT 模块的映射供 QAT 流程使用。融合Fusion相关 APIset_fuser_method、set_fused_module、_set_root_node_getter、_set_extra_inputs_getter融合是优化手段(torch.nn.Linear, torch.nn.ReLU)这类模式可在prepare 阶段由set_fuser_method指定的函数融合为nni.LinearReLUconvert 阶段再把这些融合模块由set_fused_module标识转成参考量化版本。融合函数第一个参数必须是is_qat其余参数依次对应元组模式中的元素返回融合后的模块。例如def fuse_linear_relu(is_qat, linear, relu): return torch.ao.nn.intrinsic.LinearReLU(linear, relu)仓库内置的 conv/linear 系列融合大多通过_sequential_wrapper2包装完成convbn(relu) 使用fuse_conv_bn/fuse_conv_bn_relu/fuse_convtranspose_bn/fuse_linear_bn见 _common_operator_config_utils.py 的导入及 L326-L341 的使用。在 FX 图模式中融合替换依赖两个由用户提供的辅助函数root_node_getter返回 root 节点通常是模式中带权重的模块节点如torch.nn.Linear用于在图里替换整个匹配模式extra_inputs_getter返回附加输入参数列表追加到 fused 模块从 root 节点拷贝而来的现有参数之后。这两个辅助函数可通过_set_root_node_getter、_set_extra_inputs_getter配置backend_config.py并由 utils.py 导出为get_fusion_pattern_to_root_node_getter/get_fusion_pattern_to_extra_inputs_getter供 FX fuse.py 使用。例如对模式(torch.add, MatchAllNode, (torch.nn.BatchNorm2d, torch.nn.Conv2d))root 是 Conv2dextra_inputs_getter可返回 MatchAllNode 处的额外输入节点。融合函数若以正向元组书写会被_reverse2/_reverse3转换为内部格式见 utils.py。数据类型限制Data Type Restrictions相关 APIadd_dtype_config、set_dtype_configsDTypeConfig 指定输入/输出/权重/偏置的一组受支持数据类型及关联约束。input_dtype、output_dtype、weight_dtype有两种写法简单torch.dtype或带约束的DTypeWithConstraintsimport torch from torch.ao.quantization.backend_config import DTypeConfig, DTypeWithConstraints dtype_config DTypeConfig( input_dtypetorch.quint8, output_dtypetorch.quint8, weight_dtypetorch.qint8, bias_dtypetorch.float) dtype_config_with_constraints DTypeConfig( input_dtypeDTypeWithConstraints( dtypetorch.quint8, quant_min_lower_bound0, quant_max_upper_bound255, scale_min_lower_bound2 ** -12, ), output_dtypeDTypeWithConstraints( dtypetorch.quint8, quant_min_lower_bound0, quant_max_upper_bound255, scale_min_lower_bound2 ** -12, ), weight_dtypeDTypeWithConstraints( dtypetorch.qint8, quant_min_lower_bound-128, quant_max_upper_bound127, scale_min_lower_bound2 ** -12, ), bias_dtypetorch.float)DTypeWithConstraints的全部字段定义在 backend_config.pydtype、quant_min_lower_bound、quant_max_upper_bound、scale_min_lower_bound、scale_max_upper_bound、scale_exact_match、zero_point_exact_match。prepare 阶段会把 DTypeConfig 中声明的数据类型与匹配该算子模式的 QConfig 比对若所有 DTypeConfig 都不匹配或约束不满足则该 QConfig 被忽略对应模式不会被量化。注意 DTypeConfig 的 dtype 语义与观察器一致指的是参考模型中 quantize op 的参数 dtype而非算子接口 dtype——例如动态量化的接口 dtype 是 fp32但 DTypeConfig 中input_dtype仍写 quint8见 backend_config.py 的参考模型示意。量化范围Quantization range用户的 QConfig 可指定quant_min/quant_max限制量化值范围。quant_min_lower_bound是后端允许的quant_min下界quant_max_upper_bound是允许的quant_max上界QConfig 越界即视为违反约束。例如 qnnpack/xnnpack 对对称量化权重要求量化值落在[-127, 127]排除 -128见 qnnpack.py 中qnnpack_weight_qint8_neg_127_to_127_scale_min_2_neg_12的定义。尺度范围Scale range类似地QConfig 可指定量化 scale 的最小值当前以eps暴露未来会更名以更好表达语义。scale_min_lower_bound表示后端允许的 scale 下界QConfig 的最小 scale 低于该值即违反约束。例如 xnnpack 要求 requantization scale 不低于2 ** -12qnnpack.py。注意scale_max_upper_bound目前并未实际生效因为观察器尚无对应的强制机制。固定量化参数Fixed quantization parameters对于torch.nn.Sigmoid、torch.nn.Tanh这类量化参数固定的算子BackendConfig 可以在输入/输出激活上指定精确的 scale 与 zero pointscale_exact_match、zero_point_exact_match。这些算子的用户 QConfig 必须对激活使用FixedQParamsObserver或FixedQParamsFakeQuantize且 scale/zero point 与约束一致否则 QConfig 会被忽略。仓库中为 sigmoid、hardsigmoid、softmax 配置了[0, 1]约束scale 精确匹配1.0 / 256.0、zero point 精确匹配0为 tanh 配置了[-1, 1]约束scale 精确匹配2.0 / 256.0、zero point 精确匹配128见 _common_operator_config_utils.py_add_fixed_qparams_to_dtype_configs还会在激活约束中同时设置quant_min_lower_bound0、quant_max_upper_bound255并对同时指定了 scale 上下界的配置抛异常与固定参数语义冲突。仓库内置的后端配置一览torch/ao/quantization/backend_config/目录下除核心类外还内置了多个后端的现成配置均可通过init.py 直接导入使用后端入口函数关键特性nativefbgemm/qnnpack 默认集合get_native_backend_config()/get_native_backend_config_dict()覆盖 conv/linear/二元算子/cat/固定参数/共享参数/张量信息/BN/LayerNorm/RNN/Embedding 等全系列模式见 native.pyfbgemmget_fbgemm_backend_config()激活 quint8、权重 qint8 的静态量化 int8/float16 动态量化 weight-onlyquint8/quint4x2见 fbgemm.pyqnnpack含 xnnpackget_qnnpack_backend_config()额外提供 qint8 对称量化配置激活 scale 下限2**-12、权重值域[-127, 127]见 qnnpack.pytensorrtget_tensorrt_backend_config()/get_tensorrt_backend_config_dict()qint8 输入/输出/权重、float 偏置并对torch.addmm显式指定 bias/input/weight 的输入索引API 标记为实验性见 tensorrt.pyexecutorchget_executorch_backend_config()复用 qnnpack 的对称 qint8 约束支持 quint8 静态、qint8/quint8/float16 动态与 weight-only覆盖 Conv2d/Linear/二元算子/共享参数/BN/cat/Embedding见 executorch.pyonednnget_onednn_backend_config()面向 oneDNN 后端的配置这些内置配置均复用 _common_operator_config_utils.py 中的模式工厂_get_conv_configs、_get_linear_configs、_get_binary_op_configs等以不同 dtype 配置为参数批量生成各模式的 BackendPatternConfig。Conv 系列通过_Conv1dMetadata/_Conv2dMetadata/_Conv3dMetadata三个 namedtuple 集中声明 root/transpose/bn/reference/fused/QAT/functional 的对应关系一处定义、多处复用。从 BackendConfig 到 FX 量化流程BackendConfig 最终通过 utils.py 导出的一组映射被 FX 量化 pass 消费get_pattern_to_dtype_configspattern → 支持的 DTypeConfig 列表prepare 阶段用于校验 QConfig见 prepare.pyget_qat_module_classes/get_fused_module_classesprepare 阶段收集需要处理的 QAT / 融合模块类get_root_module_to_quantized_reference_moduleconvert 阶段做 root → 参考量化模块替换convert.pyget_module_to_qat_module、get_fuser_method_mapping、get_fusion_pattern_to_root_node_getter、get_fusion_pattern_to_extra_inputs_getter分别驱动 QAT 替换、融合方法调用与图节点替换get_pattern_to_input_type_to_index为 functional 算子如F.linear、F.conv2d、F.layer_norm、F.embedding标注 weight/bias 在参数列表中的位置便于提取权重做量化统计pattern_to_human_readable/entry_to_pretty_str把内部 pattern 与配置项转换为人类可读字符串服务于量化文档/诊断输出。以准备一个自定义后端为例整体工作流可概括为用 DTypeConfig 声明数据能力 → 用 BackendPatternConfig 为每个支持的算子模式声明观察类型、融合规则与模块映射 → 聚合为 BackendConfig 传入 FX 量化入口 → prepare/convert 阶段自动消费这些声明完成观察器插入、融合与参考模块替换。若某个模式未在 BackendConfig 中登记或用户 QConfig 不满足其 dtype 约束则该模式保持 float 运行而不被量化——这正是以配置驱动一切的设计初衷。小结BackendConfig 把 PyTorch FX 图模式量化的行为差异从代码中剥离为声明式配置模式规范含反向嵌套元组表达复杂子图、观察器插入策略、root/参考/QAT 模块映射、融合方法与数据类型约束共同构成了完整的多后端量化定制面。理解并复用仓库中 native/fbgemm/qnnpack/tensorrt/executorch 等内置配置的写法即可为自有推理后端编写高契合度的量化配置实现同一套 PyTorch 量化流程服务任意目标硬件的目标。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考