From a4d50c1f47d523322da8d19c1d2354e316e3bcc0 Mon Sep 17 00:00:00 2001 From: Min Guo Date: Wed, 26 Aug 2026 20:46:23 -0700 Subject: [PATCH] Annotate conv+bn as one partition under QAT (#22194) Summary: `QnnQuantizer` has no conv-bn handling, which makes every conv+BatchNorm model fail to lower through the OSS QNN QAT path. It is the only one of the three quantizers without it -- `XNNPACKQuantizer` has `_do_annotate_conv_bn` (`backends/xnnpack/quantizer/xnnpack_quantizer_utils.py:437`) and BoltNN has the `Conv{1,2,3}dBn[Relu]Quantizer` family (`bolt/nn/executorch/quantization/_conv.py`). QNN's only BatchNorm entry is `aten.batch_norm.default`, marked `qnn_op=None` with a TODO (`annotators/htp_rules.py:186`). The per-node annotators treat conv and bn as unrelated ops. That is fine for PTQ -- bn is already folded into the conv by the time it is quantized -- but it breaks QAT in two ways. **1. BatchNorm is never folded.** `prepare_qat_pt2e` annotates *before* torchao's `_fuse_conv_bn_qat` rewrites conv+bn into the QAT folding pattern. torchao orders it that way deliberately ("Perform fusion after annotate to avoid quantizing ops in the new subgraph", `quantize_pt2e.py:63`), but the rewrite copies each matched node's `meta` onto its replacement and re-points `SharedQuantizationSpec` chains at the new nodes, so a qspec on the *conv output* lands on the scale arithmetic. `convert_pt2e`'s `_fold_conv_bn_qat` matches a pattern with q/dq at exactly three points -- conv input, scaled weight, bn output -- so the extra q/dq makes the matcher miss and bn survives. Float bn then reaches the QNN partitioner, which has no node visitor for the training variant, drops it along with its FP16 scale arithmetic, and shatters the graph into one partition per conv. A partition left with only constant inputs fails serialization as `No graph inputs present for graph [0]` -> `0x7532` / `Error 30002`, which is what this actually surfaces as. **2. The folded bias is unquantized on `bias=False` convs.** When the conv has no bias, `fold_bn_weights_into_conv_node` materializes the folded bias as a fresh fp32 parameter and explicitly leaves it alone -- *"NOTE: here we assume the bias of conv is not quantized!"* (`torchao/quantization/pt2e/utils.py`). QNN rejects a per-channel-quantized conv carrying a float bias (`0xc26`), drops it from the partition, and the orphaned weight dequant has to become a standalone Dequantize -- which QNN does not support per-channel. `conv_bn.py` adds `annotate_conv_bn_partitions`, run before the per-node pass: - claims each `conv -> bn` chain as one partition, putting the qspecs on the conv's inputs and the output qspec on the bn output, with nothing on the conv output, so the fold matches. A trailing relu is deliberately NOT absorbed: XNNPACK's equivalent puts the qspec on the relu because it fuses conv+relu into one backend op, but QNN lowers them separately and rejects a conv whose output is unquantized (`0x232 != 0x408`, FLOAT_32 meeting UFIXED_POINT_8). Keeping the qspec on the bn output -- which the fold transfers onto the conv output -- leaves the relu to the per-node annotator, which gives it its own quantized input and output; - gives a bias-less conv a zero bias first. That is a numerical no-op, but it gives the annotator a node to attach the derived bias qspec to, so the fold takes its `conv_bias_node is not None` branch and writes the folded values into an already-quantized node instead of inventing an unquantized one. Gated on `is_qat`. PTQ already works and is left untouched. **Requires D117580221.** That fixes `_fold_conv_bn_qat` failing to erase BatchNorm's `num_batches_tracked += 1` when export lifts the literal into a tensor constant. Without it the dead counter stays live as a mutated buffer, QNN maps each one to graph I/O, and a 53-BatchNorm model produces a context binary declaring 54 inputs / 56 outputs against ExecuTorch's 1 / 3 -- the runtime then indexes its 4 tensors with the QNN counts and segfaults. Both diffs must be applied for the QAT training run *and* the export. Reviewed By: YIWENX14 Differential Revision: D117419081 --- backends/qualcomm/quantizer/conv_bn.py | 208 ++++++++++++++++++++ backends/qualcomm/quantizer/quantizer.py | 9 + backends/qualcomm/tests/BUCK | 21 ++ backends/qualcomm/tests/test_conv_bn_qat.py | 172 ++++++++++++++++ 4 files changed, 410 insertions(+) create mode 100644 backends/qualcomm/quantizer/conv_bn.py create mode 100644 backends/qualcomm/tests/test_conv_bn_qat.py diff --git a/backends/qualcomm/quantizer/conv_bn.py b/backends/qualcomm/quantizer/conv_bn.py new file mode 100644 index 00000000000..4fc7875ce6c --- /dev/null +++ b/backends/qualcomm/quantizer/conv_bn.py @@ -0,0 +1,208 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Partition-level annotation for conv + batchnorm under QAT. + +The per-node annotators in ``annotators/`` treat conv and batchnorm as unrelated +ops. That is correct for PTQ, where batchnorm is folded into the conv before the +graph is quantized, but it breaks QAT. + +``prepare_qat_pt2e`` annotates the graph *before* torchao's ``_fuse_conv_bn_qat`` +rewrites each conv+bn into the QAT folding pattern:: + + scale = bn_weight / sqrt(bn_running_var + eps) + y = conv(x, w * scale.reshape(w_shape)) / scale.reshape(b_shape) + out = batch_norm(y, ...) + +torchao does that ordering deliberately -- "Perform fusion after annotate to +avoid quantizing ops in the new subgraph" -- but the rewrite copies each matched +node's ``meta`` onto its replacement and re-points ``SharedQuantizationSpec`` +chains at the new nodes, so a qspec on the *conv output* ends up on the scale +arithmetic. ``convert_pt2e``'s ``_fold_conv_bn_qat`` matches a pattern carrying +q/dq at exactly three points -- conv input, scaled weight, bn output -- so the +extra q/dq makes the subgraph matcher miss and batchnorm is never folded. Float +batchnorm then reaches the QNN partitioner, which has no node visitor for the +training variant, drops it along with its FP16 scale arithmetic, and shatters +the graph into one partition per conv. A partition left with only constant +inputs fails context-binary serialization with the opaque +``No graph inputs present for graph [0]`` / ``0x7532``. + +Annotating the partition as a unit -- qspecs on the conv's inputs, the output +qspec on the *last* node of the partition, nothing on the conv output -- gives +the fold the shape it expects. This mirrors ``_do_annotate_conv_bn`` in the +XNNPACK quantizer and the ``Conv*Bn*Quantizer`` family in BoltNN; QNN was the +only one of the three without it. +""" + +import operator +from typing import Callable, List, Optional, Set + +import torch +from torch.fx import GraphModule, Node +from torchao.quantization.pt2e.quantizer import QuantizationAnnotation + +from .qconfig import QuantizationConfig +from .rules import _is_annotated, _mark_nodes_as_annotated, Q_ANNOTATION_KEY + +CONV_TARGETS = ( + torch.ops.aten.conv1d.default, + torch.ops.aten.conv2d.default, + torch.ops.aten.conv2d.padding, + torch.ops.aten.convolution.default, + torch.ops.aten.conv_transpose1d.default, + torch.ops.aten.conv_transpose2d.input, +) + +# ``batch_norm.default`` returns a bare tensor; every other variant returns a +# tuple whose element 0 is the normalized output. +SINGLE_OUTPUT_BN_TARGETS = (torch.ops.aten.batch_norm.default,) + +BN_TARGETS = SINGLE_OUTPUT_BN_TARGETS + ( + torch.ops.aten._native_batch_norm_legit.default, + torch.ops.aten._native_batch_norm_legit_functional.default, + torch.ops.aten._native_batch_norm_legit_no_training.default, + torch.ops.aten.cudnn_batch_norm.default, +) + + +def _sole_user(node: Node) -> Optional[Node]: + return next(iter(node.users)) if len(node.users) == 1 else None + + +def _bn_output(bn: Node) -> Optional[Node]: + """The node carrying the batchnorm's normalized output.""" + if bn.target in SINGLE_OUTPUT_BN_TARGETS: + return bn + for user in bn.users: + if user.target is operator.getitem and user.args[1] == 0: + return user + return None + + +def _materialize_conv_bias(gm: GraphModule, conv: Node) -> Optional[Node]: + """Give a bias-less conv a zero bias so the fold produces a *quantized* one. + + When the conv has no bias, ``fold_bn_weights_into_conv_node`` materializes the + folded bias as a fresh fp32 parameter and explicitly leaves it unquantized + ("NOTE: here we assume the bias of conv is not quantized!"). QNN rejects a + per-channel-quantized conv carrying a float bias, so the conv drops out of the + partition and its orphaned weight dequant has to become a standalone + Dequantize -- which QNN does not support per-channel. + + Adding a zero bias is a numerical no-op that gives the annotator a node to + attach the derived bias qspec to. The fold then takes its + ``conv_bias_node is not None`` branch and writes the folded values into that + already-quantized node instead of inventing an unquantized one. + """ + args = list(conv.args) + while len(args) < 3: + args.append(None) + if isinstance(args[2], Node): + return args[2] + + val = conv.meta.get("val") + if val is None: + return None + + zeros = torch.zeros(val.shape[1], dtype=val.dtype, device=val.device) + name = f"{conv.name}_zero_bias" + gm.register_parameter(name, torch.nn.Parameter(zeros, requires_grad=False)) + with gm.graph.inserting_before(conv): + bias = gm.graph.get_attr(name) + bias.meta["val"] = val.fake_mode.from_tensor(zeros, static_shapes=True) + + args[2] = bias + conv.args = tuple(args) + return bias + + +def _annotate_partition( + gm: GraphModule, + conv: Node, + partition: List[Node], + output_node: Node, + quantization_config: QuantizationConfig, +) -> None: + input_qspec_map = {} + + act = conv.args[0] + if isinstance(act, Node) and quantization_config.input_activation is not None: + input_qspec_map[act] = quantization_config.input_activation + + weight = conv.args[1] + assert isinstance(weight, Node) + input_qspec_map[weight] = quantization_config.weight + + bias = _materialize_conv_bias(gm, conv) + if bias is not None and quantization_config.bias is not None: + input_qspec_map[bias] = ( + quantization_config.bias(conv) + if callable(quantization_config.bias) + else quantization_config.bias + ) + + # Deliberately no output_qspec on the conv -- an observer there lands inside + # the QAT fusion and stops `_fold_conv_bn_qat` from matching. + conv.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + input_qspec_map=input_qspec_map, + _annotated=True, + ) + _mark_nodes_as_annotated(partition) + output_node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( + output_qspec=quantization_config.output_activation, + _annotated=True, + ) + + +def annotate_conv_bn_partitions( + gm: GraphModule, + get_quant_config: Callable[[Node], Optional[QuantizationConfig]], + discard_nodes: Set[str], +) -> int: + """Annotate every conv+bn chain as a single partition. + + Must run before the per-node annotation pass; the nodes it claims are marked + annotated so the per-node annotators skip them. + + A trailing relu is deliberately NOT absorbed into the partition. XNNPACK's + equivalent puts the output qspec on the relu because it fuses conv+relu into + one backend op, but QNN lowers conv and relu to separate ops and rejects a + conv whose output is unquantized (`0x232 != 0x408`, FLOAT_32 meeting + UFIXED_POINT_8). Keeping the qspec on the bn output -- which the fold + transfers onto the conv output -- leaves the relu to the per-node annotator, + which gives it its own quantized input and output. + + Returns the number of partitions annotated. + """ + count = 0 + for conv in list(gm.graph.nodes): + if conv.op != "call_function" or conv.target not in CONV_TARGETS: + continue + if conv.name in discard_nodes: + continue + + bn = _sole_user(conv) + if bn is None or bn.target not in BN_TARGETS: + continue + bn_out = _bn_output(bn) + if bn_out is None: + continue + + partition = [conv, bn] if bn_out is bn else [conv, bn, bn_out] + output_node = bn_out + + if _is_annotated(partition): + continue + quantization_config = get_quant_config(conv) + if quantization_config is None: + continue + + _annotate_partition(gm, conv, partition, output_node, quantization_config) + count += 1 + + if count: + gm.recompile() + return count diff --git a/backends/qualcomm/quantizer/quantizer.py b/backends/qualcomm/quantizer/quantizer.py index 6eac3fe7e79..d7ac219d927 100644 --- a/backends/qualcomm/quantizer/quantizer.py +++ b/backends/qualcomm/quantizer/quantizer.py @@ -47,6 +47,8 @@ from torchao.quantization.pt2e import UniformQuantizationObserverBase from torchao.quantization.pt2e.quantizer import Quantizer, SharedQuantizationSpec +from .conv_bn import annotate_conv_bn_partitions + from .qconfig import ( get_16a16w_qnn_ptq_config, get_16a2w_qnn_ptq_config, @@ -513,6 +515,13 @@ def annotate(self, model: GraphModule) -> GraphModule: if self._recipe: self._recipe.annotate(model, self._rules_map) else: + if self.default_quant_config.is_qat: + # Conv+BN has to be claimed as one partition before the per-node + # pass. PTQ is left alone: batchnorm is already folded into the + # conv by the time it is quantized there. + annotate_conv_bn_partitions( + model, self._get_quant_config, self.discard_nodes + ) self._annotate(model) self._annotate_custom_annotation(model) diff --git a/backends/qualcomm/tests/BUCK b/backends/qualcomm/tests/BUCK index e572a88b9c4..19b45a02034 100644 --- a/backends/qualcomm/tests/BUCK +++ b/backends/qualcomm/tests/BUCK @@ -59,6 +59,27 @@ fbcode_target(_kind = runtime.python_library, ] ) +fbcode_target(_kind = runtime.python_test, + name = "test_conv_bn_qat", + srcs = [ + "test_conv_bn_qat.py", + ], + env = {} if runtime.is_oss else { + "LD_LIBRARY_PATH": "$(location fbsource//third-party/qualcomm/qnn/qnn-{0}:qnn_offline_compile_libs)".format(get_qnn_library_version()), + "QNN_SDK_ROOT": "$(location fbsource//third-party/qualcomm/qnn/qnn-{0}:__dir__)".format(get_qnn_library_version()), + }, + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/backends/qualcomm/_passes:passes", + "//executorch/backends/qualcomm/builders:builders", + "//executorch/backends/qualcomm/partition:partition", + "//executorch/backends/qualcomm/quantizer:quantizer", + "//executorch/backends/qualcomm/serialization:serialization", + "//executorch/backends/qualcomm/utils:utils", + ], +) + fbcode_target(_kind = runtime.python_test, name = "test_passes", srcs = [ diff --git a/backends/qualcomm/tests/test_conv_bn_qat.py b/backends/qualcomm/tests/test_conv_bn_qat.py new file mode 100644 index 00000000000..9f604e55249 --- /dev/null +++ b/backends/qualcomm/tests/test_conv_bn_qat.py @@ -0,0 +1,172 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype +from torch.fx import Node +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e +from torchao.quantization.pt2e.utils import _is_bn_node, _is_conv_or_conv_transpose_node + +_DQ_PER_CHANNEL = torch.ops.quantized_decomposed.dequantize_per_channel.default +_Q_PER_TENSOR = torch.ops.quantized_decomposed.quantize_per_tensor.default + + +class ConvBnNoBias(torch.nn.Module): + """conv(bias=False) -> bn -> relu, the shape that broke OSS QNN QAT.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1, bias=False) + self.bn = torch.nn.BatchNorm2d(8) + self.relu = torch.nn.ReLU() + + def forward(self, x): + return self.relu(self.bn(self.conv(x))) + + +class ConvBnWithBias(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1, bias=True) + self.bn = torch.nn.BatchNorm2d(8) + + def forward(self, x): + return self.bn(self.conv(x)) + + +class ConvNoBn(torch.nn.Module): + """A conv with no batchnorm must be left to the per-node annotator.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1, bias=True) + + def forward(self, x): + return torch.relu(self.conv(x)) + + +class DepthwiseSeparableConvBn(torch.nn.Module): + """The `conv_dw` block: two bias-free conv+bn+relu pairs back to back.""" + + def __init__(self): + super().__init__() + self.dw = torch.nn.Conv2d(8, 8, 3, padding=1, groups=8, bias=False) + self.bn1 = torch.nn.BatchNorm2d(8) + self.pw = torch.nn.Conv2d(8, 16, 1, bias=False) + self.bn2 = torch.nn.BatchNorm2d(16) + + def forward(self, x): + x = torch.relu(self.bn1(self.dw(x))) + return torch.relu(self.bn2(self.pw(x))) + + +def _qat_convert(module: torch.nn.Module, example_inputs): + quantizer = QnnQuantizer() + quantizer.set_default_quant_config( + QuantDtype.use_8a8w, is_qat=True, is_conv_per_channel=True + ) + exported = torch.export.export(module, example_inputs, strict=True).module() + prepared = prepare_qat_pt2e(exported, quantizer) + prepared(*example_inputs) + return convert_pt2e(prepared) + + +class TestConvBnQat(unittest.TestCase): + """The two defects behind `0x7532` on the OSS QNN QAT export. + + 1. batchnorm never folded, because an observer on the conv output landed + inside torchao's QAT conv-bn fusion and stopped `_fold_conv_bn_qat` from + matching; + 2. the folded bias was left unquantized on convs declared `bias=False`, + which QNN rejects on a per-channel-quantized conv. + """ + + def _assert_no_batch_norm(self, converted): + surviving = [n for n in converted.graph.nodes if _is_bn_node(n)] + self.assertEqual(surviving, [], f"batchnorm survived convert_pt2e: {surviving}") + + def _assert_conv_outputs_quantized(self, converted): + """QNN lowers conv and relu to separate ops, so the conv's own output has + to be quantized. Putting the partition's output qspec on a trailing relu + instead leaves the conv emitting float32 and QNN rejects it with + `0x232 != 0x408`. + """ + convs = [n for n in converted.graph.nodes if _is_conv_or_conv_transpose_node(n)] + self.assertGreater(len(convs), 0) + for conv in convs: + users = list(conv.users) + self.assertEqual( + [u.target for u in users], + [_Q_PER_TENSOR] * len(users), + f"{conv.name} output is not quantized; its users are " + f"{[(u.op, u.target) for u in users]}", + ) + + def _assert_conv_biases_quantized(self, converted): + convs = [n for n in converted.graph.nodes if _is_conv_or_conv_transpose_node(n)] + self.assertGreater(len(convs), 0) + for conv in convs: + if len(conv.args) < 3: + continue + bias = conv.args[2] + if not isinstance(bias, Node): + continue + self.assertEqual( + bias.target, + _DQ_PER_CHANNEL, + f"{conv.name} bias is {bias.op}:{bias.target}, expected a " + "dequantize_per_channel; QNN rejects a float bias on a " + "per-channel-quantized conv", + ) + + def test_conv_bn_relu_without_bias_folds_and_quantizes_bias(self): + example_inputs = (torch.randn(1, 3, 16, 16),) + converted = _qat_convert(ConvBnNoBias().eval(), example_inputs) + self._assert_no_batch_norm(converted) + self._assert_conv_biases_quantized(converted) + self._assert_conv_outputs_quantized(converted) + + def test_conv_bn_with_bias_folds(self): + example_inputs = (torch.randn(1, 3, 16, 16),) + converted = _qat_convert(ConvBnWithBias().eval(), example_inputs) + self._assert_no_batch_norm(converted) + self._assert_conv_biases_quantized(converted) + self._assert_conv_outputs_quantized(converted) + + def test_depthwise_separable_conv_bn_folds(self): + example_inputs = (torch.randn(1, 8, 16, 16),) + converted = _qat_convert(DepthwiseSeparableConvBn().eval(), example_inputs) + self._assert_no_batch_norm(converted) + self._assert_conv_biases_quantized(converted) + self._assert_conv_outputs_quantized(converted) + + def test_conv_without_bn_is_still_annotated(self): + """The partition pass must not starve the per-node conv annotator.""" + example_inputs = (torch.randn(1, 3, 16, 16),) + converted = _qat_convert(ConvNoBn().eval(), example_inputs) + convs = [n for n in converted.graph.nodes if _is_conv_or_conv_transpose_node(n)] + self.assertEqual(len(convs), 1) + self.assertEqual(convs[0].args[1].target, _DQ_PER_CHANNEL) + + def test_ptq_annotation_is_unchanged(self): + """PTQ already worked; the partition pass is gated off for it.""" + from torchao.quantization.pt2e.quantize_pt2e import prepare_pt2e + + example_inputs = (torch.randn(1, 3, 16, 16),) + quantizer = QnnQuantizer() + quantizer.set_default_quant_config( + QuantDtype.use_8a8w, is_qat=False, is_conv_per_channel=True + ) + exported = torch.export.export( + ConvBnNoBias().eval(), example_inputs, strict=True + ).module() + prepared = prepare_pt2e(exported, quantizer) + prepared(*example_inputs) + converted = convert_pt2e(prepared) + # PTQ folds batchnorm outside the quantizer; assert we did not regress it. + self._assert_no_batch_norm(converted)