[OMNIML-5477, OMNIML-5119] Add module-specific AutoQuant search spaces

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**

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.

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,

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",