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..04b7bc93 --- /dev/null +++ b/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 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