Enable weight fake quantization in calibration mode (#25)

* feat(quantization): keep weight fake quantization enabled during calibration mode

Quantizer.calibration_mode() previously disabled fake quantization on both
weights and activations. It now only disables activation fake quantization,
so activation observers see the effect of quantized weights when computing
activation ranges.

Adds enable_weight_fake_quant/disable_activation_fake_quant helpers used by
both the eager and graph quantizers, and updates docs/tests accordingly.

* fix(quantization): address review feedback on weight FQ calibration change

- Move enable_weight_fake_quant/disable_activation_fake_quant into a new
  _fake_quant_utils module so they can import FakeQuantizeImplBase and
  CompressionTargetTensor at module scope instead of via local imports,
  keeping _utils.py free of dependencies on other quantization modules.
- Use basic_config (weight + activation quantization) instead of
  input_activation_only_config in test_calibration_mode so the weight-FQ
  branch is actually exercised; only assert scale drift for activation
  modules, since weight ranges are fixed at prepare time.
- Drop the changelog.d entry.

---------

Co-authored-by: Utkarsh Simha <u_simha@apple.com>

Author: u-simha

docs/src/examples/resnet50.md

@@ -53,9 +53,9 @@ prepared_model = quantizer.prepare(example_inputs)
 
 Calibration is necessary when activation quantization is enabled. In order to determine proper quantization parameters for activation quantizers, representative data must be passed through the prepared model in the `calibration_mode()` context.
 
-Inside `calibration_mode()`, fake quantization is disabled and observers track tensor ranges seen at each activation quantizer.
-Each forward pass updates the activation scales without introducing quantization noise into the output.
-On exit, observers are disabled and fake quantization is re-enabled.
+Inside `calibration_mode()`, activation fake quantization is disabled while weight fake quantization stays on, and observers track tensor ranges seen at each activation quantizer.
+Each forward pass updates the activation scales using activations produced with quantized weights upstream, without injecting activation quantization noise into the observed values.
+On exit, observers are disabled and activation fake quantization is re-enabled.
 
 ```python
 with quantizer.calibration_mode():

docs/src/quantization/overview.md

@@ -111,9 +111,10 @@ Right after `prepare()`, the activation scales come only from `example_inputs`,
 This is what the context manager handles:
 
 - Inside the context:
-  - fake-quantization is turned **off**: forward pass gives the same output as the unquantized model. Hence without distorting the outputs, the quantization params can be computed.
-  - range observers are turned **on**: this means that each forward pass updates the observed activation ranges, and hence the activation quantization scales.
-- After exiting the context manager, observers are turned back off and fake-quantization back on, leaving the model ready for evaluation.
+  - activation fake-quantization is turned **off**: activation observers see undistorted activation values, so the observed ranges (and resulting scales) reflect the true distribution rather than already-quantized values.
+  - weight fake-quantization stays **on**: activations flowing into each observer are produced with quantized weights upstream, matching what the deployed model will actually see.
+  - range observers are turned **on**: each forward pass updates the observed activation ranges, and hence the activation quantization scales.
+- After exiting the context manager, observers are turned back off and activation fake-quantization back on, leaving the model ready for evaluation.
 
 A small amount of representative data is typically enough.
 

docs/src/tutorials/mnist_quantization.ipynb

@@ -748,7 +748,7 @@
    "source": [
     "We now call `calibration_mode()` and feed it representative inputs to populate scale and zero-point value.\n",
     "\n",
-    "The `calibration_mode()` context manager enables range observers (to collect activation statistics) while disabling fake quantization (so the forward pass is numerically identical to the unquantized model). After exiting the context, observers are frozen and fake quantization is re-enabled."
+    "The `calibration_mode()` context manager enables range observers (to collect activation statistics) and disables activation fake quantization (so the observed ranges reflect undistorted activation values), while leaving weight fake quantization enabled (so activations flowing into each observer are produced with quantized weights upstream). After exiting the context, observers are frozen and activation fake quantization is re-enabled."
    ]
   },
   {

src/coreai_opt/quantization/_eager/quantizer.py

@@ -39,6 +39,10 @@
     apply_weight_axis_defaults_eager as _apply_weight_axis_defaults,
     validate_activation_axes as _validate_activation_axes,
 )
+from coreai_opt.quantization._fake_quant_utils import (
+    disable_activation_fake_quant as _disable_activation_fake_quant,
+    enable_weight_fake_quant as _enable_weight_fake_quant,
+)
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
     ModuleQuantizerConfig,
@@ -291,8 +295,10 @@ def finalize(
     def calibration_mode(self, model: nn.Module | None = None) -> Generator:
         """Context manager for calibration phase.
 
-        Enables observers and disables fake quantization for calibration
-        data collection.
+        Enables observers and disables activation fake quantization for
+        calibration data collection. Weight fake quantization stays enabled so
+        that activation observers see the effect of quantized weights when
+        computing activation ranges.
 
         Args:
             model: Model to calibrate (uses internal model if None)
@@ -309,7 +315,8 @@ def calibration_mode(self, model: nn.Module | None = None) -> Generator:
             )
         with move_model_to_eval(self._model):
             self._model.apply(enable_observer)
-            self._model.apply(disable_fake_quant)
+            self._model.apply(_enable_weight_fake_quant)
+            self._model.apply(_disable_activation_fake_quant)
             try:
                 yield
             finally:

src/coreai_opt/quantization/_fake_quant_utils.py

@@ -0,0 +1,49 @@
+# 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
+
+"""Helpers for toggling fake quantization by quantization target."""
+
+import torch
+
+from coreai_opt.config.spec import CompressionTargetTensor
+from coreai_opt.quantization.spec.fake_quantize import FakeQuantizeImplBase
+
+
+def disable_activation_fake_quant(module: torch.nn.Module) -> None:
+    """Disable fake quantization on activation FakeQuantize modules only.
+
+    Mirrors ``torchao.quantization.pt2e.fake_quantize.disable_fake_quant`` but
+    skips weight FQ modules. Used by ``calibration_mode`` so activation
+    observers see the effect of quantized weights when collecting statistics.
+
+    Args:
+        module (torch.nn.Module): Module to (possibly) toggle. No-op for any
+            module that is not a ``FakeQuantizeImplBase`` whose
+            ``quantization_target`` is ``ACTIVATION``.
+    """
+    if (
+        isinstance(module, FakeQuantizeImplBase)
+        and module.quantization_target == CompressionTargetTensor.ACTIVATION
+    ):
+        module.disable_fake_quant()
+
+
+def enable_weight_fake_quant(module: torch.nn.Module) -> None:
+    """Enable fake quantization on weight FakeQuantize modules only.
+
+    Mirrors ``torchao.quantization.pt2e.fake_quantize.enable_fake_quant`` but
+    skips activation FQ modules. Companion to
+    :func:`disable_activation_fake_quant` used by ``calibration_mode``.
+
+    Args:
+        module (torch.nn.Module): Module to (possibly) toggle. No-op for any
+            module that is not a ``FakeQuantizeImplBase`` whose
+            ``quantization_target`` is ``WEIGHT``.
+    """
+    if (
+        isinstance(module, FakeQuantizeImplBase)
+        and module.quantization_target == CompressionTargetTensor.WEIGHT
+    ):
+        module.enable_fake_quant()

src/coreai_opt/quantization/_graph/quantizer.py

@@ -60,6 +60,10 @@
     apply_weight_axis_defaults_graph as _apply_weight_axis_defaults,
     validate_activation_axes as _validate_activation_axes,
 )
+from coreai_opt.quantization._fake_quant_utils import (
+    disable_activation_fake_quant as _disable_activation_fake_quant,
+    enable_weight_fake_quant as _enable_weight_fake_quant,
+)
 from coreai_opt.quantization.base_quantizer import _BaseQuantizer
 from coreai_opt.quantization.config import (
     KVCacheQuantConfig,
@@ -1128,9 +1132,11 @@ def calibration_mode(self, model: torch.fx.GraphModule | None = None):
         Context manager for calibration-based post-training quantization.
 
         When entering this context, observers are enabled to collect statistics
-        from calibration data, and fake quantization is disabled to get accurate
-        statistics. When exiting, observers are disabled and fake quantization
-        is re-enabled for evaluation.
+        from calibration data. Weight fake quantization stays enabled, while
+        activation fake quantization is disabled so that activation observers
+        see the effect of quantized weights when computing activation ranges.
+        When exiting, observers are disabled and fake quantization is
+        re-enabled on both weights and activations for evaluation.
 
         **When to use:**
         - Required for activation quantization to achieve good accuracy
@@ -1164,9 +1170,11 @@ def calibration_mode(self, model: torch.fx.GraphModule | None = None):
                 "Model must be prepared before entering calibration mode. Call prepare() first."
             )
 
-        # Enable observers and disable fake quantization for calibration
+        # Enable observers; keep weight FQ on, disable activation FQ so observers
+        # see the effect of quantized weights on activation ranges.
         self._model.apply(enable_observer)
-        self._model.apply(disable_fake_quant)
+        self._model.apply(_enable_weight_fake_quant)
+        self._model.apply(_disable_activation_fake_quant)
 
         # Move model to eval mode and save original state
         with move_model_to_eval(self._model):

src/coreai_opt/quantization/quantizer.py

@@ -488,9 +488,11 @@ def calibration_mode(self, model: nn.Module | fx.GraphModule | None = None):
         Context manager for calibration-based post-training quantization.
 
         When entering this context, observers are enabled to collect statistics
-        from calibration data, and fake quantization is disabled to get accurate
-        statistics. When exiting, observers are disabled and fake quantization
-        is re-enabled for evaluation.
+        from calibration data. Weight fake quantization stays enabled, while
+        activation fake quantization is disabled so that activation observers
+        see the effect of quantized weights when computing activation ranges.
+        When exiting, observers are disabled and fake quantization is
+        re-enabled on both weights and activations for evaluation.
 
         **When to use:**
 

tests/quantization/test_eager_quant.py

@@ -20,6 +20,7 @@
     FunctionRegisteredOptimizers,
     RegisteredOptimizersTracker,
 )
+from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization import (
     ModuleQuantizerConfig,
     QuantizationSpec,
@@ -529,23 +530,30 @@ def test_finalize_state_dict_safetensors_roundtrip(self, basic_config, tmp_path)
 
         assert torch.equal(out_before_roundtrip, out_after_roundtrip)
 
-    def test_calibration_mode(self, simple_model, input_activation_only_config, example_input):
+    def test_calibration_mode(self, simple_model, basic_config, example_input):
         """
         Test that calibration mode works as expected, and scales are getting updated
         """
-        quantizer = Quantizer(simple_model, input_activation_only_config)
+        quantizer = Quantizer(simple_model, basic_config)
         simple_model.eval()
         prepared_model = quantizer.prepare((example_input,))
 
         fake_quant_modules = [
             m for m in prepared_model.modules() if isinstance(m, FakeQuantizeImplBase)
         ]
+        activation_fake_quant_modules = [
+            m
+            for m in fake_quant_modules
+            if m.quantization_target == CompressionTargetTensor.ACTIVATION
+        ]
 
         for module in fake_quant_modules:
             assert module.observer_enabled.item() == 0
             assert module.fake_quant_enabled.item() == 1
 
-        pre_calibration_scales = [mod.calculate_qparams()[0].clone() for mod in fake_quant_modules]
+        pre_calibration_scales = [
+            mod.calculate_qparams()[0].clone() for mod in activation_fake_quant_modules
+        ]
 
         with quantizer.calibration_mode():
             simple_model.eval()
@@ -554,10 +562,20 @@ def test_calibration_mode(self, simple_model, input_activation_only_config, exam
 
             for module in fake_quant_modules:
                 assert module.observer_enabled.item() == 1
-                assert module.fake_quant_enabled.item() == 0
+                # Weight FQ stays on so activation observers see quantized weights;
+                # activation FQ is off so observers collect statistics on the raw
+                # (post-weight-quant) activations.
+                expected_fq = (
+                    1 if module.quantization_target == CompressionTargetTensor.WEIGHT else 0
+                )
+                assert module.fake_quant_enabled.item() == expected_fq
 
-        post_calibration_scales = [mod.calculate_qparams()[0].clone() for mod in fake_quant_modules]
+        post_calibration_scales = [
+            mod.calculate_qparams()[0].clone() for mod in activation_fake_quant_modules
+        ]
 
+        # Only activation scales are expected to move here: weight ranges are
+        # fixed at prepare time and don't depend on calibration data.
         for pre_scale, post_scale in zip(
             pre_calibration_scales, post_calibration_scales, strict=True
         ):
@@ -1732,8 +1750,10 @@ def test_e2e_workflow(self, basic_config, simple_model, example_input):
         with quantizer.calibration_mode():
             prepared_out = prepared_model(example_input_2)
             original_out = base_model(example_input_2)
-            # prepare model output should match base model, since fake quant is disabled
-            assert torch.equal(prepared_out, original_out)
+            # prepared model output should NOT match base model: weight fake
+            # quant stays on during calibration so activation observers see the
+            # effect of quantized weights. Only activation FQ is disabled.
+            assert not torch.equal(prepared_out, original_out)
 
         pre_finalize_out = prepared_model(example_input_3)
 

tests/quantization/test_graph_mode_quantizer.py

@@ -21,6 +21,7 @@
 from coreai_opt._utils.metadata_utils import (
     STATE_DICT_METADATA_BUFFER_PREFIX as _COREML_BUFFER_PREFIX,
 )
+from coreai_opt.config.spec import CompressionTargetTensor
 from coreai_opt.quantization import (
     ModuleQuantizerConfig,
     QuantizationSpec,
@@ -374,11 +375,16 @@ def test_observer_fake_quant_status(
 
         # Test entering calibration mode
         with quantizer.calibration_mode():
-            # Verify that observers are enabled and fake quant is disabled
+            # Verify that observers are enabled. Weight FQ stays on so
+            # activation observers see the effect of quantized weights;
+            # activation FQ is disabled so observers collect raw stats.
             for _name, module in prepared_model.named_modules():
                 if isinstance(module, FakeQuantizeImplBase):
                     assert module.observer_enabled.item() == 1
-                    assert module.fake_quant_enabled.item() == 0
+                    expected_fq = (
+                        1 if module.quantization_target == CompressionTargetTensor.WEIGHT else 0
+                    )
+                    assert module.fake_quant_enabled.item() == expected_fq
 
         # # Verify that observers are disabled and fake quant is enabled on exit
         for _name, module in prepared_model.named_modules():
@@ -424,12 +430,15 @@ def test_calibration_mode_with_external_model(
 
         # Test calibration_mode with external prepared model
         with quantizer.calibration_mode(other_prepared_model):
-            # Verify that observers are enabled and fake quant is disabled
-            # on the external model
+            # Verify that observers are enabled on the external model. Weight FQ
+            # stays on; activation FQ is off.
             for _name, module in other_prepared_model.named_modules():
                 if isinstance(module, FakeQuantizeImplBase):
                     assert module.observer_enabled.item() == 1
-                    assert module.fake_quant_enabled.item() == 0
+                    expected_fq = (
+                        1 if module.quantization_target == CompressionTargetTensor.WEIGHT else 0
+                    )
+                    assert module.fake_quant_enabled.item() == expected_fq
 
             # The quantizer's internal model should be updated to the provided model
             assert quantizer._model is other_prepared_model
@@ -460,8 +469,10 @@ def test_end_to_end_workflow(self, simple_conv_linear_model, basic_config, simpl
         with quantizer.calibration_mode():
             prepared_out = prepared_model(simple_model_input_2)
             original_out = simple_conv_linear_model(simple_model_input_2)
-            # prepare model output should match base model, since fake quant is disabled
-            assert torch.equal(prepared_out, original_out)
+            # prepared model output should NOT match base model: weight fake
+            # quant stays on during calibration so activation observers see the
+            # effect of quantized weights. Only activation FQ is disabled.
+            assert not torch.equal(prepared_out, original_out)
 
         pre_finalize_out = prepared_model(simple_model_input_3)