diff --git a/backends/arm/operators/ops_quant_utils.py b/backends/arm/operators/ops_quant_utils.py index fef9e3c56c1..25e35f92b20 100644 --- a/backends/arm/operators/ops_quant_utils.py +++ b/backends/arm/operators/ops_quant_utils.py @@ -20,7 +20,15 @@ def add_input_weight_zp_consts(tosa_graph, node, inputs, output_name): weight_zp = 0 if inputs[1].dtype == ts.DType.INT8: input_qparams = get_input_qparams(node) - weight_zp = input_qparams[1].zp # type: ignore[assignment] + zero_points = input_qparams[1].zp + if isinstance(zero_points, list): + if not zero_points or any(zp != zero_points[0] for zp in zero_points): + raise ValueError( + "TOSA convolution requires one common weight zero point" + ) + weight_zp = zero_points[0] + else: + weight_zp = zero_points input_zp_name = f"{output_name}_input_zp" weight_zp_name = f"{output_name}_weight_zp" @@ -29,7 +37,7 @@ def add_input_weight_zp_consts(tosa_graph, node, inputs, output_name): [1], inputs[0].dtype, [input_zp], name=input_zp_name ) weight_zp_tensor = tosa_graph.addConst( - [1], inputs[1].dtype, weight_zp, name=weight_zp_name + [1], inputs[1].dtype, [weight_zp], name=weight_zp_name ) return input_zp_tensor.name, weight_zp_tensor.name diff --git a/backends/arm/quantizer/quantization_annotator.py b/backends/arm/quantizer/quantization_annotator.py index 70f87267c65..020101a2718 100644 --- a/backends/arm/quantizer/quantization_annotator.py +++ b/backends/arm/quantizer/quantization_annotator.py @@ -129,7 +129,8 @@ def _adjust_weight_qspec_for_conv_transpose( if ( node.target != torch.ops.aten.conv_transpose2d.input or weight_qspec is None - or weight_qspec.qscheme != torch.per_channel_symmetric + or weight_qspec.qscheme + not in (torch.per_channel_symmetric, torch.per_channel_affine) ): return weight_qspec diff --git a/backends/arm/quantizer/quantization_config.py b/backends/arm/quantizer/quantization_config.py index 3cfc49c1dbe..297da259cf3 100644 --- a/backends/arm/quantizer/quantization_config.py +++ b/backends/arm/quantizer/quantization_config.py @@ -232,14 +232,13 @@ def get_weight_qspec( ) -> Optional[QuantizationSpecBase]: """Get the validated weight spec. - Validate that the weight qscheme is supported (per-tensor or - per-channel symmetric) before returning the spec. + Validate that the weight qscheme is supported before returning the spec. Returns: Optional[QuantizationSpecBase]: Weight spec, or ``None`` when unset. Raises: - ValueError: If the qscheme is not a supported symmetric scheme. + ValueError: If the qscheme is not a supported static scheme. """ if self.weight is None: @@ -248,6 +247,8 @@ def get_weight_qspec( if not hasattr(self.weight, "qscheme") or self.weight.qscheme not in [ torch.per_tensor_symmetric, torch.per_channel_symmetric, + torch.per_tensor_affine, + torch.per_channel_affine, ]: raise ValueError(f"Unsupported quantization_spec {self.weight} for weight") return self.weight @@ -354,7 +355,10 @@ def _derive_qparams_fn( ) ch_axis = None if self.weight is not None: - if qscheme == torch.per_channel_symmetric: + if qscheme in ( + torch.per_channel_symmetric, + torch.per_channel_affine, + ): ch_axis = self.weight.ch_axis if ( node.target == torch.ops.aten.conv_transpose2d.input @@ -466,7 +470,8 @@ def get_weight_qspec( node is not None and weight_qspec is not None and isinstance(weight_qspec, QuantizationSpec) - and weight_qspec.qscheme == torch.per_channel_symmetric + and weight_qspec.qscheme + in (torch.per_channel_symmetric, torch.per_channel_affine) and node.target == torch.ops.aten.conv_transpose2d.input ): # MLETORCH-1853: Fix lazy import when moving files around diff --git a/backends/arm/test/quantizer/affine_weights/BUCK b/backends/arm/test/quantizer/affine_weights/BUCK new file mode 100644 index 00000000000..54c78b6527e --- /dev/null +++ b/backends/arm/test/quantizer/affine_weights/BUCK @@ -0,0 +1,25 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = python_pytest, + name = "test_affine_weights", + srcs = ["test_affine_weights.py"], + env = {"TEST_RUNTIME_IS_NOT_OSS": "1" if not runtime.is_oss else "0"}, + preload_deps = ["//executorch/kernels/quantized:custom_ops_generated_lib"], + deps = [ + "//caffe2:torch", + "//executorch/backends/arm/_passes:passes", + "//executorch/backends/arm/operators:ops", + "//executorch/backends/arm/quantizer:lib", + "//executorch/backends/arm/test:arm_tester" if runtime.is_oss else "//executorch/backends/arm/test/tester/fb:arm_tester_fb", + "//executorch/backends/test/harness:tester", + "//pytorch/ao:torchao", + "fbsource//third-party/pypi/pytest:pytest", + "fbsource//third-party/tosa_tools:serializer", + "fbsource//third-party/tosa_tools:tosa_reference_model", + ], +) diff --git a/backends/arm/test/quantizer/affine_weights/test_affine_weights.py b/backends/arm/test/quantizer/affine_weights/test_affine_weights.py new file mode 100644 index 00000000000..4a004a8fcbd --- /dev/null +++ b/backends/arm/test/quantizer/affine_weights/test_affine_weights.py @@ -0,0 +1,169 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# 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. + +from dataclasses import replace +from unittest.mock import Mock + +import pytest +import torch +import tosa_serializer as ts +from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm.operators.ops_quant_utils import add_input_weight_zp_consts +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_quantization_config, +) +from executorch.backends.arm.quantizer.quantization_config import QuantizationConfig +from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineINT +from executorch.backends.test.harness.stages import StageType +from torchao.quantization.pt2e.quantizer import ( + DerivedQuantizationSpec, + QuantizationSpec, +) + + +def _affine_config(per_channel: bool) -> QuantizationConfig: + config = get_symmetric_quantization_config(is_per_channel=per_channel) + weight = config.weight + assert isinstance(weight, QuantizationSpec) + return replace( + config, + weight=replace( + weight, + qscheme=( + torch.per_channel_affine if per_channel else torch.per_tensor_affine + ), + quant_min=-128, + ), + ) + + +@pytest.mark.parametrize("per_channel", [False, True]) +def test_affine_weight_specs_and_derived_channel_bias(per_channel: bool) -> None: + config = _affine_config(per_channel) + assert config.get_weight_qspec() is config.weight + graph = torch.fx.Graph() + node = graph.call_function( + torch.ops.aten.linear.default, + ( + graph.placeholder("input"), + graph.placeholder("weight"), + graph.placeholder("bias"), + ), + ) + bias = config.get_bias_qspec(node) + assert isinstance(bias, DerivedQuantizationSpec) + assert bias.ch_axis == (0 if per_channel else None) + + +def test_equal_channel_weight_zero_points_serialize_as_one_scalar() -> None: + graph = torch.fx.Graph() + node = graph.placeholder("conv") + node.meta["input_qparams"] = { + 0: QuantArgs(0.1, 0, -127, 127, torch.int8), + 1: QuantArgs([0.01, 0.02], [-28, -28], -128, 127, torch.int8, per_channel=True), + } + inputs = [Mock(dtype=ts.DType.INT8), Mock(dtype=ts.DType.INT8)] + serializer = Mock() + input_const, weight_const = Mock(), Mock() + input_const.name, weight_const.name = "input_zp", "weight_zp" + serializer.addConst.side_effect = [input_const, weight_const] + add_input_weight_zp_consts(serializer, node, inputs, "conv") + assert serializer.addConst.call_args_list[1].args == ([1], ts.DType.INT8, [-28]) + node.meta["input_qparams"][1] = QuantArgs( + [0.01, 0.02], [-28, -27], -128, 127, torch.int8, per_channel=True + ) + with pytest.raises(ValueError, match="common weight zero point"): + add_input_weight_zp_consts(serializer, node, inputs, "conv") + + +def _affine_model( + operator: str, conflicting_zero_points: bool = False +) -> tuple[torch.nn.Module, tuple[torch.Tensor]]: + model: torch.nn.Linear | torch.nn.Conv2d | torch.nn.ConvTranspose2d + if operator == "linear": + model = torch.nn.Linear(8, 3) + inputs = (torch.linspace(-1, 1, 8).reshape(1, 8),) + else: + model = ( + torch.nn.ConvTranspose2d(2, 3, 2) + if operator == "conv_transpose2d" + else torch.nn.Conv2d(2, 3, 2) + ) + inputs = (torch.linspace(-1, 1, 72).reshape(1, 2, 6, 6),) + axis = 1 if operator == "conv_transpose2d" else 0 + channel_shape = model.weight.movedim(axis, 0).shape + base = torch.linspace(-0.25, 1, model.weight.numel() // 3).reshape( + channel_shape[1:] + ) + weights = torch.stack([base * scale for scale in (1, 2, 3)]) + if conflicting_zero_points: + weights[0] += 0.25 + with torch.no_grad(): + model.weight.copy_(weights.movedim(0, axis)) + assert model.bias is not None + model.bias.copy_(torch.tensor([0.1, 0.2, 0.3])) + return model.eval(), inputs + + +def _check_affine_qparams( + pipeline: TosaPipelineINT[tuple[torch.Tensor]], per_channel: bool +) -> None: + graph_module = pipeline.tester.get_artifact(StageType.QUANTIZE) + op = ( + torch.ops.quantized_decomposed.dequantize_per_channel.default + if per_channel + else torch.ops.quantized_decomposed.dequantize_per_tensor.default + ) + weights = [ + node + for node in graph_module.graph.nodes + if node.target == op and node.args[-1] == torch.int8 + ] + if per_channel: + assert len(weights) == 1 + scales = graph_module.get_buffer(str(weights[0].args[1].target)) + zero_points = graph_module.get_buffer(str(weights[0].args[2].target)) + assert scales.numel() == 3 + assert torch.unique(scales).numel() == 3 + assert torch.unique(zero_points).numel() == 1 + assert zero_points[0] != 0 + biases = [ + node + for node in graph_module.graph.nodes + if node.target == op and node.args[-1] == torch.int32 + ] + assert len(biases) == 1 + assert biases[0].args[3] == 0 + assert graph_module.get_buffer(str(biases[0].args[1].target)).numel() == 3 + else: + weight = next(node for node in weights if node.args[0].op == "get_attr") + assert weight.args[2] != 0 + + +@pytest.mark.parametrize("operator", ["linear", "conv2d", "conv_transpose2d"]) +@pytest.mark.parametrize("per_channel", [False, True]) +def test_affine_weights_tosa_pipeline(operator: str, per_channel: bool) -> None: + model, inputs = _affine_model(operator) + pipeline = TosaPipelineINT[tuple[torch.Tensor]]( + model, + inputs, + aten_op=f"torch.ops.aten.{operator}.{'input' if operator == 'conv_transpose2d' else 'default'}", + per_channel_quantization=per_channel, + ) + pipeline.quantizer.set_global(_affine_config(per_channel)) + pipeline.add_stage_after("quantize", _check_affine_qparams, pipeline, per_channel) + pipeline.run() + assert pipeline.has_stage("run_method_and_compare_outputs") + + +def test_conflicting_affine_weight_zero_points_rejected_by_pipeline() -> None: + model, inputs = _affine_model("conv2d", conflicting_zero_points=True) + pipeline = TosaPipelineINT[tuple[torch.Tensor]]( + model, inputs, aten_op="torch.ops.aten.conv2d.default" + ) + pipeline.quantizer.set_global(_affine_config(True)) + with pytest.raises(ValueError, match="common weight zero point"): + pipeline.run()