Fix output spec adjustment for fixed qparams ops (#22)

- Update documentation

Author: crowbat

changelog.d/180525445.fixed

@@ -0,0 +1 @@
+Fix setting of qscheme and float_range for fixed output range ops

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.
 

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,27 @@
 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.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,
+    }
+)
+
 
 def _get_aten_graph_module_for_pattern(
     pattern: Callable,
@@ -131,38 +155,30 @@ 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,
+    # 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
+    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 +223,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 +267,22 @@ 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.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.
             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 +313,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.reconstruct_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(

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 reconstruct_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,

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."""