From 9e67e862422ebff354ea05e765c40507d536e4e3 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Wed, 1 Jul 2026 08:36:54 -0700 Subject: [PATCH 1/3] - Fix output spec adjustment for fixed qparams ops - Update documentation --- changelog.d/180525445.fixed | 1 + docs/src/quantization/advanced.md | 27 ++- .../quantization/_graph/_annotation_utils.py | 160 +++++++------ src/coreai_opt/quantization/spec/factory.py | 45 +++- .../quantization/spec/fake_quantize.py | 9 +- .../quantization/spec/qparams_calculator.py | 8 +- tests/quantization/test_annotation_utils.py | 111 +++++++++ tests/quantization/test_factory.py | 94 +++++++- .../quantization/test_graph_mode_quantizer.py | 215 ++++++++++++++++++ .../test_graph_mode_quantizer_mnist.py | 4 +- 10 files changed, 587 insertions(+), 87 deletions(-) create mode 100644 changelog.d/180525445.fixed create mode 100644 tests/quantization/test_annotation_utils.py diff --git a/changelog.d/180525445.fixed b/changelog.d/180525445.fixed new file mode 100644 index 0000000..dd550ce --- /dev/null +++ b/changelog.d/180525445.fixed @@ -0,0 +1 @@ +Fix setting of qscheme and float_range for fixed output range ops diff --git a/docs/src/quantization/advanced.md b/docs/src/quantization/advanced.md index 447b470..09fbaf8 100644 --- a/docs/src/quantization/advanced.md +++ b/docs/src/quantization/advanced.md @@ -166,17 +166,28 @@ The `qscheme` controls how these bins are distributed around zero, by determinin ## Quantization Defaults for Known-Range Activations -In graph mode, certain ops have known output ranges. For these ops, the user's `qscheme` setting is not respected — the activation is always treated as asymmetric or symmetric depending on the op, regardless of what the user configured. The treatment also differs between `relu` and the `sigmoid` / `tanh` family. For `relu`, only `qscheme` is overridden; `dtype`, scale, and zero point are still derived from the user's spec and calibration data. For `sigmoid` and `tanh`, scale, zero point, **and** `dtype` are pinned to fixed values (always `torch.uint8`, ignoring whatever `dtype` the user configured). +In graph mode, certain activation ops have analytically known output ranges. For these ops, the quantizer overrides the `qscheme` and `float_range` of the qparams calculator at prepare time, regardless of what the user configured. The user's `dtype` is always preserved — these adjustments do not change the number of bits or the signed/unsigned choice. -| Op | Output range | Always treated as | Scale | Zero point | -| --------- | ------------ | ----------------- | --------- | ---------- | -| `relu` | \[0, ∞) | asymmetric | dynamic | dynamic | -| `sigmoid` | [0, 1] | asymmetric | `1 / 256` | `0` | -| `tanh` | [-1, 1] | symmetric | `2 / 256` | `128` | +The scale and zero point values in the table below assume the default `int8` dtype. For other dtypes, the same formulas apply with the appropriate `quant_min` / `quant_max`. -**Relu**: Treated as asymmetric. The user's `qscheme` is ignored, but `dtype`, scale, and zero point are still derived from the user's spec and calibration data. The zero point follows `zero_point = quant_min - round(min_val / scale)`. Since `relu`'s observed min is always `0`, the zero point very commonly ends up near `quant_min` (e.g., `-128` for `int8`). +| Op | Output range | `qscheme` | `float_range` | Scale (int8) | Zero point (int8) | +| ------------- | ------------------- | ---------- | ------------- | ------------ | ----------------- | +| `hardsigmoid` | [0, 1] | asymmetric | (0, 1) | 1 / 255 | −128 | +| `hardtanh` | Depends (see below) | Depends | Depends | Depends | Depends | +| `relu` | \[0, ∞) | asymmetric | (0, None) | dynamic | −128 | +| `relu6` | [0, 6] | asymmetric | (0, 6) | 6 / 255 | −128 | +| `sigmoid` | [0, 1] | asymmetric | (0, 1) | 1 / 255 | −128 | +| `tanh` | [−1, 1] | symmetric | (−1, 1) | 2 / 255 | 0 | -> **Motivation for asymmetric `relu` and `sigmoid`**: Both ops produce non-negative outputs. With symmetric quantization, the zero point sits at the center of the quantized range, placing half the bins in negative territory that these ops never produce. Those bins are effectively wasted — no floating-point value will ever map to them, reducing quantization resolution by half. Asymmetric treatment shifts the zero point toward the edge of the range so all bins cover values the op actually produces. +**Relu**: The lower bound of `float_range` is pinned to 0 and `qscheme` is set to asymmetric. Because the observed minimum is always 0, the zero point is fixed at `quant_min` (−128 for int8) and stays there regardless of calibration data. The upper bound remains `None` (data-driven), so the scale continues to update during calibration. + +**Sigmoid and hardsigmoid**: Both `qscheme` and `float_range` are fully pinned. Scale and zero point are entirely determined by the dtype and the fixed output range — calibration data has no effect on them. + +**Tanh**: `qscheme` (symmetric) and `float_range` (−1, 1) are fully pinned. Scale and zero point are entirely determined by the dtype and range. + +**Hardtanh**: Bounds are read from the op's node arguments at prepare time, so the effective range and qscheme depend on how the op was configured. If `min_val == −max_val` the range is symmetric around zero and `qscheme` is set to symmetric; otherwise `qscheme` is set to asymmetric. Both ends of `float_range` are pinned to the configured bounds. `relu6` is a special case of `hardtanh(0, 6)` and is handled identically. + +> **Motivation for asymmetric treatment**: Symmetric quantization places the zero point at the center of the quantized range. For `relu`, `sigmoid`, and `hardsigmoid`, whose outputs are always non-negative, symmetric quantization places half the bins in negative territory that the op never produces — wasting half the available resolution. Asymmetric quantization shifts the zero point to the edge of the range so that all bins cover values the op actually generates. For `tanh` and symmetric `hardtanh`, the output is centered at zero so both halves of the range are used equally, and symmetric quantization is appropriate. Eager mode does not perform these adjustments — all activations are quantized uniformly using the user-configured spec. diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index fe304ff..d5c66f3 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -19,7 +19,6 @@ from torch.fx.passes.utils.source_matcher_utils import SourcePartition from torchao.quantization.pt2e import WrapperModule, find_sequential_partitions from torchao.quantization.pt2e.quantizer import ( - FixedQParamsQuantizationSpec, QuantizationAnnotation, QuantizationSpec as TorchAOQuantizationSpec, SharedQuantizationSpec as _SharedQuantizationSpec, @@ -46,7 +45,11 @@ _ACTIVATION_SPEC_DICT, _STATE_SPEC_DICT, ) -from coreai_opt.quantization.spec import QuantizationSpec +from coreai_opt.quantization.spec import ( + QuantizationComponentFactory, + QuantizationScheme, + QuantizationSpec, +) from ._annotation_config import AnnotationConfig, AnnotationContext @@ -56,6 +59,22 @@ INPUT_NODE_PREFIX = "input::" PARAM_NODE_PREFIX = "param::" +# Ops that are transparent to quantization range propagation: they don't alter +# the numeric range of their inputs, so we traverse through them when propagating +# adjusted qspecs to child nodes. +_PASSTHROUGH_OP_OVERLOADS: frozenset = frozenset( + { + torch.ops.aten.dropout, + torch.ops.aten.feature_dropout, + torch.ops.aten.permute, + torch.ops.aten.reshape, + torch.ops.aten.squeeze, + torch.ops.aten.transpose, + torch.ops.aten.unsqueeze, + torch.ops.aten.view, + } +) + def _get_aten_graph_module_for_pattern( pattern: Callable, @@ -131,38 +150,27 @@ class OpsListPattern: F.hardsigmoid, ) -_tanh_qspec = FixedQParamsQuantizationSpec( - dtype=torch.uint8, - scale=2.0 / 256.0, - zero_point=128, - quant_min=0, - quant_max=255, - qscheme=torch.per_tensor_symmetric, -) - -_sigmoid_qspec = FixedQParamsQuantizationSpec( - dtype=torch.uint8, - scale=1.0 / 256.0, - zero_point=0, - quant_min=0, - quant_max=255, - qscheme=torch.per_tensor_affine, -) - +# Dictionary mapping ops with known output bounds to (qscheme, float_range). +# float_range elements may be None to leave that side data-driven. _fixed_q_params_ops = { - torch.ops.aten.tanh.default: _tanh_qspec, - torch.ops.aten.tanh_.default: _tanh_qspec, - torch.ops.aten.sigmoid.default: _sigmoid_qspec, - torch.ops.aten.sigmoid_.default: _sigmoid_qspec, - torch.ops.aten.hardsigmoid.default: _sigmoid_qspec, - torch.ops.aten.hardsigmoid_.default: _sigmoid_qspec, + torch.ops.aten.tanh.default: (QuantizationScheme.SYMMETRIC, (-1.0, 1.0)), + torch.ops.aten.tanh_.default: (QuantizationScheme.SYMMETRIC, (-1.0, 1.0)), + torch.ops.aten.sigmoid.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), + torch.ops.aten.sigmoid_.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), + torch.ops.aten.hardsigmoid.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), + torch.ops.aten.hardsigmoid_.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), + # relu: always >= 0, upper bound is data-driven + torch.ops.aten.relu.default: (QuantizationScheme.ASYMMETRIC, (0.0, None)), + torch.ops.aten.relu_.default: (QuantizationScheme.ASYMMETRIC, (0.0, None)), + # relu6: clipped to [0, 6] + torch.ops.aten.relu6.default: (QuantizationScheme.ASYMMETRIC, (0.0, 6.0)), + torch.ops.aten.relu6_.default: (QuantizationScheme.ASYMMETRIC, (0.0, 6.0)), } -_always_affine_ops = ( - torch.ops.aten.relu.default, - torch.ops.aten.relu_.default, - torch.ops.aten.relu6.default, - torch.ops.aten.relu6_.default, +# hardtanh bounds are configurable via node arguments; handled separately. +_hardtanh_ops = ( + torch.ops.aten.hardtanh.default, + torch.ops.aten.hardtanh_.default, ) @@ -207,19 +215,34 @@ def mark_nodes_as_annotated(nodes: Iterable[Node]) -> None: node.meta[Q_ANNOTATION_KEY]._annotated = True -def _propagate_qscheme_to_child_nodes( +def _propagate_adjusted_spec_to_child_nodes( root_node: torch.fx.Node, - qscheme: torch.qscheme, - shared_observer_nodes: set[torch.fx.Node] | None = None, + qscheme: QuantizationScheme | None, + float_range: tuple[float, float] | None, + shared_observer_nodes: set[torch.fx.Node], ) -> None: """ - Given a qscheme, propagate the qscheme to all applicable children. Any input qspecs - which are not shared qspecs will have qschemes updated. The propagation logic + Given a qscheme or float_range, propagate the info to all applicable children. Any input qspecs + which are not shared qspecs will have specs updated. The propagation logic continues downwards through the graph until we encounter a non-shared observer op. """ + # Set of op types for which we want to propagate the updated spec through, even though they + # are not registered ops with quantizers themselves. + # This is a temporary solution. Adding them as SharedObserverPatterns may make sense, but + # additional consideration is needed as to whether it makes sense to have quantizers in between + # multiple shared observer ops. + # To minimize the impact of this change to quantization behavior as a whole, use the below + # set to skip these ops while continuing to traverse through the graph. nodes_to_propagate = [(root_node, user) for user in root_node.users.keys()] while nodes_to_propagate: parent, curr_node = nodes_to_propagate.pop(0) + if ( + curr_node.op == "call_function" + and getattr(curr_node.target, "overloadpacket", None) in _PASSTHROUGH_OP_OVERLOADS + ): + assert curr_node not in shared_observer_nodes + nodes_to_propagate.extend([(curr_node, user) for user in curr_node.users.keys()]) + continue if not is_node_annotated(curr_node): continue curr_input_qspec = curr_node.meta[Q_ANNOTATION_KEY].input_qspec_map.get(parent) @@ -236,10 +259,20 @@ def _propagate_qscheme_to_child_nodes( # dequantize ops inserted. continue if not isinstance(curr_input_qspec, _SharedQuantizationSpec): + ctr = curr_input_qspec.observer_or_fake_quant_ctr + kwargs = {} + if qscheme is not None: + kwargs["qscheme"] = qscheme + if float_range is not None: + kwargs["float_range"] = float_range + if kwargs: + ctr = QuantizationComponentFactory.update_partial_qparams_calculator(ctr, **kwargs) + + # qscheme in TorchAOQuantizationSpec is not read by coreai-opt later on so we omit it. + # Only the qscheme contained within observer_or_fake_quant_ctr matters. adjusted_qspec = TorchAOQuantizationSpec( - observer_or_fake_quant_ctr=curr_input_qspec.observer_or_fake_quant_ctr, + observer_or_fake_quant_ctr=ctr, dtype=curr_input_qspec.dtype, - qscheme=qscheme, quant_min=curr_input_qspec.quant_min, quant_max=curr_input_qspec.quant_max, ) @@ -270,39 +303,30 @@ def adjust_output_qspec_for_qscheme_and_propagate( if qspec is None: return - # ReLU6 activation maps to torch.ops.aten.hardtanh.default with - # min_val = 0 and max_val = 6 - is_always_affine_op = node.target in _always_affine_ops or ( - node.target in [torch.ops.aten.hardtanh.default, torch.ops.aten.hardtanh_.default] - and node.args[1] == 0 # min_val, corresponding to ReLU6 - and node.args[2] == 6 # max_val, corresponding to ReLU6 - ) - - adjusted_qspec = None if node.target in _fixed_q_params_ops: - adjusted_qspec = TorchAOQuantizationSpec( - observer_or_fake_quant_ctr=qspec.observer_or_fake_quant_ctr, - dtype=qspec.dtype, - qscheme=_fixed_q_params_ops[node.target].qscheme, - quant_min=qspec.quant_min, - quant_max=qspec.quant_max, - ) - # FIXME: Because of a bug in PyTorch in function _create_obs_or_fq_from_qspec - # in module torch/ao/quantization/fx/prepare.py which creates a - # FixedQParamsFakeQuantize partial, instead of an instance, we cannot - # actually create FixedQParamsQuantizationSpec - elif is_always_affine_op: - adjusted_qspec = TorchAOQuantizationSpec( - observer_or_fake_quant_ctr=qspec.observer_or_fake_quant_ctr, - dtype=qspec.dtype, - qscheme=torch.per_tensor_affine, - quant_min=qspec.quant_min, - quant_max=qspec.quant_max, + qscheme, float_range = _fixed_q_params_ops[node.target] + elif node.target in _hardtanh_ops: + min_val, max_val = node.args[1], node.args[2] + float_range = (min_val, max_val) + qscheme = ( + QuantizationScheme.SYMMETRIC if min_val == -max_val else QuantizationScheme.ASYMMETRIC ) + else: + return - if adjusted_qspec is not None: - node.meta[Q_ANNOTATION_KEY].output_qspec = adjusted_qspec - _propagate_qscheme_to_child_nodes(node, adjusted_qspec.qscheme, shared_observer_nodes) + ctr = QuantizationComponentFactory.update_partial_qparams_calculator( + qspec.observer_or_fake_quant_ctr, qscheme=qscheme, float_range=float_range + ) + + # qscheme in TorchAOQuantizationSpec is not read by coreai-opt later on so we omit it. + # Only the qscheme contained within observer_or_fake_quant_ctr matters. + node.meta[Q_ANNOTATION_KEY].output_qspec = TorchAOQuantizationSpec( + observer_or_fake_quant_ctr=ctr, + dtype=qspec.dtype, + quant_min=qspec.quant_min, + quant_max=qspec.quant_max, + ) + _propagate_adjusted_spec_to_child_nodes(node, qscheme, float_range, shared_observer_nodes) def _get_weighted_mod_pattern( diff --git a/src/coreai_opt/quantization/spec/factory.py b/src/coreai_opt/quantization/spec/factory.py index 9a449b0..ef023d2 100644 --- a/src/coreai_opt/quantization/spec/factory.py +++ b/src/coreai_opt/quantization/spec/factory.py @@ -5,6 +5,8 @@ from __future__ import annotations +from typing import Any + from coreai_opt._utils.spec_utils import PartialConstructor as _PartialConstructor from coreai_opt.config.spec import ( CompressionComponentFactoryBase, @@ -194,7 +196,6 @@ def create_fake_quantizer( # Standard arguments that all fake quantizers need common_args = { "dtype": spec.dtype, - "qscheme": spec.qscheme, "qformulation": spec.qformulation, "granularity": spec.granularity, "target_dtype": spec.target_dtype, @@ -211,6 +212,47 @@ def create_fake_quantizer( # Create instance with all arguments return spec.fake_quantize_cls(**common_args, **extra_args) + @classmethod + def update_partial_qparams_calculator( + cls, + partial_ctr: _PartialConstructor, + **kwargs: Any, + ) -> _PartialConstructor: + """Return a new PartialConstructor whose qparams_calculator has attributes overridden. + + The replacement wraps the existing ``qparams_calculator`` callable arg so that + each freshly constructed calculator has the given attributes set before it is + returned. All overridden attributes (e.g. ``float_range``, ``qscheme``) are + plain instance attributes on ``QParamsCalculatorBase`` that are read lazily + during ``forward()``, so post-construction mutation is safe as long as no + forward pass has run yet. + + Args: + partial_ctr: The existing fake-quantizer partial to update. + **kwargs: Attribute name/value pairs to set on the calculator instance. + + Returns: + A new PartialConstructor whose qparams_calculator factory applies the + overrides. + """ + old_factory = partial_ctr.callable_args.get("qparams_calculator") + if old_factory is None: + return partial_ctr + + def _new_factory(): + calculator: QParamsCalculatorBase = old_factory() + for attr, value in kwargs.items(): + if not hasattr(calculator, attr): + msg = ( + f"Cannot override unknown attribute '{attr}' on " + f"{type(calculator).__name__}; expected an existing calculator attribute." + ) + raise AttributeError(msg) + setattr(calculator, attr, value) + return calculator + + return partial_ctr.with_callable_args(qparams_calculator=_new_factory) + @classmethod def create_fake_quantizer_partial( cls, spec: QuantizationSpec, quantization_target: CompressionTargetTensor @@ -236,7 +278,6 @@ def create_fake_quantizer_partial( # (excluding qparams_calculator) common_args = { "dtype": spec.dtype, - "qscheme": spec.qscheme, "qformulation": spec.qformulation, "granularity": spec.granularity, "target_dtype": spec.target_dtype, diff --git a/src/coreai_opt/quantization/spec/fake_quantize.py b/src/coreai_opt/quantization/spec/fake_quantize.py index aafd476..1c0bfaa 100644 --- a/src/coreai_opt/quantization/spec/fake_quantize.py +++ b/src/coreai_opt/quantization/spec/fake_quantize.py @@ -28,11 +28,11 @@ from coreai_opt.config.spec import CompressionSimulatorBase, CompressionTargetTensor from coreai_opt.quantization._utils import get_quantization_shapes as _get_quantization_shapes from coreai_opt.quantization.spec.errors import _BlockSizeMismatchError +from coreai_opt.quantization.spec.qscheme import QuantizationScheme from .granularity import QuantizationGranularity from .qformulation import QuantizationFormulation from .qparams_calculator import QParamsCalculatorBase, StatelessQParamsCalculatorBase -from .qscheme import QuantizationScheme __all__ = ["FakeQuantizeImplBase"] @@ -47,7 +47,6 @@ class FakeQuantizeImplBase(CompressionSimulatorBase, FakeQuantizeBase): def __init__( self, dtype: torch.dtype, - qscheme: QuantizationScheme, qformulation: QuantizationFormulation, granularity: QuantizationGranularity, target_dtype: torch.dtype, @@ -60,7 +59,6 @@ def __init__( ): super().__init__() self.dtype = dtype - self.qscheme = qscheme self.qformulation = qformulation self._granularity = granularity self.target_dtype = target_dtype @@ -75,6 +73,11 @@ def __init__( n_bits = _get_n_bits_from_dtype(dtype) self.n_bits = n_bits + @property + def qscheme(self) -> QuantizationScheme: + """The quantization scheme, delegated to the qparams_calculator.""" + return self.qparams_calculator.qscheme + @property def granularity(self) -> QuantizationGranularity: """Getter for granularity.""" diff --git a/src/coreai_opt/quantization/spec/qparams_calculator.py b/src/coreai_opt/quantization/spec/qparams_calculator.py index 0e9faa7..b110292 100644 --- a/src/coreai_opt/quantization/spec/qparams_calculator.py +++ b/src/coreai_opt/quantization/spec/qparams_calculator.py @@ -130,13 +130,17 @@ def _get_min_and_max_val(self, tensor: torch.Tensor) -> tuple[torch.Tensor, torc taken from float_range setting. """ min_val = ( - self._get_tensor_with_granularity_from_scalar(self.float_range[0], tensor) + torch.clamp( + self._get_tensor_with_granularity_from_scalar(self.float_range[0], tensor), max=0 + ) if self.float_range[0] is not None else None ) max_val = ( - self._get_tensor_with_granularity_from_scalar(self.float_range[1], tensor) + torch.clamp( + self._get_tensor_with_granularity_from_scalar(self.float_range[1], tensor), min=0 + ) if self.float_range[1] is not None else None ) diff --git a/tests/quantization/test_annotation_utils.py b/tests/quantization/test_annotation_utils.py new file mode 100644 index 0000000..6f1f05a --- /dev/null +++ b/tests/quantization/test_annotation_utils.py @@ -0,0 +1,111 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-Clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""Tests for _annotation_utils.""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch._ops import OpOverloadPacket + +from coreai_opt.quantization._graph._annotation_utils import _PASSTHROUGH_OP_OVERLOADS + + +def _overloadpackets_in_graph(model: nn.Module, example_inputs: tuple) -> set[OpOverloadPacket]: + """Return the set of OpOverloadPackets for all call_function nodes in the exported graph.""" + exported = torch.export.export(model, example_inputs, strict=False) + return { + node.target.overloadpacket + for node in exported.graph_module.graph.nodes + if node.op == "call_function" and hasattr(node.target, "overloadpacket") + } + + +class _Permute(nn.Module): + def forward(self, x): + return x.permute(0, 2, 1) + + +class _Reshape(nn.Module): + def forward(self, x): + return x.reshape(-1) + + +class _Squeeze(nn.Module): + def forward(self, x): + return x.squeeze(1) + + +class _Transpose(nn.Module): + def forward(self, x): + return x.transpose(0, 1) + + +class _Unsqueeze(nn.Module): + def forward(self, x): + return x.unsqueeze(0) + + +class _View(nn.Module): + def forward(self, x): + return x.view(-1) + + +class _Dropout(nn.Module): + def forward(self, x): + # training=True ensures the op is not folded away at export time + return F.dropout(x, p=0.5, training=True) + + +class _FeatureDropout(nn.Module): + def forward(self, x): + # training=True ensures the op is not folded away at export time + return F.dropout2d(x, p=0.5, training=True) + + +# Maps each member of _PASSTHROUGH_OP_OVERLOADS to a (model, example_inputs) pair +# that exercises it. If an op is added to _PASSTHROUGH_OP_OVERLOADS without a +# corresponding entry here, test_passthrough_ops_all_covered will fail. +_PASSTHROUGH_OP_TEST_CASES: dict[OpOverloadPacket, tuple[nn.Module, tuple]] = { + torch.ops.aten.dropout: ( + _Dropout(), + ( + torch.randn( + 4, + ), + ), + ), + torch.ops.aten.feature_dropout: (_FeatureDropout(), (torch.randn(1, 4, 4, 4),)), + torch.ops.aten.permute: (_Permute(), (torch.randn(1, 4, 8),)), + torch.ops.aten.reshape: (_Reshape(), (torch.randn(1, 4),)), + torch.ops.aten.squeeze: (_Squeeze(), (torch.randn(1, 1, 4),)), + torch.ops.aten.transpose: (_Transpose(), (torch.randn(2, 4),)), + torch.ops.aten.unsqueeze: ( + _Unsqueeze(), + ( + torch.randn( + 4, + ), + ), + ), + torch.ops.aten.view: (_View(), (torch.randn(1, 4),)), +} + + +def test_passthrough_ops_all_covered(): + """Every op in _PASSTHROUGH_OP_OVERLOADS must have a test case in _PASSTHROUGH_OP_TEST_CASES.""" + assert set(_PASSTHROUGH_OP_TEST_CASES.keys()) == _PASSTHROUGH_OP_OVERLOADS + + +@pytest.mark.parametrize( + "overloadpacket, model_inputs", + _PASSTHROUGH_OP_TEST_CASES.items(), + ids=list(_PASSTHROUGH_OP_TEST_CASES.keys()), +) +def test_passthrough_op_produces_expected_overloadpacket(overloadpacket, model_inputs): + """Each passthrough op must lower to its expected ATen OpOverloadPacket after export.""" + model, inputs = model_inputs + assert overloadpacket in _overloadpackets_in_graph(model, inputs) diff --git a/tests/quantization/test_factory.py b/tests/quantization/test_factory.py index 7e8bf8b..5ff7f46 100644 --- a/tests/quantization/test_factory.py +++ b/tests/quantization/test_factory.py @@ -24,6 +24,7 @@ StaticQParamsCalculator, _DefaultQParamsCalculator, ) +from coreai_opt.quantization.spec.qscheme import QuantizationScheme from coreai_opt.quantization.spec.range_calculator import ( MinMaxRangeCalculator, RangeCalculatorBase, @@ -286,7 +287,6 @@ class ExtraArgFakeQuantizeImpl(_DefaultFakeQuantizeImpl): def __init__( self, dtype, - qscheme, qformulation, granularity, target_dtype, @@ -299,7 +299,6 @@ def __init__( ): super().__init__( dtype, - qscheme, qformulation, granularity, target_dtype, @@ -605,6 +604,97 @@ def test_partial_vs_direct_instantiation(self): assert fq_partial_out.dtype == x.dtype assert fq_direct_out.dtype == x.dtype + def test_update_partial_qparams_calculator_single_attr(self): + """ + update_partial_qparams_calculator overrides one attribute on each constructed calculator. + """ + spec = QuantizationSpec( + dtype=torch.int8, + qscheme="symmetric", + granularity=PerTensorGranularity(), + qparam_calculator_cls=StaticQParamsCalculator, + range_calculator_cls=MinMaxRangeCalculator, + ) + partial = QuantizationComponentFactory.create_fake_quantizer_partial( + spec, quantization_target=CompressionTargetTensor.WEIGHT + ) + + updated = QuantizationComponentFactory.update_partial_qparams_calculator( + partial, qscheme=QuantizationScheme.ASYMMETRIC + ) + + fq = updated() + assert fq.qparams_calculator.qscheme == QuantizationScheme.ASYMMETRIC + # qscheme property on the fake quantizer delegates to the calculator + assert fq.qscheme == QuantizationScheme.ASYMMETRIC + + def test_update_partial_qparams_calculator_multiple_attrs(self): + """update_partial_qparams_calculator overrides multiple attributes in one call.""" + spec = QuantizationSpec( + dtype=torch.int8, + qscheme="symmetric", + granularity=PerTensorGranularity(), + qparam_calculator_cls=StaticQParamsCalculator, + range_calculator_cls=MinMaxRangeCalculator, + ) + partial = QuantizationComponentFactory.create_fake_quantizer_partial( + spec, quantization_target=CompressionTargetTensor.WEIGHT + ) + + updated = QuantizationComponentFactory.update_partial_qparams_calculator( + partial, + qscheme=QuantizationScheme.ASYMMETRIC, + float_range=(0.0, 1.0), + ) + + fq = updated() + assert fq.qparams_calculator.qscheme == QuantizationScheme.ASYMMETRIC + assert fq.qparams_calculator.float_range == (0.0, 1.0) + + def test_update_partial_qparams_calculator_does_not_mutate_original(self): + """update_partial_qparams_calculator returns a new partial; the original is unchanged.""" + spec = QuantizationSpec( + dtype=torch.int8, + qscheme="symmetric", + granularity=PerTensorGranularity(), + qparam_calculator_cls=StaticQParamsCalculator, + range_calculator_cls=MinMaxRangeCalculator, + ) + partial = QuantizationComponentFactory.create_fake_quantizer_partial( + spec, quantization_target=CompressionTargetTensor.WEIGHT + ) + + _ = QuantizationComponentFactory.update_partial_qparams_calculator( + partial, qscheme=QuantizationScheme.ASYMMETRIC + ) + + # Original partial still produces calculators with the original qscheme. + fq_original = partial() + assert fq_original.qparams_calculator.qscheme == QuantizationScheme.SYMMETRIC + + def test_update_partial_qparams_calculator_independent_instances(self): + """Each call to an updated partial creates an independent calculator with the override.""" + spec = QuantizationSpec( + dtype=torch.int8, + qscheme="symmetric", + granularity=PerTensorGranularity(), + qparam_calculator_cls=MovingAverageQParamsCalculator, + range_calculator_cls=MinMaxRangeCalculator, + ) + partial = QuantizationComponentFactory.create_fake_quantizer_partial( + spec, quantization_target=CompressionTargetTensor.ACTIVATION + ) + updated = QuantizationComponentFactory.update_partial_qparams_calculator( + partial, float_range=(0.0, 1.0) + ) + + fq1 = updated() + fq2 = updated() + + assert id(fq1.qparams_calculator) != id(fq2.qparams_calculator) + assert fq1.qparams_calculator.float_range == (0.0, 1.0) + assert fq2.qparams_calculator.float_range == (0.0, 1.0) + class TestCompressionTargetTensorAttribute: """Test the quantization_target attribute functionality""" diff --git a/tests/quantization/test_graph_mode_quantizer.py b/tests/quantization/test_graph_mode_quantizer.py index f70e24d..149b7a2 100644 --- a/tests/quantization/test_graph_mode_quantizer.py +++ b/tests/quantization/test_graph_mode_quantizer.py @@ -1495,3 +1495,218 @@ def test_all_nodes_disabled_finalize_raises_for_no_fq_nodes(self, caplog, backen # Finalize should raise since there are no FQ nodes to export with pytest.raises(ValueError, match="no fake quantization nodes"): quantizer.finalize(backend=backend) + + +class TestFixedQParamsActivations: + """ + Tests that activations with analytically known output ranges receive correct fixed qparams. + """ + + @staticmethod + def _assert_fq_properties( + fq: FakeQuantizeImplBase, + expected_qscheme: QuantizationScheme, + expected_float_range: tuple, + expected_scale: float, + ) -> None: + assert fq.qscheme == expected_qscheme + assert fq.qparams_calculator.qscheme == expected_qscheme + assert fq.qscheme == fq.qparams_calculator.qscheme + assert fq.qparams_calculator.float_range == expected_float_range + scale, _, _ = fq.calculate_qparams() + torch.testing.assert_close(scale, torch.full_like(scale, expected_scale)) + + @pytest.mark.parametrize( + "activation_fn, expected_qscheme, expected_float_range, expected_scale", + [ + pytest.param( + torch.sigmoid, + QuantizationScheme.ASYMMETRIC, + (0.0, 1.0), + # Asymmetric [0, 1] over 255 int8 steps: scale = 1/255. + 1.0 / 255.0, + id="sigmoid", + ), + pytest.param( + torch.tanh, + QuantizationScheme.SYMMETRIC, + (-1.0, 1.0), + # Symmetric [-1, 1] over 255 int8 steps: scale = 2/255. + 2.0 / 255.0, + id="tanh", + ), + pytest.param( + nn.Hardtanh(-3.0, 3.0), + QuantizationScheme.SYMMETRIC, + (-3.0, 3.0), + # Symmetric [-3, 3] over 255 int8 steps: scale = 6/255. + 6.0 / 255.0, + id="hardtanh_symmetric", + ), + pytest.param( + nn.Hardtanh(0.0, 6.0), + QuantizationScheme.ASYMMETRIC, + (0.0, 6.0), + # Asymmetric [0, 6] over 255 int8 steps: scale = 6/255. + 6.0 / 255.0, + id="hardtanh_asymmetric", + ), + ], + ) + def test_fixed_activation_qparams_stable_through_prepare_calibration_and_qat( + self, + activation_fn, + expected_qscheme, + expected_float_range, + expected_scale, + ): + """Fixed-range activation ops should get the correct qscheme and float_range, + with a scale that remains unchanged through calibration and QAT.""" + + class ActivationLinearActivation(nn.Module): + def __init__(self, fn): + super().__init__() + self.linear = nn.Linear(2, 2, bias=False) + self._fn = fn + + def forward(self, x): + x = self.linear(x) + x = self._fn(x) + return x + + torch.manual_seed(42) + model = ActivationLinearActivation(activation_fn).eval() + example_input = torch.randn(1, 2) + + quantizer = Quantizer(model, QuantizerConfig()) + prepared_model = quantizer.prepare((example_input,)) + + # Run one forward to initialize the qparams calculators. + prepared_model(example_input) + + fq = prepared_model.activation_post_process_2 + + # After preparation. + self._assert_fq_properties(fq, expected_qscheme, expected_float_range, expected_scale) + + # After calibration — qparams should be unchanged. + with quantizer.calibration_mode(): + for _ in range(5): + prepared_model(torch.randn(1, 2)) + + self._assert_fq_properties(fq, expected_qscheme, expected_float_range, expected_scale) + + # After QAT — qparams should still be unchanged. + optimizer = torch.optim.SGD(prepared_model.parameters(), lr=1e-3) + criterion = nn.MSELoss() + with quantizer.training_mode(): + for _ in range(5): + x = torch.randn(1, 2) + optimizer.zero_grad() + loss = criterion(prepared_model(x), torch.zeros(1, 2)) + loss.backward() + optimizer.step() + + self._assert_fq_properties(fq, expected_qscheme, expected_float_range, expected_scale) + + def test_relu_output_qparams(self): + """ReLU output fake quant should have ASYMMETRIC qscheme with float_range=(0, None). + + The min is analytically pinned at 0 (zero_point stays at -128 across all phases), + while the max remains data-driven (float_range[1] is None). + """ + + class LinearRelu(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(2, 2, bias=False) + + def forward(self, x): + return torch.relu(self.linear(x)) + + torch.manual_seed(42) + model = LinearRelu().eval() + example_input = torch.randn(1, 2) + + quantizer = Quantizer(model, QuantizerConfig()) + prepared_model = quantizer.prepare((example_input,)) + prepared_model(example_input) + + fq = prepared_model.activation_post_process_2 + + assert fq.qscheme == QuantizationScheme.ASYMMETRIC + assert fq.qparams_calculator.qscheme == QuantizationScheme.ASYMMETRIC + assert fq.qscheme == fq.qparams_calculator.qscheme + assert fq.qparams_calculator.float_range == (0.0, None) + + # Min is pinned to 0 — zero_point should stay at -128 across all phases. + _, zp, _ = fq.calculate_qparams() + assert torch.all(zp == -128) + + with quantizer.calibration_mode(): + for _ in range(5): + prepared_model(torch.randn(1, 2)) + + _, zp, _ = fq.calculate_qparams() + assert torch.all(zp == -128) + + optimizer = torch.optim.SGD(prepared_model.parameters(), lr=1e-3) + criterion = nn.MSELoss() + with quantizer.training_mode(): + for _ in range(5): + x = torch.randn(1, 2) + optimizer.zero_grad() + loss = criterion(prepared_model(x), torch.zeros(1, 2)) + loss.backward() + optimizer.step() + + _, zp, _ = fq.calculate_qparams() + assert torch.all(zp == -128) + + def test_relu_qparams_propagated_through_passthrough_ops(self): + """Relu's ASYMMETRIC qscheme and pinned-min float_range propagate through + passthrough ops (view, squeeze, unsqueeze) to the following quantizable op's input FQ.""" + + class LinearReluPassthroughLinear(nn.Module): + def __init__(self): + super().__init__() + self.linear1 = nn.Linear(2, 2, bias=False) + self.linear2 = nn.Linear(2, 2, bias=False) + + def forward(self, x): + x = torch.relu(self.linear1(x)) + x = x.view(1, 2) + x = x.squeeze(0) + x = x.unsqueeze(0) + return self.linear2(x) + + torch.manual_seed(42) + model = LinearReluPassthroughLinear().eval() + example_input = torch.randn(1, 2) + + quantizer = Quantizer(model, QuantizerConfig()) + prepared_model = quantizer.prepare((example_input,)) + prepared_model(example_input) + + # Both relu's output FQ and linear2's input FQ (reached through the passthrough chain) + # should carry the relu-adjusted spec. + fqs_with_relu_range = [ + (name, m) + for name, m in prepared_model.named_modules() + if isinstance(m, FakeQuantizeImplBase) + and getattr(m.qparams_calculator, "float_range", None) == (0.0, None) + ] + assert len(fqs_with_relu_range) == 2, ( + f"Expected 2 FQ modules with float_range=(0.0, None) " + f"(relu output + linear2 input via passthrough propagation), " + f"got {len(fqs_with_relu_range)}: {[n for n, _ in fqs_with_relu_range]}" + ) + for name, fq in fqs_with_relu_range: + assert fq.qscheme == QuantizationScheme.ASYMMETRIC, ( + f"{name}: expected ASYMMETRIC qscheme" + ) + assert fq.qparams_calculator.qscheme == QuantizationScheme.ASYMMETRIC, ( + f"{name}: qparams_calculator qscheme mismatch" + ) + _, zp, _ = fq.calculate_qparams() + assert torch.all(zp == -128), f"{name}: expected zero_point=-128" diff --git a/tests/quantization/test_graph_mode_quantizer_mnist.py b/tests/quantization/test_graph_mode_quantizer_mnist.py index a9c1cfd..adcc3ce 100644 --- a/tests/quantization/test_graph_mode_quantizer_mnist.py +++ b/tests/quantization/test_graph_mode_quantizer_mnist.py @@ -296,8 +296,8 @@ def test_weight_and_activation_qat_mnist(mnist_pretrained_model, mnist_dataset, example_inputs=(torch.ones(1, 1, 28, 28, dtype=torch.float),) ) post_prepare_accuracy = utils.eval_model(prepared_model, test_loader) - assert post_prepare_accuracy < 88, ( - "Expect accuracy to drop below 88% after preparation with an all ones data sample" + assert post_prepare_accuracy < 94, ( + "Expect accuracy to drop below 94% after preparation with an all ones data sample" ) # Fine tune the model From 291217d0c2a20d8c2cf0f13ca741b1441872e495 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:07:43 -0700 Subject: [PATCH 2/3] Add additional passthrough ops --- .../quantization/_graph/_annotation_utils.py | 5 ++++ tests/quantization/test_annotation_utils.py | 30 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index d5c66f3..55a1608 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -64,11 +64,16 @@ # adjusted qspecs to child nodes. _PASSTHROUGH_OP_OVERLOADS: frozenset = frozenset( { + torch.ops.aten.clone, torch.ops.aten.dropout, + torch.ops.aten.expand, torch.ops.aten.feature_dropout, torch.ops.aten.permute, torch.ops.aten.reshape, + torch.ops.aten.select, + torch.ops.aten.slice, torch.ops.aten.squeeze, + torch.ops.aten.t, torch.ops.aten.transpose, torch.ops.aten.unsqueeze, torch.ops.aten.view, diff --git a/tests/quantization/test_annotation_utils.py b/tests/quantization/test_annotation_utils.py index 6f1f05a..05b77ef 100644 --- a/tests/quantization/test_annotation_utils.py +++ b/tests/quantization/test_annotation_utils.py @@ -66,10 +66,36 @@ def forward(self, x): return F.dropout2d(x, p=0.5, training=True) +class _Clone(nn.Module): + def forward(self, x): + return x.clone() + + +class _Expand(nn.Module): + def forward(self, x): + return x.expand(3, -1) + + +class _Select(nn.Module): + def forward(self, x): + return x.select(1, 0) + + +class _Slice(nn.Module): + def forward(self, x): + return x[:, :2] + + +class _T(nn.Module): + def forward(self, x): + return x.t() + + # Maps each member of _PASSTHROUGH_OP_OVERLOADS to a (model, example_inputs) pair # that exercises it. If an op is added to _PASSTHROUGH_OP_OVERLOADS without a # corresponding entry here, test_passthrough_ops_all_covered will fail. _PASSTHROUGH_OP_TEST_CASES: dict[OpOverloadPacket, tuple[nn.Module, tuple]] = { + torch.ops.aten.clone: (_Clone(), (torch.randn(1, 4),)), torch.ops.aten.dropout: ( _Dropout(), ( @@ -78,10 +104,14 @@ def forward(self, x): ), ), ), + torch.ops.aten.expand: (_Expand(), (torch.randn(1, 4),)), torch.ops.aten.feature_dropout: (_FeatureDropout(), (torch.randn(1, 4, 4, 4),)), torch.ops.aten.permute: (_Permute(), (torch.randn(1, 4, 8),)), torch.ops.aten.reshape: (_Reshape(), (torch.randn(1, 4),)), + torch.ops.aten.select: (_Select(), (torch.randn(1, 4, 4),)), + torch.ops.aten.slice: (_Slice(), (torch.randn(1, 4),)), torch.ops.aten.squeeze: (_Squeeze(), (torch.randn(1, 1, 4),)), + torch.ops.aten.t: (_T(), (torch.randn(2, 4),)), torch.ops.aten.transpose: (_Transpose(), (torch.randn(2, 4),)), torch.ops.aten.unsqueeze: ( _Unsqueeze(), From e9e8e7bee05b65fadfb9862b1cb44fa03e18fed2 Mon Sep 17 00:00:00 2001 From: Kevin Hsieh <2467001+crowbat@users.noreply.github.com> Date: Mon, 13 Jul 2026 15:07:31 -0700 Subject: [PATCH 3/3] Address review comments --- .../quantization/_graph/_annotation_utils.py | 9 ++- src/coreai_opt/quantization/spec/factory.py | 2 +- tests/quantization/test_factory.py | 73 ++++--------------- 3 files changed, 22 insertions(+), 62 deletions(-) diff --git a/src/coreai_opt/quantization/_graph/_annotation_utils.py b/src/coreai_opt/quantization/_graph/_annotation_utils.py index 55a1608..ce41371 100644 --- a/src/coreai_opt/quantization/_graph/_annotation_utils.py +++ b/src/coreai_opt/quantization/_graph/_annotation_utils.py @@ -158,10 +158,13 @@ class OpsListPattern: # Dictionary mapping ops with known output bounds to (qscheme, float_range). # float_range elements may be None to leave that side data-driven. _fixed_q_params_ops = { + # tanh: bounded to [-1, 1] torch.ops.aten.tanh.default: (QuantizationScheme.SYMMETRIC, (-1.0, 1.0)), torch.ops.aten.tanh_.default: (QuantizationScheme.SYMMETRIC, (-1.0, 1.0)), + # sigmoid: bounded to [0, 1] torch.ops.aten.sigmoid.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), torch.ops.aten.sigmoid_.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), + # hardsigmoid: bounded to [0, 1] torch.ops.aten.hardsigmoid.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), torch.ops.aten.hardsigmoid_.default: (QuantizationScheme.ASYMMETRIC, (0.0, 1.0)), # relu: always >= 0, upper bound is data-driven @@ -271,7 +274,9 @@ def _propagate_adjusted_spec_to_child_nodes( if float_range is not None: kwargs["float_range"] = float_range if kwargs: - ctr = QuantizationComponentFactory.update_partial_qparams_calculator(ctr, **kwargs) + ctr = QuantizationComponentFactory.reconstruct_partial_qparams_calculator( + ctr, **kwargs + ) # qscheme in TorchAOQuantizationSpec is not read by coreai-opt later on so we omit it. # Only the qscheme contained within observer_or_fake_quant_ctr matters. @@ -319,7 +324,7 @@ def adjust_output_qspec_for_qscheme_and_propagate( else: return - ctr = QuantizationComponentFactory.update_partial_qparams_calculator( + ctr = QuantizationComponentFactory.reconstruct_partial_qparams_calculator( qspec.observer_or_fake_quant_ctr, qscheme=qscheme, float_range=float_range ) diff --git a/src/coreai_opt/quantization/spec/factory.py b/src/coreai_opt/quantization/spec/factory.py index ef023d2..fbd1ac2 100644 --- a/src/coreai_opt/quantization/spec/factory.py +++ b/src/coreai_opt/quantization/spec/factory.py @@ -213,7 +213,7 @@ def create_fake_quantizer( return spec.fake_quantize_cls(**common_args, **extra_args) @classmethod - def update_partial_qparams_calculator( + def reconstruct_partial_qparams_calculator( cls, partial_ctr: _PartialConstructor, **kwargs: Any, diff --git a/tests/quantization/test_factory.py b/tests/quantization/test_factory.py index 5ff7f46..1afb627 100644 --- a/tests/quantization/test_factory.py +++ b/tests/quantization/test_factory.py @@ -604,9 +604,9 @@ def test_partial_vs_direct_instantiation(self): assert fq_partial_out.dtype == x.dtype assert fq_direct_out.dtype == x.dtype - def test_update_partial_qparams_calculator_single_attr(self): - """ - update_partial_qparams_calculator overrides one attribute on each constructed calculator. + def test_reconstruct_partial_qparams_calculator(self): + """reconstruct_partial_qparams_calculator overrides attributes without mutating the original + partial. """ spec = QuantizationSpec( dtype=torch.int8, @@ -619,78 +619,33 @@ def test_update_partial_qparams_calculator_single_attr(self): spec, quantization_target=CompressionTargetTensor.WEIGHT ) - updated = QuantizationComponentFactory.update_partial_qparams_calculator( + # Overriding a single attribute updates the constructed calculator. + updated_single = QuantizationComponentFactory.reconstruct_partial_qparams_calculator( partial, qscheme=QuantizationScheme.ASYMMETRIC ) - - fq = updated() + fq = updated_single() assert fq.qparams_calculator.qscheme == QuantizationScheme.ASYMMETRIC # qscheme property on the fake quantizer delegates to the calculator assert fq.qscheme == QuantizationScheme.ASYMMETRIC - def test_update_partial_qparams_calculator_multiple_attrs(self): - """update_partial_qparams_calculator overrides multiple attributes in one call.""" - spec = QuantizationSpec( - dtype=torch.int8, - qscheme="symmetric", - granularity=PerTensorGranularity(), - qparam_calculator_cls=StaticQParamsCalculator, - range_calculator_cls=MinMaxRangeCalculator, - ) - partial = QuantizationComponentFactory.create_fake_quantizer_partial( - spec, quantization_target=CompressionTargetTensor.WEIGHT - ) - - updated = QuantizationComponentFactory.update_partial_qparams_calculator( + # Overriding multiple attributes in one call updates all of them. + updated_multiple = QuantizationComponentFactory.reconstruct_partial_qparams_calculator( partial, qscheme=QuantizationScheme.ASYMMETRIC, float_range=(0.0, 1.0), ) - - fq = updated() + fq = updated_multiple() assert fq.qparams_calculator.qscheme == QuantizationScheme.ASYMMETRIC assert fq.qparams_calculator.float_range == (0.0, 1.0) - def test_update_partial_qparams_calculator_does_not_mutate_original(self): - """update_partial_qparams_calculator returns a new partial; the original is unchanged.""" - spec = QuantizationSpec( - dtype=torch.int8, - qscheme="symmetric", - granularity=PerTensorGranularity(), - qparam_calculator_cls=StaticQParamsCalculator, - range_calculator_cls=MinMaxRangeCalculator, - ) - partial = QuantizationComponentFactory.create_fake_quantizer_partial( - spec, quantization_target=CompressionTargetTensor.WEIGHT - ) - - _ = QuantizationComponentFactory.update_partial_qparams_calculator( - partial, qscheme=QuantizationScheme.ASYMMETRIC - ) - - # Original partial still produces calculators with the original qscheme. + # The original partial is unaffected by either override. fq_original = partial() assert fq_original.qparams_calculator.qscheme == QuantizationScheme.SYMMETRIC - def test_update_partial_qparams_calculator_independent_instances(self): - """Each call to an updated partial creates an independent calculator with the override.""" - spec = QuantizationSpec( - dtype=torch.int8, - qscheme="symmetric", - granularity=PerTensorGranularity(), - qparam_calculator_cls=MovingAverageQParamsCalculator, - range_calculator_cls=MinMaxRangeCalculator, - ) - partial = QuantizationComponentFactory.create_fake_quantizer_partial( - spec, quantization_target=CompressionTargetTensor.ACTIVATION - ) - updated = QuantizationComponentFactory.update_partial_qparams_calculator( - partial, float_range=(0.0, 1.0) - ) - - fq1 = updated() - fq2 = updated() - + # Each call to an updated partial creates an independent calculator with the override + # applied. + fq1 = updated_multiple() + fq2 = updated_multiple() assert id(fq1.qparams_calculator) != id(fq2.qparams_calculator) assert fq1.qparams_calculator.float_range == (0.0, 1.0) assert fq2.qparams_calculator.float_range == (0.0, 1.0)