From d4e9d752e23bf30ce97efc733802dc7b0adabab0 Mon Sep 17 00:00:00 2001 From: Min Guo Date: Thu, 27 Aug 2026 15:39:32 -0700 Subject: [PATCH] Only claim conv overloads the QAT fold can actually match Summary: `annotate_conv_bn_partitions` listed `aten.conv2d.padding` and `aten.convolution.default` in `CONV_TARGETS`, but torchao can never fold either. `_fuse_conv_bn_qat` and `_fold_conv_bn_qat` both trace their patterns from `F.conv1d`, `F.conv2d`, `F.conv_transpose1d` and `F.conv_transpose2d`, so only those four overloads can match. Claiming a conv the fold cannot match is worse than not claiming it. The pass drops the conv's output qspec on the assumption the fold will move the bn output observer onto the conv; when the fold misses, the conv is left emitting float32 into a quantized graph and QNN rejects it with `0x232 != 0x408` (FLOAT_32 meeting UFIXED_POINT_8) -- the very failure the module docstring describes. `nn.Conv2d(..., padding="same")` exports to `aten.conv2d.padding` and hit exactly this. `aten.convolution.default` only appears after decomposition, long after this pass runs, so it never matched anything. Measured on master, conv+bn+relu through `prepare_qat_pt2e` -> `convert_pt2e` with `QnnQuantizer(8a8w, is_qat=True)`: | model | exported target | in CONV_TARGETS | bn survives | conv output unquantized | | --- | --- | --- | --- | --- | | `Conv2d(padding=1)` | `conv2d.default` | yes | no | 0 | | `Conv2d(padding="same")` | `conv2d.padding` | **yes** | yes | **1** | | `Conv1d(padding="same")` | `conv1d.padding` | no | yes | 0 | | `ConvTranspose2d` | `conv_transpose2d.input` | yes | no | 0 | `conv1d.padding` is the control: it also fails to fold -- that is a torchao limitation this pass cannot fix -- but because it was never claimed, the per-node annotator still gives the conv an output qspec and the graph stays valid. Only the overload this pass claimed lost its output. After narrowing `CONV_TARGETS`, the `conv2d.padding` case reports 0 unquantized conv outputs. Note `_is_conv_or_conv_transpose_node` is not a usable guard here -- it returns True for `conv2d.padding` (`torchao/quantization/pt2e/utils.py`, `_is_conv_node`). Differential Revision: D117778197 --- backends/qualcomm/quantizer/conv_bn.py | 18 +++++- backends/qualcomm/tests/test_conv_bn_qat.py | 66 +++++++++++++++++++++ 2 files changed, 82 insertions(+), 2 deletions(-) diff --git a/backends/qualcomm/quantizer/conv_bn.py b/backends/qualcomm/quantizer/conv_bn.py index 4fc7875ce6c..f54ac577131 100644 --- a/backends/qualcomm/quantizer/conv_bn.py +++ b/backends/qualcomm/quantizer/conv_bn.py @@ -47,11 +47,25 @@ from .qconfig import QuantizationConfig from .rules import _is_annotated, _mark_nodes_as_annotated, Q_ANNOTATION_KEY +# Only the overloads torchao's `_fuse_conv_bn_qat` / `_fold_conv_bn_qat` can +# actually match. Both trace their patterns from `F.conv1d`, `F.conv2d`, +# `F.conv_transpose1d` and `F.conv_transpose2d`, which export to exactly these +# four targets. +# +# Claiming a conv the fold cannot match is worse than not claiming it: this pass +# drops the conv's output qspec on the assumption the fold will move the bn +# output observer onto it, so if the fold misses, the conv is left emitting +# float32 into a quantized graph and QNN rejects it (`0x232 != 0x408`, FLOAT_32 +# meeting UFIXED_POINT_8). `aten.conv2d.padding` -- produced by +# `nn.Conv2d(..., padding="same")` -- hit exactly that. `aten.convolution.default` +# only appears after decomposition, well after this pass runs, so it never +# matched anything here either. +# +# Note `_is_conv_or_conv_transpose_node` is NOT a usable guard: it returns True +# for `conv2d.padding`. 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, ) diff --git a/backends/qualcomm/tests/test_conv_bn_qat.py b/backends/qualcomm/tests/test_conv_bn_qat.py index 9f604e55249..1b17e789c60 100644 --- a/backends/qualcomm/tests/test_conv_bn_qat.py +++ b/backends/qualcomm/tests/test_conv_bn_qat.py @@ -65,6 +65,41 @@ def forward(self, x): return torch.relu(self.bn2(self.pw(x))) +class ConvBnSamePadding(torch.nn.Module): + """`padding="same"` exports to `aten.conv2d.padding`, which torchao cannot fold.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding="same", 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 Conv1dBn(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv1d(3, 8, 3, padding=1, bias=False) + self.bn = torch.nn.BatchNorm1d(8) + self.relu = torch.nn.ReLU() + + def forward(self, x): + return self.relu(self.bn(self.conv(x))) + + +class ConvTranspose2dBn(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.ConvTranspose2d(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))) + + def _qat_convert(module: torch.nn.Module, example_inputs): quantizer = QnnQuantizer() quantizer.set_default_quant_config( @@ -145,6 +180,37 @@ def test_depthwise_separable_conv_bn_folds(self): self._assert_conv_biases_quantized(converted) self._assert_conv_outputs_quantized(converted) + def test_conv1d_bn_folds(self): + example_inputs = (torch.randn(1, 3, 16),) + converted = _qat_convert(Conv1dBn().eval(), example_inputs) + self._assert_no_batch_norm(converted) + self._assert_conv_outputs_quantized(converted) + + def test_conv_transpose2d_bn_folds(self): + example_inputs = (torch.randn(1, 3, 16, 16),) + converted = _qat_convert(ConvTranspose2dBn().eval(), example_inputs) + self._assert_no_batch_norm(converted) + self._assert_conv_outputs_quantized(converted) + + def test_unfoldable_conv_overload_keeps_quantized_output(self): + """A conv the fold cannot match must keep its own output qspec. + + `padding="same"` exports to `aten.conv2d.padding`. torchao's + `_fuse_conv_bn_qat` / `_fold_conv_bn_qat` trace their patterns from + `F.conv1d`/`F.conv2d`/`F.conv_transpose{1,2}d`, so only those four + overloads can ever match and BatchNorm survives here. + + That is a torchao limitation this pass cannot fix, but it must not make it + worse: if the partition pass claims the conv it drops the conv's output + qspec expecting the fold to move the bn output observer onto it, and the + conv is left emitting float32 into a quantized graph -- which QNN rejects + with `0x232 != 0x408`. So the conv must stay quantized even when the fold + misses. + """ + example_inputs = (torch.randn(1, 3, 16, 16),) + converted = _qat_convert(ConvBnSamePadding().eval(), example_inputs) + 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),)