From 0cc88fa0f991548cc86116b1068a16d5fe0fe577 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Wed, 26 Aug 2026 23:43:35 -0700 Subject: [PATCH] Qualcomm: let the 8a8w QAT activation spec share an observer with per-channel QNN carries two activation specs: the default config for most ops and the per-channel config for conv and linear. PT2E gives an edge one observer only when the producer's output qspec and the consumer's input qspec agree on every attribute _union_input_edge_with compares -- dtype, is_dynamic, quant_min, quant_max, qscheme, ch_axis, scale, zero_point. observer_or_fake_quant_ctr is deliberately excluded, so this is about attribute values, not spec identity or SharedQuantizationSpec. get_8a8w_qnn_qat_config built its activation spec unconditionally in the symmetric shape: ch_axis=0, no quant_min/quant_max, with only qscheme switching on act_symmetric. get_qat_per_channel_quant_config branches properly, so on the default asymmetric path the two disagree on quant_min, quant_max and ch_axis. Every per-channel conv output meeting a default-config consumer then gets two observers instead of one, which convert_pt2e materializes as a dequantize immediately followed by a requantize. On a conv-relu-conv model with no BatchNorm, QAT produced three such clusters where PTQ produced none. Split on act_symmetric the way get_8a8w_qnn_ptq_config already does. Passing explicit uint8 bounds for per_tensor_affine is a numerical no-op -- the observer already defaults to 0/255, and scale, zero_point and fake-quant output are bit-identical either way -- it only makes the attributes visible to the comparison. Two tests. The first compares the default and per-channel activation specs across every QUANT_CONFIG_DICT entry using PT2E's own _has_same_attr; it carries a _KNOWN_UNSHAREABLE list for the eight 16-bit act_symmetric combinations, where the default configs set quant_min/quant_max unconditionally while the per-channel config omits them. Repairing those changes the observed zero_point on a path that works today, so it is left for a follow-up; the list asserts they still mismatch, so it fails if anything shifts. The second test is the end-to-end guard, counting quantize nodes fed directly by a dequantize. Both fail without the fix: 3 != 0, and mismatched on ['quant_min', 'quant_max', 'ch_axis']. test_passes.py already runs in OSS CI by name, so no workflow change is needed. Authored with assistance from Claude Code. --- backends/qualcomm/quantizer/qconfig.py | 50 +++++--- backends/qualcomm/tests/test_passes.py | 158 +++++++++++++++++++++++++ 2 files changed, 193 insertions(+), 15 deletions(-) diff --git a/backends/qualcomm/quantizer/qconfig.py b/backends/qualcomm/quantizer/qconfig.py index 2f8ce1c972e..5cf00e1474f 100644 --- a/backends/qualcomm/quantizer/qconfig.py +++ b/backends/qualcomm/quantizer/qconfig.py @@ -813,21 +813,41 @@ def get_8a8w_qnn_qat_config( ) -> QuantizationConfig: # the smallest scale defaults to DEFAULT_EPS_8BIT extra_args: Dict[str, Any] = {"eps": eps if eps else DEFAULT_EPS_8BIT} - act_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( - dtype=torch.uint8, - qscheme=( - torch.per_tensor_symmetric if act_symmetric else torch.per_tensor_affine - ), - observer=act_observer.with_args(**extra_args), - ) - act_quantization_spec = QuantizationSpec( - dtype=torch.uint8, - qscheme=( - torch.per_tensor_symmetric if act_symmetric else torch.per_tensor_affine - ), - ch_axis=0, - observer_or_fake_quant_ctr=act_fake_quant_ctr, - ) + # These must match get_qat_per_channel_quant_config attribute for attribute. + # PT2E merges the observer on a per-channel conv output with the observer on + # its consumer's input edge only when dtype, quant_min, quant_max, qscheme and + # ch_axis all agree (_union_input_edge_with in torchao pt2e/prepare.py). Any + # drift leaves two observers on one edge, i.e. a redundant dequantize-requantize + # at every conv boundary. + if act_symmetric: + # If we keep quant_min and quant_max none, observer will default use 128 as zero_point. + # If we provide uint8 quant_min/max, it will use 127 as zero_point, which is undesired. + act_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( + dtype=torch.uint8, + qscheme=torch.per_tensor_symmetric, + observer=act_observer.with_args(**extra_args), + ) + act_quantization_spec = QuantizationSpec( + dtype=torch.uint8, + qscheme=torch.per_tensor_symmetric, + ch_axis=0, + observer_or_fake_quant_ctr=act_fake_quant_ctr, + ) + else: + act_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( + dtype=torch.uint8, + quant_min=torch.iinfo(torch.uint8).min, + quant_max=torch.iinfo(torch.uint8).max, + qscheme=torch.per_tensor_affine, + observer=act_observer.with_args(**extra_args), + ) + act_quantization_spec = QuantizationSpec( + dtype=torch.uint8, + quant_min=torch.iinfo(torch.uint8).min, + quant_max=torch.iinfo(torch.uint8).max, + qscheme=torch.per_tensor_affine, + observer_or_fake_quant_ctr=act_fake_quant_ctr, + ) weight_fake_quant_ctr = FusedMovingAvgObsFakeQuantize.with_args( dtype=torch.int8, diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 9085b4eb7ed..5a6d000d198 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -848,5 +848,163 @@ def test_backend_bundle_cache_survives_an_expired_entry(self): self.assertGreaterEqual(reused, 1, "third manager must reuse the live bundle") +# Attributes PT2E compares in _union_input_edge_with when deciding whether an +# input edge and its producer's output can share one observer. +_SHARING_ATTRS = ( + "dtype", + "is_dynamic", + "quant_min", + "quant_max", + "qscheme", + "ch_axis", + "scale", + "zero_point", +) + +# (quant_dtype, is_qat, act_symmetric) combinations whose default and per-channel +# activation specs are known not to match yet. Every 16-bit config sets +# quant_min/quant_max unconditionally while the per-channel config omits them when +# act_symmetric, so the two cannot share an observer. Fixing that changes the +# observed zero_point on a path that works today, so it is left for a follow-up. +_KNOWN_UNSHAREABLE = { + (QuantDtype.use_16a16w, False, True), + (QuantDtype.use_16a2w, False, True), + (QuantDtype.use_16a4w, False, True), + (QuantDtype.use_16a4w, True, True), + (QuantDtype.use_16a4w_block, False, True), + (QuantDtype.use_16a4w_block, True, True), + (QuantDtype.use_16a8w, False, True), + (QuantDtype.use_16a8w, True, True), +} + +_Q_PER_TENSOR = { + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.quantize_per_tensor.tensor, +} +_DQ_PER_TENSOR = { + torch.ops.quantized_decomposed.dequantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.tensor, +} + + +def _count_requantize_clusters(graph_module): + """Quantize nodes fed directly by a dequantize: two observers on one edge.""" + return sum( + 1 + for node in graph_module.graph.nodes + if node.op == "call_function" + and node.target in _Q_PER_TENSOR + and isinstance(node.args[0], torch.fx.Node) + and node.args[0].target in _DQ_PER_TENSOR + ) + + +class ConvReluConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 8, 3, padding=1) + self.conv2 = torch.nn.Conv2d(8, 8, 3, padding=1) + + def forward(self, x): + return torch.relu(self.conv2(torch.relu(self.conv1(x)))) + + +class TestActivationSpecSharing(unittest.TestCase): + """A per-channel op's activation spec must match the default one. + + QNN uses two activation specs: the default config for most ops and the + per-channel config for conv and linear. PT2E gives an edge one observer only + when the producer's output qspec and the consumer's input qspec agree on every + attribute in _SHARING_ATTRS -- observer_or_fake_quant_ctr is deliberately not + compared, so this is about attributes, not spec identity or + SharedQuantizationSpec. Any drift leaves two observers on one edge, which + convert_pt2e materializes as a dequantize immediately followed by a quantize + at every conv boundary. + """ + + def test_default_and_per_channel_act_specs_can_share_an_observer(self): + from executorch.backends.qualcomm.quantizer.quantizer import QUANT_CONFIG_DICT + from torchao.quantization.pt2e.prepare import _has_same_attr + + for (quant_dtype, is_qat), funcs in QUANT_CONFIG_DICT.items(): + default_fn, per_channel_fn = funcs[0], funcs[1] + if default_fn is None or per_channel_fn is None: + continue + for act_symmetric in (False, True): + combination = (quant_dtype, is_qat, act_symmetric) + with self.subTest(combination=combination): + default_act = default_fn( + act_symmetric=act_symmetric + ).output_activation + per_channel_act = per_channel_fn( + act_symmetric=act_symmetric, ch_axis=0 + ).input_activation + if default_act is None or per_channel_act is None: + # fp16a8w keeps activations in float. + continue + mismatched = [ + attr + for attr in _SHARING_ATTRS + if not _has_same_attr(default_act, per_channel_act, attr) + ] + if combination in _KNOWN_UNSHAREABLE: + self.assertNotEqual( + mismatched, + [], + f"{combination} is listed in _KNOWN_UNSHAREABLE but now " + "matches; remove it from the list", + ) + continue + self.assertEqual( + mismatched, + [], + f"{combination} cannot share an observer, mismatched on " + + ", ".join( + f"{attr}: {getattr(default_act, attr, None)} vs " + f"{getattr(per_channel_act, attr, None)}" + for attr in mismatched + ), + ) + + def test_8a8w_qat_does_not_double_quantize_conv_boundaries(self): + """End-to-end guard for the same invariant. + + This model has no BatchNorm, so it isolates the activation-spec mismatch + from anything conv-bn related. Before the act_symmetric split in + get_8a8w_qnn_qat_config, QAT produced three redundant clusters here while + PTQ produced none. + """ + from torchao.quantization.pt2e.quantize_pt2e import prepare_qat_pt2e + + example_inputs = (torch.randn(2, 3, 8, 8),) + counts = {} + for is_qat in (False, True): + quantizer = QnnQuantizer() + quantizer.set_default_quant_config( + QuantDtype.use_8a8w, + is_qat=is_qat, + is_conv_per_channel=True, + is_linear_per_channel=True, + ) + module = ConvReluConv() + module = module.train() if is_qat else module.eval() + exported = torch.export.export(module, example_inputs, strict=True).module() + prepared = ( + prepare_qat_pt2e(exported, quantizer) + if is_qat + else prepare_pt2e(exported, quantizer) + ) + prepared(*example_inputs) + counts[is_qat] = _count_requantize_clusters(convert_pt2e(prepared)) + + self.assertEqual(counts[False], 0, "PTQ regressed") + self.assertEqual( + counts[True], + counts[False], + "QAT inserts a dequantize->quantize round trip that PTQ does not; the " + "default and per-channel activation specs have drifted apart", + ) + + if __name__ == "__main__": unittest.main()