Validate AutoQuant search checkpoint inputs

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

Author: meenchen

modelopt/torch/quantization/algorithms.py

@@ -508,6 +508,11 @@ def _module_search_space_signature(module_search_spaces) -> tuple:
     )
 
 
+def _quantization_formats_signature(quant_recipes) -> tuple[str, ...]:
+    """Return a checkpoint-stable description of the global candidate formats."""
+    return tuple(sorted(recipe.checkpoint_signature for recipe in quant_recipes))
+
+
 class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
     """Base searcher for AutoQuantize algorithm."""
 
@@ -569,6 +574,7 @@ def default_state_dict(self) -> SearchStateDict:
             "cost": {},
             "active_moe_expert_ratio": None,
             "cost_denominator": None,
+            "quantization_formats_signature": None,
             "module_search_space_signature": None,
             "disabled_layers": None,
             "candidate_stats": defaultdict(dict),
@@ -911,11 +917,29 @@ def before_search(self):
         module_search_spaces = self._normalize_module_search_spaces(
             self.config["module_search_spaces"]
         )
+        default_search_recipes = self._get_search_recipes(self.config["quantization_formats"])
+        quantization_formats_signature = _quantization_formats_signature(default_search_recipes)
         module_search_space_signature = _module_search_space_signature(module_search_spaces)
+        restored_quantization_formats_signature = getattr(
+            self, "quantization_formats_signature", None
+        )
         restored_module_search_space_signature = getattr(
             self, "module_search_space_signature", None
         )
         has_restored_calibration_or_scores = bool(self.quantizer_states or self.candidate_stats)
+        if has_restored_calibration_or_scores and restored_quantization_formats_signature is None:
+            raise ValueError(
+                "Checkpoint does not record its quantization_formats signature and cannot be "
+                "safely reused. Use a different checkpoint path."
+            )
+        if (
+            has_restored_calibration_or_scores
+            and restored_quantization_formats_signature != quantization_formats_signature
+        ):
+            raise ValueError(
+                "Checkpoint quantization_formats do not match the current search config. "
+                "Use a different checkpoint path."
+            )
         if has_restored_calibration_or_scores and (
             (restored_module_search_space_signature is None and module_search_space_signature)
             or (
@@ -927,9 +951,9 @@ def before_search(self):
                 "Checkpoint module_search_spaces do not match the current search config. "
                 "Use a different checkpoint path."
             )
+        self.quantization_formats_signature = quantization_formats_signature
         self.module_search_space_signature = module_search_space_signature
 
-        default_search_recipes = self._get_search_recipes(self.config["quantization_formats"])
         search_recipes = sorted(
             {
                 *default_search_recipes,

modelopt/torch/quantization/model_quant.py

@@ -519,9 +519,13 @@ def _process_quantization_formats(formats, custom_name_prefix):
             processed.append((quant_cfg, name))
         return processed
 
+    if not isinstance(quantization_formats, list):
+        raise TypeError("`quantization_formats` must be a list.")
+    if not quantization_formats:
+        raise ValueError("`quantization_formats` must be a non-empty list.")
     processed_quantization_formats = _process_quantization_formats(quantization_formats, "CUSTOM")
-
-    assert len(processed_quantization_formats) > 0, "`quantization_formats` should not be empty"
+    if not processed_quantization_formats:
+        raise ValueError("`quantization_formats` must contain at least one non-None format.")
 
     processed_module_search_spaces = []
     for idx, search_space in enumerate(module_search_spaces or []):
@@ -545,12 +549,15 @@ def _process_quantization_formats(formats, custom_name_prefix):
             raise ValueError(
                 "module_search_spaces.module_name_patterns must be a non-empty string list."
             )
-        formats = _process_quantization_formats(
-            search_space.get("quantization_formats") or [], f"CUSTOM_MODULE_{idx}"
-        )
+        raw_formats = search_space.get("quantization_formats")
+        if not isinstance(raw_formats, list):
+            raise TypeError("module_search_spaces.quantization_formats must be a list.")
+        if not raw_formats:
+            raise ValueError("module_search_spaces.quantization_formats must be a non-empty list.")
+        formats = _process_quantization_formats(raw_formats, f"CUSTOM_MODULE_{idx}")
         if not formats:
             raise ValueError(
-                "module_search_spaces.quantization_formats must contain at least one format."
+                "module_search_spaces.quantization_formats must contain at least one non-None format."
             )
         allow_no_quant = search_space.get("allow_no_quant", True)
         if not isinstance(allow_no_quant, bool):

tests/unit/torch/quantization/test_autoquant.py

@@ -24,6 +24,7 @@
 
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
+import modelopt.torch.quantization.model_quant as model_quant
 from modelopt.torch.quantization._auto_quantize_cost import (
     EXCLUDED_MODULE_NAME_PATTERNS_KEY,
     _get_module_weight_numel,
@@ -39,7 +40,7 @@
     estimate_quant_compression,
 )
 from modelopt.torch.quantization.config import _base_disable_all, _default_disabled_quantizer_cfg
-from modelopt.torch.utils import safe_load
+from modelopt.torch.utils import safe_load, safe_save
 from modelopt.torch.utils.distributed import DistributedProcessGroup
 
 
@@ -366,9 +367,51 @@ def test_auto_quantize_module_search_spaces_keep_fixed_routed_experts_costed():
     assert routed_hparam.active == int4_recipe
 
 
-def test_auto_quantize_fixed_module_isolated_from_unrelated_calibration(monkeypatch):
-    import modelopt.torch.quantization.model_quant as model_quant
+@pytest.mark.parametrize("formats", ["FP8_DEFAULT_CFG", mtq.FP8_DEFAULT_CFG, ()])
+def test_auto_quantize_rejects_non_list_global_formats(formats):
+    with pytest.raises(TypeError, match="`quantization_formats` must be a list"):
+        mtq.auto_quantize(TransformerBlock(), quantization_formats=formats)
+
+
+@pytest.mark.parametrize("formats", [[], [None]])
+def test_auto_quantize_rejects_empty_global_formats(formats):
+    with pytest.raises(ValueError, match="`quantization_formats` must"):
+        mtq.auto_quantize(TransformerBlock(), quantization_formats=formats)
+
+
+@pytest.mark.parametrize("formats", ["FP8_DEFAULT_CFG", mtq.FP8_DEFAULT_CFG, ()])
+def test_auto_quantize_rejects_non_list_module_formats(formats):
+    with pytest.raises(
+        TypeError, match=r"module_search_spaces\.quantization_formats must be a list"
+    ):
+        mtq.auto_quantize(
+            TransformerBlock(),
+            quantization_formats=[mtq.INT8_DEFAULT_CFG],
+            module_search_spaces=[
+                {
+                    "module_name_patterns": ["*mlp*"],
+                    "quantization_formats": formats,
+                }
+            ],
+        )
+
+
+@pytest.mark.parametrize("formats", [[], [None]])
+def test_auto_quantize_rejects_empty_module_formats(formats):
+    with pytest.raises(ValueError, match=r"module_search_spaces\.quantization_formats must"):
+        mtq.auto_quantize(
+            TransformerBlock(),
+            quantization_formats=[mtq.INT8_DEFAULT_CFG],
+            module_search_spaces=[
+                {
+                    "module_name_patterns": ["*mlp*"],
+                    "quantization_formats": formats,
+                }
+            ],
+        )
 
+
+def test_auto_quantize_fixed_module_isolated_from_unrelated_calibration(monkeypatch):
     model = TransformerBlock()
     calibration_states = []
     original_calibrate = model_quant.calibrate
@@ -1011,6 +1054,82 @@ def interrupt_after_calibration(self):
         )
 
 
+def test_auto_quantize_calibration_only_checkpoint_validates_global_formats_and_legacy(
+    tmp_path, monkeypatch
+):
+    checkpoint_path = str(tmp_path / "autoquant_calibration_only_checkpoint.pth")
+    legacy_checkpoint_path = str(tmp_path / "autoquant_legacy_calibration_only_checkpoint.pth")
+    original_estimate_scores = AutoQuantizeGradientSearcher.estimate_sensitivity_scores
+
+    def interrupt_after_calibration(self):
+        raise RuntimeError("interrupt after calibration")
+
+    monkeypatch.setattr(
+        AutoQuantizeGradientSearcher,
+        "estimate_sensitivity_scores",
+        interrupt_after_calibration,
+    )
+    model = TransformerBlock()
+    with pytest.raises(RuntimeError, match="interrupt after calibration"):
+        mtq.auto_quantize(
+            model,
+            constraints={"effective_bits": 6.0},
+            quantization_formats=[
+                mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
+                mtq.INT8_DEFAULT_CFG,
+            ],
+            data_loader=[model.get_input()],
+            forward_step=lambda model, batch: model(batch),
+            loss_func=lambda output, data: output.sum(),
+            num_calib_steps=1,
+            num_score_steps=1,
+            checkpoint=checkpoint_path,
+        )
+
+    saved = safe_load(checkpoint_path)
+    assert saved["quantizer_states"]
+    assert saved["quantization_formats_signature"]
+    legacy_saved = dict(saved)
+    legacy_saved.pop("quantization_formats_signature")
+    safe_save(legacy_saved, legacy_checkpoint_path)
+
+    monkeypatch.setattr(
+        AutoQuantizeGradientSearcher,
+        "estimate_sensitivity_scores",
+        original_estimate_scores,
+    )
+    mismatched_model = TransformerBlock()
+    with pytest.raises(ValueError, match="quantization_formats do not match"):
+        mtq.auto_quantize(
+            mismatched_model,
+            constraints={"effective_bits": 6.0},
+            quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.FP8_DEFAULT_CFG],
+            data_loader=[mismatched_model.get_input()],
+            forward_step=lambda model, batch: model(batch),
+            loss_func=lambda output, data: output.sum(),
+            num_calib_steps=1,
+            num_score_steps=1,
+            checkpoint=checkpoint_path,
+        )
+
+    legacy_model = TransformerBlock()
+    with pytest.raises(ValueError, match="does not record its quantization_formats signature"):
+        mtq.auto_quantize(
+            legacy_model,
+            constraints={"effective_bits": 6.0},
+            quantization_formats=[
+                mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
+                mtq.INT8_DEFAULT_CFG,
+            ],
+            data_loader=[legacy_model.get_input()],
+            forward_step=lambda model, batch: model(batch),
+            loss_func=lambda output, data: output.sum(),
+            num_calib_steps=1,
+            num_score_steps=1,
+            checkpoint=legacy_checkpoint_path,
+        )
+
+
 @pytest.mark.parametrize("method", ["gradient", "kl_div"])
 def test_get_auto_quantize_config(method):
     model = TransformerBlock()