From dc051ab5e8bf38588e71dc3a0cf6e513f7828212 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Wed, 8 Jul 2026 17:49:31 -0700
Subject: [PATCH 1/5] Add module-specific AutoQuant search spaces

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/hf_ptq.py                     |  44 ++++--
 modelopt/recipe/config.py                     |  47 ++++++
 modelopt/torch/quantization/algorithms.py     | 143 ++++++++++++++++--
 modelopt/torch/quantization/model_quant.py    |  89 +++++++++--
 ...8_module_spaces_at_6p0bits-active_moe.yaml |  66 ++++++++
 tests/examples/hf_ptq/test_hf_ptq_args.py     |  27 ++++
 tests/unit/recipe/test_loader.py              |  21 +++
 .../unit/torch/quantization/test_autoquant.py |  93 ++++++++++++
 8 files changed, 494 insertions(+), 36 deletions(-)
 create mode 100644 modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe.yaml

diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py
index 8dcc78afa27..b37a00e939a 100755
--- a/examples/hf_ptq/hf_ptq.py
+++ b/examples/hf_ptq/hf_ptq.py
@@ -303,6 +303,25 @@ def _match_candidate_to_preset(fmt) -> tuple[str | None, dict]:
     return None, fmt.model_dump()
 
 
+def _mtq_candidate_formats(formats) -> list[dict]:
+    """Translate recipe candidate formats to export-compatible mtq configs."""
+    quantization_formats = []
+    for fmt in formats:
+        preset_name, quant_cfg = _match_candidate_to_preset(fmt)
+        if preset_name is not None and preset_name not in _AUTO_QUANTIZE_QFORMATS:
+            raise ValueError(
+                f"AutoQuantize candidate_formats entry '{preset_name}' is not supported for "
+                "unified checkpoint export. Use an export-compatible format."
+            )
+        if preset_name is None:
+            warnings.warn(
+                "An AutoQuantize candidate_formats entry matches no shipped preset; its export "
+                "compatibility cannot be verified. Ensure it is safe for HF checkpoint export."
+            )
+        quantization_formats.append(quant_cfg)
+    return quantization_formats
+
+
 def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) -> dict:
     """Map a resolved AutoQuantizeConfig to mtq.auto_quantize inputs.
 
@@ -327,23 +346,19 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) -
     # Translate each candidate to its mtq preset dict and, in the same pass, guard export
     # compatibility (fails fast, before the expensive search). Custom configs matching no shipped
     # preset can't be verified, so warn rather than block.
-    quantization_formats = []
-    for fmt in aq_config.candidate_formats:
-        preset_name, quant_cfg = _match_candidate_to_preset(fmt)
-        if preset_name is not None and preset_name not in _AUTO_QUANTIZE_QFORMATS:
-            raise ValueError(
-                f"AutoQuantize candidate_formats entry '{preset_name}' is not supported for "
-                "unified checkpoint export. Use an export-compatible format."
-            )
-        if preset_name is None:
-            warnings.warn(
-                "An AutoQuantize candidate_formats entry matches no shipped preset; its export "
-                "compatibility cannot be verified. Ensure it is safe for HF checkpoint export."
-            )
-        quantization_formats.append(quant_cfg)
+    quantization_formats = _mtq_candidate_formats(aq_config.candidate_formats)
+    module_search_spaces = [
+        {
+            "module_name_patterns": search_space.module_name_patterns,
+            "quantization_formats": _mtq_candidate_formats(search_space.candidate_formats),
+            "allow_no_quant": search_space.allow_no_quant,
+        }
+        for search_space in aq_config.module_search_spaces
+    ]
     return {
         "constraints": constraints,
         "quantization_formats": quantization_formats,
+        "module_search_spaces": module_search_spaces,
         "disabled_layers": aq_config.disabled_layers,
         "kv_cache_quant_cfg": kv_cache_quant_cfg,
         "method": aq_config.auto_quantize_method,
@@ -467,6 +482,7 @@ def forward_step(model, batch):
         forward_step=forward_step,
         loss_func=loss_func,
         quantization_formats=inputs["quantization_formats"],
+        module_search_spaces=inputs["module_search_spaces"],
         num_calib_steps=len(calib_dataloader),
         num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)),
         verbose=True,
diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py
index 2cf5a2f0cfb..4d0473e0de2 100644
--- a/modelopt/recipe/config.py
+++ b/modelopt/recipe/config.py
@@ -38,6 +38,7 @@
     "AutoQuantizeConfig",
     "AutoQuantizeConstraints",
     "AutoQuantizeCost",
+    "AutoQuantizeModuleSearchSpace",
     "ModelOptAutoQuantizeRecipe",
     "ModelOptDFlashRecipe",
     "ModelOptEagleRecipe",
@@ -191,6 +192,46 @@ def _validate_effective_bits(cls, v: float) -> float:
         return v
 
 
+class AutoQuantizeModuleSearchSpace(ModeloptBaseConfig):
+    """Candidate formats selectable for modules matching one or more name patterns."""
+
+    module_name_patterns: LayerPatternList = ModeloptField(
+        default=[],
+        title="Module name patterns",
+        description="Glob patterns matched against quantizable module names. A grouped AutoQuantize "
+        "decision must match a rule for every module in the group or for none of them.",
+        validate_default=True,
+    )
+    candidate_formats: list[QuantizeConfig] = ModeloptField(
+        default=[],
+        title="Module candidate quantization formats",
+        description="Formats selectable for matching modules. These override the top-level "
+        "candidate_formats for the matching AutoQuantize decision group.",
+        validate_default=True,
+    )
+    allow_no_quant: bool = ModeloptField(
+        default=True,
+        title="Allow no-quant selection",
+        description="Whether BF16/no-quant is selectable for matching modules. AutoQuantize keeps "
+        "an internal no-quant baseline for sensitivity scoring and cost normalization even when "
+        "this is false.",
+    )
+
+    @field_validator("module_name_patterns")
+    @classmethod
+    def _at_least_one_module_pattern(cls, v: list[str]) -> list[str]:
+        if not v:
+            raise ValueError("module_search_spaces requires at least 1 module_name_pattern")
+        return v
+
+    @field_validator("candidate_formats")
+    @classmethod
+    def _at_least_one_module_candidate(cls, v: list[QuantizeConfig]) -> list[QuantizeConfig]:
+        if not v:
+            raise ValueError("module_search_spaces requires at least 1 candidate_format")
+        return v
+
+
 class AutoQuantizeConfig(ModeloptBaseConfig):
     """Schema for the ``auto_quantize`` block of an AutoQuantize recipe."""
 
@@ -206,6 +247,12 @@ class AutoQuantizeConfig(ModeloptBaseConfig):
         "(e.g. [fp8]) yields a {fp8, bf16} per-layer search.",
         validate_default=True,
     )
+    module_search_spaces: list[AutoQuantizeModuleSearchSpace] = ModeloptField(
+        default=[],
+        title="Module-specific search spaces",
+        description="Optional per-module overrides for candidate formats and BF16/no-quant "
+        "selectability. Matching is performed after runtime-fusion grouping.",
+    )
     auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField(
         default="gradient",
         title="Sensitivity scoring method",
diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index c7803ecf838..54daeef1bb5 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -294,9 +294,19 @@ def __init__(
         name: str | None = None,
         quant_module_names: list[str] | None = None,
         cost_weight: float = 1.0,
+        allow_no_quant: bool = True,
     ) -> None:
-        """Initializes Hparam with original value and choices."""
-        choices = sorted({*(choices or []), QuantRecipe(quant_cfg=None)})
+        """Initializes Hparam with internal scoring choices and solver selectability."""
+        candidate_choices = sorted(set(choices or []))
+        # A one-format rule with no-quant disallowed is genuinely fixed: keep that
+        # format active while other groups are scored. Multi-format rules retain an
+        # internal no-quant reference for sensitivity estimation, then filter it out
+        # before LP selection.
+        choices = (
+            candidate_choices
+            if not allow_no_quant and len(candidate_choices) == 1
+            else sorted({*candidate_choices, QuantRecipe(quant_cfg=None)})
+        )
         super().__init__(choices, original=choices[0])
 
         self.name = name
@@ -307,6 +317,7 @@ def __init__(
         }
         assert cost_weight >= 0.0, "cost_weight must be non-negative."
         self.cost_weight = cost_weight
+        self.allow_no_quant = allow_no_quant
 
         self.quant_modules = list(set(quant_modules or []))
         self.score_modules = list(set(score_modules or self.quant_modules))
@@ -440,7 +451,7 @@ def get_cost(self, recipe: QuantRecipe, cost_weight: float | None = None) -> flo
     @property
     def attrs(self) -> list[str]:
         """Return the attributes of the hparam for repr."""
-        return ["name", "cost_weight", *super().attrs]
+        return ["name", "cost_weight", "allow_no_quant", *super().attrs]
 
 
 _LINEAR_ATTN_QKVZ_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_qkv|in_proj_z)$")
@@ -457,6 +468,18 @@ def _linear_attn_ba_group_key(_model, name: str) -> str | None:
     return f"{m.group(1)}/ba" if m else None
 
 
+def _module_search_space_signature(module_search_spaces) -> tuple:
+    """Return a checkpoint-stable description of module-specific candidate spaces."""
+    return tuple(
+        (
+            tuple(search_space["module_name_patterns"]),
+            tuple(str(recipe) for recipe in search_space["quant_recipes"]),
+            search_space["allow_no_quant"],
+        )
+        for search_space in module_search_spaces
+    )
+
+
 class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
     """Base searcher for AutoQuantize algorithm."""
 
@@ -496,6 +519,7 @@ def default_search_config(self):
         """Get the default config for the searcher."""
         return {
             "quantization_formats": ["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"],
+            "module_search_spaces": [],
             "data_loader": None,
             "num_calib_steps": 512,
             "num_score_steps": 128,
@@ -517,6 +541,7 @@ def default_state_dict(self) -> SearchStateDict:
             "cost": {},
             "active_moe_expert_ratio": None,
             "cost_denominator": None,
+            "module_search_space_signature": None,
             "disabled_layers": None,
             "candidate_stats": defaultdict(dict),
             "quantizer_states": {},
@@ -626,7 +651,49 @@ def _get_score_module_from_name(
             )
             return quant_module
 
-    def insert_hparams_after_merge_rules(self, model, quant_recipes, disabled_layers=None):
+    def _normalize_module_search_spaces(self, module_search_spaces):
+        """Convert processed API search spaces to QuantRecipe-based rules."""
+        return [
+            {
+                "module_name_patterns": tuple(search_space["module_name_patterns"]),
+                "quant_recipes": self._get_search_recipes(search_space["quantization_formats"]),
+                "allow_no_quant": search_space["allow_no_quant"],
+            }
+            for search_space in module_search_spaces
+        ]
+
+    @staticmethod
+    def _match_module_search_space(quant_module_names, module_search_spaces):
+        """Return the unique rule that fully covers a runtime-grouped decision."""
+        matched_search_spaces = []
+        for search_space in module_search_spaces:
+            matches = [
+                any(
+                    fnmatch.fnmatch(module_name, pattern)
+                    for pattern in search_space["module_name_patterns"]
+                )
+                for module_name in quant_module_names
+            ]
+            if not any(matches):
+                continue
+            if not all(matches):
+                raise ValueError(
+                    "A module_search_spaces rule partially matches runtime-grouped modules "
+                    f"{quant_module_names}. Update its module_name_patterns so the rule covers "
+                    "the entire group or none of it."
+                )
+            matched_search_spaces.append(search_space)
+
+        if len(matched_search_spaces) > 1:
+            raise ValueError(
+                "Multiple module_search_spaces rules match runtime-grouped modules "
+                f"{quant_module_names}. Make the module_name_patterns disjoint."
+            )
+        return matched_search_spaces[0] if matched_search_spaces else None
+
+    def insert_hparams_after_merge_rules(
+        self, model, quant_recipes, disabled_layers=None, module_search_spaces=None
+    ):
         """Restrict the search space using the merge rules and insert the hparams for the model."""
         # TRTLLM fuses linear layers such as q_proj, k_proj, v_proj into same layer
         # Hence we need to restrict the search space so that all these layers share the same recipe
@@ -686,7 +753,18 @@ def insert_hparams_after_merge_rules(self, model, quant_recipes, disabled_layers
                 quant_module_names, self.config["cost"]
             )
 
-            _quant_recipes = None if disabled else quant_recipes
+            search_space = self._match_module_search_space(
+                quant_module_names, module_search_spaces or []
+            )
+            if disabled:
+                _quant_recipes = None
+                allow_no_quant = True
+            elif search_space is not None:
+                _quant_recipes = search_space["quant_recipes"]
+                allow_no_quant = search_space["allow_no_quant"]
+            else:
+                _quant_recipes = quant_recipes
+                allow_no_quant = True
             hparam = QuantRecipeHparam(
                 _quant_recipes,
                 quant_modules=quant_modules,
@@ -694,6 +772,7 @@ def insert_hparams_after_merge_rules(self, model, quant_recipes, disabled_layers
                 name=str(group_key),
                 quant_module_names=quant_module_names,
                 cost_weight=cost_weight,
+                allow_no_quant=allow_no_quant,
             )
 
             for module in quant_modules:
@@ -721,6 +800,7 @@ def estimate_sensitivity_scores(self) -> None:
 
     def initialize_candidate_stats(self):
         """Initialize the candidate stats for the model."""
+        no_quant_recipe = QuantRecipe(quant_cfg=None)
         for name, hparam in named_hparams(self.model, unique=True):
             if not isinstance(hparam, QuantRecipeHparam):
                 continue
@@ -728,6 +808,8 @@ def initialize_candidate_stats(self):
             formats, scores, costs = [], [], []
             prev_score = float("inf")
             for recipe in hparam.choices:
+                if recipe == no_quant_recipe and not hparam.allow_no_quant:
+                    continue
                 formats.append(recipe)
 
                 score = hparam.get_score(recipe)  # type: ignore [arg-type]
@@ -744,6 +826,10 @@ def initialize_candidate_stats(self):
             self.candidate_stats[name]["module_names"] = hparam.quant_module_names
             self.candidate_stats[name]["quantizer_attrs"] = hparam.quant_module_replay_attrs
             self.candidate_stats[name]["cost_weight"] = hparam.cost_weight
+            self.candidate_stats[name]["allow_no_quant"] = hparam.allow_no_quant
+            # Keep the no-quant cost as denominator metadata even when no-quant is not
+            # solver-selectable for this hparam. Fixed formats must remain in the cost model.
+            self.candidate_stats[name]["uncompressed_cost"] = hparam.get_cost(no_quant_recipe)
 
     def _run_func(self, func, num_iters=1, desc=""):
         for i, data in tqdm(
@@ -794,11 +880,44 @@ def before_search(self):
         self.disabled_layers = self.config["disabled_layers"]
         self.cost_denominator = getattr(self, "cost_denominator", None)
 
-        search_recipes = self._get_search_recipes(self.config["quantization_formats"])
+        module_search_spaces = self._normalize_module_search_spaces(
+            self.config["module_search_spaces"]
+        )
+        module_search_space_signature = _module_search_space_signature(module_search_spaces)
+        restored_module_search_space_signature = getattr(
+            self, "module_search_space_signature", None
+        )
+        if self.candidate_stats and (
+            (restored_module_search_space_signature is None and module_search_space_signature)
+            or (
+                restored_module_search_space_signature is not None
+                and restored_module_search_space_signature != module_search_space_signature
+            )
+        ):
+            raise ValueError(
+                "Checkpoint module_search_spaces do not match the current search config. "
+                "Use a different checkpoint path."
+            )
+        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,
+                *(
+                    recipe
+                    for search_space in module_search_spaces
+                    for recipe in search_space["quant_recipes"]
+                ),
+            }
+        )
         self._verify_constraint(search_recipes)
         self._cost_model = cost_model
         self.insert_hparams_after_merge_rules(
-            self.model, search_recipes, self.config["disabled_layers"]
+            self.model,
+            default_search_recipes,
+            self.config["disabled_layers"],
+            module_search_spaces,
         )
 
         QuantRecipe.disable_folding_pqs_to_weights()
@@ -809,10 +928,13 @@ def before_search(self):
             if recipe == QuantRecipe(quant_cfg=None):  # No-quant format
                 continue
 
+            no_quant_recipe = QuantRecipe(quant_cfg=None)
             for name, hparam in named_hparams(self.model, configurable=True):
                 if not isinstance(hparam, QuantRecipeHparam):
                     continue
-                hparam.active = recipe
+                hparam.active = no_quant_recipe
+                if recipe in hparam.choices:
+                    hparam.active = recipe
 
             if recipe in self.quantizer_states:
                 saved = self.quantizer_states[recipe]
@@ -891,6 +1013,9 @@ def _get_total_weight_size_from_candidate_stats(candidate_stats):
         no_quant_recipe = QuantRecipe(quant_cfg=None)
         total_weight_size = 0
         for candidate_stat in candidate_stats.values():
+            if "uncompressed_cost" in candidate_stat:
+                total_weight_size += candidate_stat["uncompressed_cost"]
+                continue
             no_quant_idx = candidate_stat["formats"].index(no_quant_recipe)
             total_weight_size += candidate_stat["costs"][no_quant_idx]
         return total_weight_size
@@ -1618,7 +1743,7 @@ def _resolve_best_recipe(search_state, constraints, verbose=False):
     compression = effective_bits / 16.0
     candidate_stats = search_state["candidate_stats"]
     total_weight_size = search_state.get("cost_denominator") or sum(
-        s["costs"][-1] for s in candidate_stats.values()
+        s.get("uncompressed_cost", max(s["costs"])) for s in candidate_stats.values()
     )
     max_weight_size = total_weight_size * compression
     method = search_state["method"]
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 1adda0511a6..8747adaf4d2 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/modelopt/torch/quantization/model_quant.py
@@ -283,6 +283,7 @@ def auto_quantize(
     verbose: bool = False,
     method: str = "gradient",
     checkpoint: str | None = None,
+    module_search_spaces: list[dict[str, Any]] | None = None,
 ):
     r"""Perform optimal per-layer quantization by searching for the best quantization formats per-layer.
 
@@ -446,6 +447,14 @@ def forward_backward_step(model, batch) -> None:
         checkpoint: (Optional) Path to checkpoint file for saving/restoring auto_quantize search state.
             If the checkpoint file exists, the search state will be restored from it, skipping the
             expensive score estimation step.
+        module_search_spaces: Optional module-specific candidate overrides. Each entry contains
+            ``module_name_patterns`` (one or more glob patterns), ``quantization_formats``, and
+            optional ``allow_no_quant`` (default True). A matching entry replaces the global
+            candidate set for that runtime-grouped decision. Setting ``allow_no_quant=False``
+            keeps BF16/no-quant as an internal sensitivity and cost baseline but prevents the
+            solver from selecting it. A single candidate with ``allow_no_quant=False`` fixes the
+            matching module group to that format while retaining its cost in the effective-bits
+            constraint.
 
     Returns: A tuple (model, state_dict) where ``model`` is the searched and quantized model and
         ``state_dict`` contains the history and detailed stats of the search procedure.
@@ -493,23 +502,76 @@ def forward_backward_step(model, batch) -> None:
         might not be readily deployable to TensorRT-LLM yet.
 
     """
-    processed_quantization_formats = []
-    for i, quant_cfg in enumerate(quantization_formats):
-        if quant_cfg is None:
-            continue
 
-        name = QuantRecipe.get_auto_name_for_config(quant_cfg)
-        if name is None:
-            name = f"CUSTOM_{i}"
-            warnings.warn(
-                f"Received custom quantization formats for search, auto_quantize results may not be optimal. "
-                f"This config will be displayed as {name}"
-            )
-        processed_quantization_formats.append((quant_cfg, name))
+    def _process_quantization_formats(formats, custom_name_prefix):
+        processed = []
+        for i, quant_cfg in enumerate(formats):
+            if quant_cfg is None:
+                continue
+
+            name = QuantRecipe.get_auto_name_for_config(quant_cfg)
+            if name is None:
+                name = f"{custom_name_prefix}_{i}"
+                warnings.warn(
+                    "Received custom quantization formats for search, auto_quantize results may "
+                    f"not be optimal. This config will be displayed as {name}"
+                )
+            processed.append((quant_cfg, name))
+        return processed
+
+    processed_quantization_formats = _process_quantization_formats(quantization_formats, "CUSTOM")
 
     assert len(processed_quantization_formats) > 0, "`quantization_formats` should not be empty"
 
-    for quant_cfg, name in processed_quantization_formats:
+    processed_module_search_spaces = []
+    for idx, search_space in enumerate(module_search_spaces or []):
+        if not isinstance(search_space, dict):
+            raise TypeError("Each module_search_spaces entry must be a dict.")
+        unknown_keys = set(search_space) - {
+            "module_name_patterns",
+            "quantization_formats",
+            "allow_no_quant",
+        }
+        if unknown_keys:
+            raise ValueError(f"Unsupported module_search_spaces keys: {sorted(unknown_keys)}")
+        patterns = search_space.get("module_name_patterns")
+        if isinstance(patterns, str):
+            patterns = [patterns]
+        if (
+            not isinstance(patterns, list)
+            or not patterns
+            or not all(isinstance(pattern, str) for pattern in patterns)
+        ):
+            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}"
+        )
+        if not formats:
+            raise ValueError(
+                "module_search_spaces.quantization_formats must contain at least one format."
+            )
+        allow_no_quant = search_space.get("allow_no_quant", True)
+        if not isinstance(allow_no_quant, bool):
+            raise TypeError("module_search_spaces.allow_no_quant must be a bool.")
+        processed_module_search_spaces.append(
+            {
+                "module_name_patterns": patterns,
+                "quantization_formats": formats,
+                "allow_no_quant": allow_no_quant,
+            }
+        )
+
+    all_processed_formats = [
+        *processed_quantization_formats,
+        *(
+            quant_format
+            for search_space in processed_module_search_spaces
+            for quant_format in search_space["quantization_formats"]
+        ),
+    ]
+    for quant_cfg, name in all_processed_formats:
         algo = QuantRecipe(quant_cfg, name=name).config.algorithm
         algo_method = algo["method"] if isinstance(algo, dict) else algo
         if algo_method not in _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS:
@@ -534,6 +596,7 @@ def forward_backward_step(model, batch) -> None:
     )
     search_config = {
         "quantization_formats": processed_quantization_formats,
+        "module_search_spaces": processed_module_search_spaces,
         "data_loader": data_loader,
         "forward_step": forward_step,
         "loss_func": loss_func,
diff --git a/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe.yaml b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe.yaml
new file mode 100644
index 00000000000..8a3f4ce38d2
--- /dev/null
+++ b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe.yaml
@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Qwen3.6 MoE AutoQuantize search with module-family-specific candidate spaces.
+# Routed experts are fixed to W4A16 NVFP4 but remain in active-MoE effective-bit
+# accounting. Shared experts and hybrid-attention projections select between
+# W4A16 NVFP4 and FP8. Routed experts remain W4A16 during scoring; searched
+# multi-format families retain BF16 only as an internal scoring baseline.
+
+# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe
+imports:
+  base_disabled_layers: configs/auto_quantize/units/base_disabled_layers
+  fp8: configs/ptq/presets/model/fp8
+  w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4
+
+metadata:
+  recipe_type: auto_quantize
+  description: >-
+    Qwen3.6 MoE module-specific search: routed experts fixed to W4A16 NVFP4,
+    shared experts and attention searched over W4A16 NVFP4 and FP8, at 6.0
+    active-MoE effective bits.
+
+auto_quantize:
+  constraints:
+    effective_bits: 6.0
+    cost_model: active_moe
+    cost:
+      active_moe_expert_ratio: 0.03125
+
+  # Fallback for unmatched searchable modules, including lm_head. BF16/no-quant
+  # remains an implicit choice only for this fallback search space.
+  candidate_formats:
+    - $import: w4a16_nvfp4
+    - $import: fp8
+
+  module_search_spaces:
+    # One selectable format makes routed experts fixed W4A16. allow_no_quant
+    # controls solver choices only; their BF16-equivalent denominator cost and
+    # active-MoE ratio remain part of effective-bit accounting.
+    - module_name_patterns:
+        - "*mlp.experts*"
+        - "*mixer.experts*"
+      candidate_formats:
+        - $import: w4a16_nvfp4
+      allow_no_quant: false
+
+    # Search these active/sensitive families only between W4A16 and W8A8/FP8.
+    - module_name_patterns:
+        - "*mlp.shared_expert*"
+        - "*linear_attn*"
+        - "*self_attn*"
+      candidate_formats:
+        - $import: w4a16_nvfp4
+        - $import: fp8
+      allow_no_quant: false
+
+  auto_quantize_method: gradient
+  score_size: 128
+
+  disabled_layers:
+    - $import: base_disabled_layers
+
+  cost_excluded_layers:
+    - "*visual*"
+    - "*mtp*"
+    - "*vision_tower*"
diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py
index 7160dd38cb4..3e2f9047b7e 100644
--- a/tests/examples/hf_ptq/test_hf_ptq_args.py
+++ b/tests/examples/hf_ptq/test_hf_ptq_args.py
@@ -57,6 +57,7 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch):
     assert inputs["kv_cache_quant_cfg"] is None
     assert inputs["method"] == "gradient"
     assert inputs["score_size"] == 128
+    assert inputs["module_search_spaces"] == []
     # disabled_layers come straight from the recipe (no model introspection).
     assert inputs["disabled_layers"] == aq.disabled_layers
     assert "*output_layer*" in inputs["disabled_layers"]
@@ -87,6 +88,32 @@ def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch):
     assert "*visual*" in inputs["disabled_layers"]
 
 
+def test_autoquant_recipe_maps_module_search_spaces(monkeypatch):
+    """Module-specific recipe candidates map to public mtq.auto_quantize inputs."""
+    hf_ptq, args = _parse_hf_ptq_args(
+        monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none"
+    )
+    aq = load_recipe(
+        "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe"
+    ).auto_quantize
+    inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args)
+
+    routed, searched = inputs["module_search_spaces"]
+    assert routed["module_name_patterns"] == ["*mlp.experts*", "*mixer.experts*"]
+    assert routed["quantization_formats"] == [QUANT_CFG_CHOICES["w4a16_nvfp4"]]
+    assert routed["allow_no_quant"] is False
+    assert searched["module_name_patterns"] == [
+        "*mlp.shared_expert*",
+        "*linear_attn*",
+        "*self_attn*",
+    ]
+    assert searched["quantization_formats"] == [
+        QUANT_CFG_CHOICES["w4a16_nvfp4"],
+        QUANT_CFG_CHOICES["fp8"],
+    ]
+    assert searched["allow_no_quant"] is False
+
+
 def test_autoquant_rejects_non_export_safe_candidate(monkeypatch):
     """A candidate that resolves to a preset outside the export-safe set is rejected before search."""
     hf_ptq, args = _parse_hf_ptq_args(
diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py
index 3aaacaa3e0e..4e4d599dcff 100644
--- a/tests/unit/recipe/test_loader.py
+++ b/tests/unit/recipe/test_loader.py
@@ -1726,6 +1726,7 @@ def test_load_recipe_autoquantize_minimal(tmp_path):
     assert aq.constraints.cost_model == "weight"
     assert aq.constraints.cost is None
     assert len(aq.candidate_formats) == 2
+    assert aq.module_search_spaces == []
 
 
 def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path):
@@ -1812,6 +1813,26 @@ def test_load_recipe_autoquantize_builtin_active_moe():
     assert all(c.effective_bits is None for c in aq.candidate_formats)
 
 
+def test_load_recipe_autoquantize_module_search_spaces():
+    """Qwen module rules resolve imported formats and no-quant selectability."""
+    recipe = load_recipe(
+        "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_module_spaces_at_6p0bits-active_moe"
+    )
+    aq = recipe.auto_quantize
+    assert len(aq.module_search_spaces) == 2
+    routed, searched = aq.module_search_spaces
+    assert routed.module_name_patterns == ["*mlp.experts*", "*mixer.experts*"]
+    assert len(routed.candidate_formats) == 1
+    assert routed.allow_no_quant is False
+    assert searched.module_name_patterns == [
+        "*mlp.shared_expert*",
+        "*linear_attn*",
+        "*self_attn*",
+    ]
+    assert len(searched.candidate_formats) == 2
+    assert searched.allow_no_quant is False
+
+
 @pytest.mark.parametrize(
     "recipe_path",
     [
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index b85feb32649..c87d961d89e 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -269,6 +269,99 @@ def test_auto_quantize_active_moe_cost_model(num_experts_attr):
     assert all("active_costs" not in stats for stats in search_history["candidate_stats"].values())
 
 
+def test_auto_quantize_module_search_spaces_keep_fixed_routed_experts_costed():
+    """A one-format routed rule is fixed in the LP but retained in active-MoE cost."""
+    model = _AutoQuantMoeModel()
+    int4_recipe = QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    int8_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG)
+    no_quant_recipe = QuantRecipe(None)
+
+    searched_model, search_history = mtq.auto_quantize(
+        model,
+        constraints={"effective_bits": 6.0, "cost_model": "active_moe"},
+        quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG],
+        module_search_spaces=[
+            {
+                "module_name_patterns": ["*mlp.experts*"],
+                "quantization_formats": [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG],
+                "allow_no_quant": False,
+            },
+            {
+                "module_name_patterns": ["*mlp.shared_expert*"],
+                "quantization_formats": [
+                    mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
+                    mtq.INT8_DEFAULT_CFG,
+                ],
+                "allow_no_quant": False,
+            },
+        ],
+        data_loader=[model.get_input() for _ in range(2)],
+        forward_step=lambda model, batch: model(batch),
+        loss_func=lambda output, data: output.sum(),
+        num_calib_steps=2,
+        num_score_steps=2,
+    )
+
+    stats = search_history["candidate_stats"]
+    routed = next(
+        candidate
+        for candidate in stats.values()
+        if any("mlp.experts" in name for name in candidate["module_names"])
+    )
+    shared = next(
+        candidate
+        for candidate in stats.values()
+        if any("mlp.shared_expert" in name for name in candidate["module_names"])
+    )
+
+    assert routed["formats"] == [int4_recipe]
+    assert routed["allow_no_quant"] is False
+    assert routed["cost_weight"] == pytest.approx(0.25)
+    assert routed["costs"] == pytest.approx([routed["uncompressed_cost"] * 0.25])
+    assert no_quant_recipe not in routed["formats"]
+    assert shared["formats"] == [int4_recipe, int8_recipe]
+    assert shared["allow_no_quant"] is False
+    assert no_quant_recipe not in shared["formats"]
+    assert search_history["cost_denominator"] == pytest.approx(
+        sum(candidate["uncompressed_cost"] for candidate in stats.values())
+    )
+    routed_best = [
+        recipe
+        for name, recipe in search_history["best"]["recipe"].items()
+        if any("mlp.experts" in module for module in stats[name]["module_names"])
+    ]
+    assert routed_best and all(recipe == int4_recipe for recipe in routed_best)
+    routed_hparam = searched_model.mlp.experts[0].gate_proj.get_hparam("quant_recipe")
+    assert not routed_hparam.is_configurable
+    assert routed_hparam.active == int4_recipe
+
+
+def test_auto_quantize_module_search_space_cannot_split_runtime_group():
+    model = TransformerBlock()
+
+    with pytest.raises(ValueError, match="partially matches runtime-grouped modules"):
+        mtq.auto_quantize(
+            model,
+            constraints={"effective_bits": 6.0},
+            quantization_formats=[
+                mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
+                mtq.INT8_DEFAULT_CFG,
+            ],
+            module_search_spaces=[
+                {
+                    "module_name_patterns": ["*q_proj"],
+                    "quantization_formats": [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG],
+                    "allow_no_quant": False,
+                }
+            ],
+            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,
+        )
+
+
 def test_active_moe_ratio_requires_single_config_object():
     model = torch.nn.Module()
     model.config = SimpleNamespace(

From 5041a98be74a994726cff453638985933bf022e3 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Wed, 8 Jul 2026 18:07:50 -0700
Subject: [PATCH 2/5] Document module-specific AutoQuant choices

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 examples/hf_ptq/README.md | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md
index 3a7380a9f9e..09aeaa29f98 100755
--- a/examples/hf_ptq/README.md
+++ b/examples/hf_ptq/README.md
@@ -382,12 +382,16 @@ The recipe quantizes the less accuracy-sensitive layers with the more aggressive
 keeps the more sensitive ones at higher precision (or unquantized), so the model meets the recipe's
 `effective_bits` target. To author your own, copy a shipped recipe and adjust `candidate_formats`,
 `constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div`), `score_size`,
-`disabled_layers` (excluded from the search), and `cost_excluded_layers` (kept out of the bit-budget
-accounting — e.g. VL vision towers). Recipes can splice a shared base `disabled_layers` set via
-`$import` (see `modelopt_recipes/configs/auto_quantize/units/base_disabled_layers`).
-
-bf16 (no quantization) is always an implicit per-layer choice, so `candidate_formats` need only list
-the quantized options — a single format (e.g. `[fp8]`) gives a `{fp8, bf16}` per-layer search.
+`module_search_spaces` (optional per-module candidate overrides), `disabled_layers` (excluded from
+the search), and `cost_excluded_layers` (kept out of the bit-budget accounting — e.g. VL vision
+towers). Recipes can splice a shared base `disabled_layers` set via `$import` (see
+`modelopt_recipes/configs/auto_quantize/units/base_disabled_layers`).
+
+bf16 (no quantization) is an implicit per-layer choice for the top-level `candidate_formats`, so a
+single format (e.g. `[fp8]`) gives a `{fp8, bf16}` per-layer search. A `module_search_spaces` rule can
+set `allow_no_quant: false` to exclude bf16 from the solver choices for matching modules. A rule with
+one candidate format then fixes those modules to that format while retaining their uncompressed
+weight in the effective-bit denominator and their selected-format cost in the numerator.
 
 For models without backprop support (e.g. Llama-4), use the `kl_div` scoring method — see the shipped
 `general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits` recipe.

From c6154844c9e397128740fec3052c6212444eb168 Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 9 Jul 2026 16:17:31 -0700
Subject: [PATCH 3/5] Fix module search calibration and checkpoint identity

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/algorithms.py     | 145 ++++++++++-------
 .../unit/torch/quantization/test_autoquant.py | 154 ++++++++++++++++++
 2 files changed, 243 insertions(+), 56 deletions(-)

diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index 54daeef1bb5..50a7a1133d1 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/modelopt/torch/quantization/algorithms.py
@@ -223,6 +223,12 @@ def __init__(self, quant_cfg: str | dict[str, Any] | None = None, name: str | No
         self.compression = estimate_quant_compression(self.config)
 
         self._str_repr: str = f"{name}(effective-bits: {self.compression * 16})"
+        self._config_signature = self.config.model_dump_json()
+
+    @property
+    def checkpoint_signature(self) -> str:
+        """Return the canonical identity used for ordering and checkpoint validation."""
+        return getattr(self, "_config_signature", self.config.model_dump_json())
 
     @staticmethod
     def get_auto_name_for_config(quant_cfg: str | dict[str, Any] | None) -> str | None:
@@ -248,14 +254,19 @@ def __repr__(self) -> str:
         return self._str_repr
 
     def __lt__(self, other: "QuantRecipe"):
-        return self.compression < other.compression
+        return (self.compression, self.checkpoint_signature) < (
+            other.compression,
+            other.checkpoint_signature,
+        )
 
     def __eq__(self, other: object):
-        assert isinstance(other, QuantRecipe)
-        return self._str_repr == other._str_repr
+        return (
+            isinstance(other, QuantRecipe)
+            and self.checkpoint_signature == other.checkpoint_signature
+        )
 
     def __hash__(self) -> int:
-        return hash(self._str_repr)
+        return hash(self.checkpoint_signature)
 
     @staticmethod
     def disable_folding_pqs_to_weights():
@@ -329,10 +340,12 @@ def __init__(
         # (``*_input_quantizer`` + ``*_weight_quantizers`` ModuleList) — see
         # ``_get_quantizer_attrs``. Both layouts share the same snapshot dict
         # shape so ``active.setter`` swaps the right child modules.
-        self._all_quantizer_choices = {quant_recipe: {} for quant_recipe in self.choices}
+        no_quant_recipe = QuantRecipe(quant_cfg=None)
+        calibration_recipes = sorted({*self.choices, no_quant_recipe})
+        self._all_quantizer_choices = {quant_recipe: {} for quant_recipe in calibration_recipes}
 
         quant_recipe: QuantRecipe
-        for quant_recipe in self.choices:
+        for quant_recipe in calibration_recipes:
             for quant_module in self.quant_modules:
                 attr_names = _get_quantizer_attrs(quant_module)
                 for attr_name in attr_names:
@@ -369,16 +382,31 @@ def active(self) -> HPType:
     def active(self, val: HPType | None):
         """Set the active value with a sanity check for choices and dynamic hparams."""
         val = self.original if val is None else val
+        assert isinstance(val, QuantRecipe)
         assert val in self._choices, f"val = {val}, choices = {self.choices}"
         if self.is_configurable:
             self._active = val
         else:
             assert self._active == val
 
-        for nn_module, quantizer_choices in self._all_quantizer_choices[val].items():
+        self._apply_quantizer_choice(val)
+
+    def _apply_quantizer_choice(self, recipe: QuantRecipe) -> None:
+        for nn_module, quantizer_choices in self._all_quantizer_choices[recipe].items():
             for quantizer_attr_name, quantizer in quantizer_choices.items():
                 setattr(nn_module, quantizer_attr_name, quantizer)
 
+    def set_calibration_recipe(self, recipe: QuantRecipe) -> None:
+        """Enable ``recipe`` for this calibration pass or isolate the group."""
+        calibration_recipe = recipe if recipe in self.choices else QuantRecipe(quant_cfg=None)
+        self._apply_quantizer_choice(calibration_recipe)
+
+    def restore_active_quantizers(self) -> None:
+        """Restore quantizer objects corresponding to the solver-visible active recipe."""
+        active = self.active
+        assert isinstance(active, QuantRecipe)
+        self._apply_quantizer_choice(active)
+
     @property
     def importance(self) -> dict:
         """Raises an error since this is not a useful abstraction for AutoQuantize."""
@@ -473,7 +501,7 @@ def _module_search_space_signature(module_search_spaces) -> tuple:
     return tuple(
         (
             tuple(search_space["module_name_patterns"]),
-            tuple(str(recipe) for recipe in search_space["quant_recipes"]),
+            tuple(sorted(recipe.checkpoint_signature for recipe in search_space["quant_recipes"])),
             search_space["allow_no_quant"],
         )
         for search_space in module_search_spaces
@@ -887,7 +915,8 @@ def before_search(self):
         restored_module_search_space_signature = getattr(
             self, "module_search_space_signature", None
         )
-        if self.candidate_stats and (
+        has_restored_calibration_or_scores = bool(self.quantizer_states or self.candidate_stats)
+        if has_restored_calibration_or_scores and (
             (restored_module_search_space_signature is None and module_search_space_signature)
             or (
                 restored_module_search_space_signature is not None
@@ -924,60 +953,64 @@ def before_search(self):
 
         # Iterate over the search recipes and calibrate the quantizers for each recipe
         calibrated_new = False
-        for recipe in search_recipes:
-            if recipe == QuantRecipe(quant_cfg=None):  # No-quant format
-                continue
+        quant_recipe_hparams = [
+            hparam
+            for _, hparam in named_hparams(self.model, unique=True)
+            if isinstance(hparam, QuantRecipeHparam)
+        ]
+        try:
+            for recipe in search_recipes:
+                if recipe == QuantRecipe(quant_cfg=None):  # No-quant format
+                    continue
 
-            no_quant_recipe = QuantRecipe(quant_cfg=None)
-            for name, hparam in named_hparams(self.model, configurable=True):
-                if not isinstance(hparam, QuantRecipeHparam):
+                for hparam in quant_recipe_hparams:
+                    hparam.set_calibration_recipe(recipe)
+
+                if recipe in self.quantizer_states:
+                    saved = self.quantizer_states[recipe]
+                    # config is unused by restore_quantizer_state
+                    restore_quantizer_state(
+                        self.model, QuantizeConfig(), {"quantizer_state": saved["metadata"]}
+                    )
+                    set_quantizer_state_dict(self.model, saved["state_dict"])
+                    if self.config["verbose"]:
+                        print_rank_0(f"AutoQuantize: Restored calibration for {recipe}")
                     continue
-                hparam.active = no_quant_recipe
-                if recipe in hparam.choices:
-                    hparam.active = recipe
 
-            if recipe in self.quantizer_states:
-                saved = self.quantizer_states[recipe]
-                # config is unused by restore_quantizer_state
-                restore_quantizer_state(
-                    self.model, QuantizeConfig(), {"quantizer_state": saved["metadata"]}
+                # Lets reduce the number of calibration steps for AWQ since it takes longer
+                num_calib_steps = (
+                    self.config["num_calib_steps"]
+                    if "awq" not in str(recipe.config.algorithm)
+                    else max(1, self.config["num_calib_steps"] // 4)
                 )
-                set_quantizer_state_dict(self.model, saved["state_dict"])
-                if self.config["verbose"]:
-                    print_rank_0(f"AutoQuantize: Restored calibration for {recipe}")
-                continue
 
-            # Lets reduce the number of calibration steps for AWQ since it takes longer
-            num_calib_steps = (
-                self.config["num_calib_steps"]
-                if "awq" not in str(recipe.config.algorithm)
-                else max(1, self.config["num_calib_steps"] // 4)
-            )
+                def forward_loop(model):
+                    self._run_func(
+                        self.config["forward_step"],
+                        num_iters=num_calib_steps,
+                        desc=f"Calibrating for {recipe}",
+                    )
 
-            def forward_loop(model):
-                self._run_func(
-                    self.config["forward_step"],
-                    num_iters=num_calib_steps,
-                    desc=f"Calibrating for {recipe}",
+                calibrate(
+                    self.model,
+                    algorithm=recipe.config.algorithm,
+                    forward_loop=forward_loop,
                 )
-
-            calibrate(
-                self.model,
-                algorithm=recipe.config.algorithm,
-                forward_loop=forward_loop,
-            )
-            # Calibrate adds a new mode to the model. Since auto_quantize mixes the quantization recipes
-            # across layers, lets not save this new mode in the modelopt state.
-            # TODO: This is a hack. We need to create a mode for auto_quantize to handle this in a clean way.
-            ModeloptStateManager(self.model).state_dict().pop()
-            metadata: dict = {}
-            # config is unused by update_quantize_metadata
-            update_quantize_metadata(self.model, QuantizeConfig(), metadata)
-            self.quantizer_states[recipe] = {
-                "metadata": metadata["quantizer_state"],
-                "state_dict": get_quantizer_state_dict(self.model),
-            }
-            calibrated_new = True
+                # Calibrate adds a new mode to the model. Since auto_quantize mixes the quantization recipes
+                # across layers, lets not save this new mode in the modelopt state.
+                # TODO: This is a hack. We need to create a mode for auto_quantize to handle this in a clean way.
+                ModeloptStateManager(self.model).state_dict().pop()
+                metadata: dict = {}
+                # config is unused by update_quantize_metadata
+                update_quantize_metadata(self.model, QuantizeConfig(), metadata)
+                self.quantizer_states[recipe] = {
+                    "metadata": metadata["quantizer_state"],
+                    "state_dict": get_quantizer_state_dict(self.model),
+                }
+                calibrated_new = True
+        finally:
+            for hparam in quant_recipe_hparams:
+                hparam.restore_active_quantizers()
 
         if calibrated_new:
             self.save_search_checkpoint(verbose=self.config["verbose"])
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index c87d961d89e..06d8258cd37 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/tests/unit/torch/quantization/test_autoquant.py
@@ -35,6 +35,7 @@
     QuantRecipe,
     QuantRecipeHparam,
     _AutoQuantizeBaseSearcher,
+    _module_search_space_signature,
     estimate_quant_compression,
 )
 from modelopt.torch.quantization.config import _base_disable_all, _default_disabled_quantizer_cfg
@@ -113,6 +114,35 @@ def test_quant_recipe(quant_cfg, other_quant_cfg, is_less_than):
     assert qr_this_duplicate == qr_this
 
 
+def test_module_search_space_signature_is_canonical_and_uses_full_configs():
+    fp8_recipe = QuantRecipe(mtq.FP8_DEFAULT_CFG)
+    int8_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG)
+
+    def signature(recipes):
+        return _module_search_space_signature(
+            [
+                {
+                    "module_name_patterns": ("*mlp*",),
+                    "quant_recipes": recipes,
+                    "allow_no_quant": False,
+                }
+            ]
+        )
+
+    assert signature([fp8_recipe, int8_recipe]) == signature([int8_recipe, fp8_recipe])
+
+    custom_max_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    custom_max_cfg["quant_cfg"].append({"quantizer_name": "unused_quantizer", "enable": False})
+    custom_mse_cfg = copy.deepcopy(custom_max_cfg)
+    custom_mse_cfg["algorithm"] = "mse"
+    max_recipe = QuantRecipe(custom_max_cfg, name="CUSTOM_MODULE_0_0")
+    mse_recipe = QuantRecipe(custom_mse_cfg, name="CUSTOM_MODULE_0_0")
+
+    assert str(max_recipe) == str(mse_recipe)
+    assert max_recipe != mse_recipe
+    assert signature([max_recipe]) != signature([mse_recipe])
+
+
 def test_quant_recipe_hparam():
     model_test = torch.nn.Linear(4, 16)
     model_ref = torch.nn.Linear(4, 16)
@@ -336,6 +366,58 @@ 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
+
+    model = TransformerBlock()
+    calibration_states = []
+    original_calibrate = model_quant.calibrate
+
+    def recording_calibrate(model, algorithm="max", forward_loop=None):
+        algorithm_name = algorithm.get("method") if isinstance(algorithm, dict) else algorithm
+        weight_quantizer = model.mlp.weight_quantizer
+        calibration_states.append(
+            {
+                "algorithm": algorithm_name,
+                "enabled": weight_quantizer.is_enabled,
+                "num_bits": weight_quantizer.num_bits,
+                "quantizer_id": id(weight_quantizer),
+            }
+        )
+        return original_calibrate(model, algorithm=algorithm, forward_loop=forward_loop)
+
+    monkeypatch.setattr(model_quant, "calibrate", recording_calibrate)
+
+    searched_model, _ = mtq.auto_quantize(
+        model,
+        constraints={"effective_bits": 6.0},
+        quantization_formats=[mtq.INT4_AWQ_CFG],
+        module_search_spaces=[
+            {
+                "module_name_patterns": ["*mlp*"],
+                "quantization_formats": [mtq.FP8_DEFAULT_CFG],
+                "allow_no_quant": False,
+            }
+        ],
+        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,
+    )
+
+    awq_state = next(state for state in calibration_states if "awq" in state["algorithm"])
+    fp8_state = next(state for state in calibration_states if state["algorithm"] == "max")
+    assert not awq_state["enabled"]
+    assert fp8_state["enabled"]
+    assert fp8_state["num_bits"] == (4, 3)
+    assert awq_state["quantizer_id"] != fp8_state["quantizer_id"]
+    assert id(searched_model.mlp.weight_quantizer) == fp8_state["quantizer_id"]
+    assert searched_model.mlp.weight_quantizer.pre_quant_scale is None
+    assert not searched_model.mlp.get_hparam("quant_recipe").is_configurable
+    assert searched_model.mlp.get_hparam("quant_recipe").active == QuantRecipe(mtq.FP8_DEFAULT_CFG)
+
+
 def test_auto_quantize_module_search_space_cannot_split_runtime_group():
     model = TransformerBlock()
 
@@ -857,6 +939,78 @@ def test_auto_quantize_checkpoint_resume(method, tmp_path, capsys):
                 )
 
 
+def test_auto_quantize_calibration_only_checkpoint_validates_module_search_spaces(
+    tmp_path, monkeypatch
+):
+    checkpoint_path = str(tmp_path / "autoquant_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,
+            ],
+            module_search_spaces=[
+                {
+                    "module_name_patterns": ["*mlp*"],
+                    "quantization_formats": [mtq.FP8_DEFAULT_CFG],
+                    "allow_no_quant": False,
+                }
+            ],
+            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 not saved["candidate_stats"]
+
+    monkeypatch.setattr(
+        AutoQuantizeGradientSearcher,
+        "estimate_sensitivity_scores",
+        original_estimate_scores,
+    )
+    resumed_model = TransformerBlock()
+    with pytest.raises(ValueError, match="module_search_spaces do not match"):
+        mtq.auto_quantize(
+            resumed_model,
+            constraints={"effective_bits": 6.0},
+            quantization_formats=[
+                mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG,
+                mtq.INT8_DEFAULT_CFG,
+            ],
+            module_search_spaces=[
+                {
+                    "module_name_patterns": ["*mlp*"],
+                    "quantization_formats": [mtq.INT8_DEFAULT_CFG],
+                    "allow_no_quant": False,
+                }
+            ],
+            data_loader=[resumed_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,
+        )
+
+
 @pytest.mark.parametrize("method", ["gradient", "kl_div"])
 def test_get_auto_quantize_config(method):
     model = TransformerBlock()

From 4a250925c2abc72ced14db615ca0dc7d4a1be1fd Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Thu, 9 Jul 2026 16:54:56 -0700
Subject: [PATCH 4/5] Document module-specific AutoQuant search spaces

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 CHANGELOG.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 48aed5eb8b9..2d2d21e0fc5 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -42,6 +42,7 @@ Changelog
 - Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet.
 - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline.
 - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface/<model>/auto_quantize/``.
+- Add module-specific AutoQuantize search spaces through ``mtq.auto_quantize(..., module_search_spaces=...)`` and recipe-level ``auto_quantize.module_search_spaces``. Glob-matched runtime decision groups can override the global candidate formats and control whether BF16/no-quant is solver-selectable with ``allow_no_quant``; one candidate with ``allow_no_quant: false`` fixes the matching group to that format while retaining its effective-bits cost. Rules cannot partially split runtime-fused groups, fixed groups are isolated from unrelated calibration algorithms, and checkpoint replay validates the complete module-specific candidate configuration before reusing calibration or sensitivity state.
 - Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer.
 
 **Bug Fixes**

From 058601ea19615748999e6f8d083fba53eb4d5b5b Mon Sep 17 00:00:00 2001
From: weimingc <17592131+meenchen@users.noreply.github.com>
Date: Tue, 14 Jul 2026 14:38:17 -0700
Subject: [PATCH 5/5] Validate AutoQuant search checkpoint inputs

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
---
 modelopt/torch/quantization/algorithms.py     |  26 +++-
 modelopt/torch/quantization/model_quant.py    |  19 ++-
 .../unit/torch/quantization/test_autoquant.py | 125 +++++++++++++++++-
 3 files changed, 160 insertions(+), 10 deletions(-)

diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py
index 50a7a1133d1..520774117ce 100644
--- a/modelopt/torch/quantization/algorithms.py
+++ b/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,
diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py
index 8747adaf4d2..e981c948837 100644
--- a/modelopt/torch/quantization/model_quant.py
+++ b/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):
diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py
index 06d8258cd37..629482da1b3 100644
--- a/tests/unit/torch/quantization/test_autoquant.py
+++ b/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()