From 653b476a43a477fc51aafe1c9a7d613a7f68958e Mon Sep 17 00:00:00 2001 From: M Platypus Date: Tue, 4 Aug 2026 18:30:26 -0400 Subject: [PATCH 1/2] fix: Hessian auto-quant crash and qconfig_dict mutation bugs Two bugs in Hessian-based auto-quantization: 1. run_auto_quantization() crashed with a bare TypeError instead of failing gracefully. optimal_bitwidth stays None whenever no calibration_dataloader is supplied, or compute_hessian_sensitivity can't compute sensitivities (returns ({}, {}) when inputs/targets/criterion aren't available) -- a path callers like audio_classification/train.py already anticipate and try to degrade gracefully from, logging "Could not obtain sample data... Proceeding without it". That graceful-degradation intent was defeated a few lines later by `total_bit_budget = optimal_bitwidth * total_params`, raising TypeError: unsupported operand type(s) for *: 'NoneType' and 'int'. Fixed by returning the caller's already-built qconfig_mapping unchanged when optimal_bitwidth is None -- this is the qconfig_mapping the caller constructed from the manually-specified uniform bitwidth before auto- quantization was attempted, making it the correct, non-crashing fallback (not a forced all-float32 fallback, and not a crash). 2. apply_mixed_precision() wrote qconfig_dict['weight']['bitwidth'] / ['activation']['bitwidth'] directly onto whatever dict it was handed. Two of its four call sites (qconfig_types.py's manual mixed-precision path, and auto_quantization.py's default Hessian auto-quantization path -- i.e. essentially every successful quantized run) passed the original, un-copied qconfig_dict, which is frequently the same dict object retained as self.qconfig_type on the live model wrapper (quant_base.py). Any downstream logging, checkpoint metadata, or reuse of that config silently observed whichever bitwidth tier this function processed last, not the model's actual/intended default. Two other call sites already worked around this by building a throwaway copy before calling in -- confirming the mutation was an oversight, not intended behavior. Fixed in the function itself (not just the two unsafe call sites): builds a local copy per bit-width tier instead of mutating the caller's dict, so every caller is protected regardless of whether it already built its own defensive copy. Verified: 3 of 4 added tests fail on the unmodified code (both crash tests reproduce the exact TypeError; the mutation test observes the caller's dict change from 8 to 4) and pass after the fix. The 4th test confirms the mutation fix doesn't break the actual mixed-precision assignment. Co-Authored-By: Claude Sonnet 5 --- .../tests/test_auto_quantization_bugs.py | 91 +++++++++++++++++++ .../quantization/base/fx/auto_quantization.py | 16 ++++ .../quantization/base/fx/qconfig_types.py | 16 +++- 3 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py diff --git a/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py b/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py new file mode 100644 index 00000000..d2cec5c9 --- /dev/null +++ b/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py @@ -0,0 +1,91 @@ +"""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 -- set_module_name must still be called with a + qconfig built at the requested bitwidth for each layer.""" + qconfig_mapping = QConfigMapping().set_global(get_default_qconfig()) + + result = apply_mixed_precision( + qconfig_mapping, {"weight": {"bitwidth": 8}, "activation": {"bitwidth": 8}}, {32: ["fc"]} + ) + + # bit_width == 32 takes the "disable quantization for this layer" path. + assert result.module_name_qconfigs["fc"] is None diff --git a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/auto_quantization.py b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/auto_quantization.py index c1f9ace7..02085a40 100644 --- a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/auto_quantization.py +++ b/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}") diff --git a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/qconfig_types.py b/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/qconfig_types.py index 7034b52e..91582812 100644 --- a/tinyml-modeloptimization/torchmodelopt/tinyml_torchmodelopt/quantization/base/fx/qconfig_types.py +++ b/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 From 0c62fc803264b5f3acd5c10b0f4791440a5683af Mon Sep 17 00:00:00 2001 From: M Platypus Date: Tue, 4 Aug 2026 21:45:43 -0400 Subject: [PATCH 2/2] test: exercise the sub-32-bit branch apply_mixed_precision actually rewrote Both CodeRabbit and the earlier Opus review flagged that test_apply_mixed_precision_still_applies_the_requested_bitwidth_per_layer used {32: ["fc"]}, which takes the bit_width == 32 "disable quantization" path -- a branch the mutation fix never touched -- so nothing verified the rewritten bw_qconfig_dict construction actually produces a qconfig at the requested bitwidth. Now uses {4: ["fc"]} and asserts the produced qconfig's real quantization ranges (signed 4-bit weights: quant_max 7; unsigned 4-bit activations: quant_max 15, values confirmed empirically against get_default_qconfig before asserting). The 32-bit disable path keeps its own separate test. Co-Authored-By: Claude Sonnet 5 --- .../tests/test_auto_quantization_bugs.py | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py b/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py index d2cec5c9..04b7bc93 100644 --- a/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py +++ b/tinyml-modeloptimization/torchmodelopt/tests/test_auto_quantization_bugs.py @@ -79,13 +79,28 @@ def test_apply_mixed_precision_does_not_mutate_callers_qconfig_dict(): 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 -- set_module_name must still be called with a - qconfig built at the requested bitwidth for each layer.""" + 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"]} ) - # bit_width == 32 takes the "disable quantization for this layer" path. assert result.module_name_qconfigs["fc"] is None