fix: Hessian auto-quant crash and qconfig_dict mutation bugs

tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py

@@ -0,0 +1,106 @@
+"""Regression tests for two bugs in Hessian-based auto-quantization:
+
+1. run_auto_quantization() crashed with a bare TypeError (None * int) instead
+   of failing gracefully when the Hessian sensitivity search couldn't run (no
+   calibration_dataloader, or sample inputs/targets/criterion unavailable --
+   a path callers like audio_classification/train.py already anticipate and
+   try to degrade gracefully from).
+2. apply_mixed_precision() mutated the caller's qconfig_dict in place, which
+   is frequently the same dict object retained as self.qconfig_type on the
+   live model wrapper.
+"""
+import torch
+import torch.nn as nn
+from torch.ao.quantization import QConfigMapping
+
+from tinyml_torchmodelopt.quantization.base.fx.auto_quantization import run_auto_quantization
+from tinyml_torchmodelopt.quantization.base.fx.qconfig_types import apply_mixed_precision, get_default_qconfig
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = nn.Linear(4, 2)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+def test_run_auto_quantization_falls_back_when_sensitivities_unavailable():
+    """No 'inputs'/'targets'/'criterion' in qconfig_dict -> compute_hessian_sensitivity
+    returns ({}, {}) -> optimal_bitwidth stays None. Must fall back to the
+    caller's already-built qconfig_mapping, not crash on None * total_params."""
+    model = _TinyModel()
+    fallback_mapping = QConfigMapping().set_global(get_default_qconfig())
+    qconfig_dict = {"weight": {"bitwidth": 8}, "activation": {"bitwidth": 8}}
+
+    result = run_auto_quantization(
+        model, qconfig_dict, fallback_mapping,
+        get_default_qconfig_fn=get_default_qconfig,
+        apply_mixed_precision_fn=apply_mixed_precision,
+    )
+
+    assert result is fallback_mapping
+
+
+def test_run_auto_quantization_falls_back_when_no_calibration_dataloader():
+    """inputs/targets/criterion present (so sensitivities *could* compute) but
+    no calibration_dataloader -- the binary search branch is also skipped, so
+    optimal_bitwidth still ends up None via a different path than the test
+    above. Must not crash either."""
+    model = _TinyModel()
+    fallback_mapping = QConfigMapping().set_global(get_default_qconfig())
+    qconfig_dict = {
+        "weight": {"bitwidth": 8}, "activation": {"bitwidth": 8},
+        "inputs": torch.randn(4, 4), "targets": torch.zeros(4, dtype=torch.long),
+        "criterion": nn.CrossEntropyLoss(),
+        "calibration_dataloader": None,
+    }
+
+    result = run_auto_quantization(
+        model, qconfig_dict, fallback_mapping,
+        get_default_qconfig_fn=get_default_qconfig,
+        apply_mixed_precision_fn=apply_mixed_precision,
+    )
+
+    assert result is fallback_mapping
+
+
+def test_apply_mixed_precision_does_not_mutate_callers_qconfig_dict():
+    qconfig_mapping = QConfigMapping().set_global(get_default_qconfig())
+    qconfig_dict = {"weight": {"bitwidth": 8}, "activation": {"bitwidth": 8}}
+    original_weight_bitwidth = qconfig_dict["weight"]["bitwidth"]
+
+    apply_mixed_precision(qconfig_mapping, qconfig_dict, {4: ["fc"]})
+
+    assert qconfig_dict["weight"]["bitwidth"] == original_weight_bitwidth
+    assert qconfig_dict["activation"]["bitwidth"] == original_weight_bitwidth
+
+
+def test_apply_mixed_precision_still_applies_the_requested_bitwidth_per_layer():
+    """Regression safety: the mutation fix must not break the actual mixed-
+    precision assignment. bit_width < 32 takes the branch the fix rewrote
+    (the local bw_qconfig_dict construction), so assert the produced qconfig
+    genuinely reflects the requested 4-bit width -- signed 4-bit weights span
+    [-7, 7] and unsigned 4-bit activations span [0, 15]."""
+    qconfig_mapping = QConfigMapping().set_global(get_default_qconfig())
+
+    result = apply_mixed_precision(
+        qconfig_mapping, {"weight": {"bitwidth": 8}, "activation": {"bitwidth": 8}}, {4: ["fc"]}
+    )
+
+    qconfig = result.module_name_qconfigs["fc"]
+    assert qconfig is not None
+    assert qconfig.weight().quant_max == 7
+    assert qconfig.activation().quant_max == 15
+
+
+def test_apply_mixed_precision_bitwidth_32_disables_quantization_for_the_layer():
+    """bit_width == 32 takes the separate "disable quantization" path."""
+    qconfig_mapping = QConfigMapping().set_global(get_default_qconfig())
+
+    result = apply_mixed_precision(
+        qconfig_mapping, {"weight": {"bitwidth": 8}, "activation": {"bitwidth": 8}}, {32: ["fc"]}
+    )
+
+    assert result.module_name_qconfigs["fc"] is None

tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/auto_quantization.py

@@ -455,6 +455,22 @@ def run_auto_quantization(model, qconfig_dict, qconfig_mapping, get_default_qcon
         bsearch_threshold = None
         higher_is_better = None
 
+    if optimal_bitwidth is None:
+        # Either no calibration_dataloader was supplied, or compute_hessian_sensitivity
+        # couldn't compute sensitivities (e.g. sample inputs/targets/criterion weren't
+        # available -- callers like audio_classification/train.py already anticipate
+        # this and try to degrade gracefully, logging "Proceeding without it" before
+        # calling in here). Without a sensitivity-driven bitwidth search there is
+        # nothing to allocate, so fall back to the qconfig_mapping the caller already
+        # built from the manually-specified uniform bitwidth, unchanged, instead of
+        # crashing on `None * total_params` a few lines below.
+        logger.warning(
+            "Hessian-based auto-quantization could not run (no calibration_dataloader "
+            "or no computable sensitivities) -- falling back to the manually-specified "
+            "uniform bitwidth qconfig instead of a mixed-precision search."
+        )
+        return qconfig_mapping
+
     total_params = sum(module_params.values())
     total_bit_budget = optimal_bitwidth * total_params
     logger.info(f"Total parameters: {total_params}")

tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/qconfig_types.py

@@ -166,9 +166,19 @@ def apply_mixed_precision(qconfig_mapping, qconfig_dict, mixed_precision):
             for layer in layers:
                 qconfig_mapping.set_module_name(layer, None)
         else:
-            qconfig_dict['weight']['bitwidth'] = bit_width
-            qconfig_dict['activation']['bitwidth'] = bit_width
-            qconfig = get_default_qconfig(qconfig_dict=qconfig_dict)
+            # Build a local copy instead of mutating the caller's qconfig_dict in
+            # place. qconfig_dict is frequently the same dict object retained as
+            # self.qconfig_type on the model wrapper (quant_base.py) -- writing
+            # qconfig_dict['weight']['bitwidth'] directly here left that live,
+            # retained object holding whichever bitwidth tier this loop processed
+            # last, rather than the model's actual/intended default. Some callers
+            # already built their own throwaway copy before calling in here as a
+            # workaround; fixing it in this function protects every caller.
+            bw_qconfig_dict = {
+                'weight': {**qconfig_dict.get('weight', {}), 'bitwidth': bit_width},
+                'activation': {**qconfig_dict.get('activation', {}), 'bitwidth': bit_width},
+            }
+            qconfig = get_default_qconfig(qconfig_dict=bw_qconfig_dict)
             for layer in layers:
                 qconfig_mapping.set_module_name(layer, qconfig)
     return qconfig_mapping