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

CHANGELOG.rst

@@ -52,6 +52,7 @@ Experimental
 - 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 global candidate formats and control whether BF16/no-quant is solver-selectable with ``allow_no_quant``. Recipes can instead reuse a normal PTQ ``quantize`` config as the fixed baseline and explicitly list only genuinely searched modules; fixed and searched groups remain in one calibration, scoring, effective-bits, checkpoint, and export flow. Rules cannot partially split runtime-fused groups, fixed groups are isolated from unrelated calibration algorithms, infeasible resolved budgets fail before calibration, and checkpoint replay validates the fixed baseline, resolved groups, candidate choices, scoring boundaries, and cost weights 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.
 - Add a ``constant_amax`` ``QuantizerAttributeConfig`` field that pins a quantizer's ``amax`` to a fixed value and skips activation calibration (no forward statistics collected). Unlike ``use_constant_amax`` (which hardcodes 448.0 for KV-cache cast math and registers no buffer), ``constant_amax`` stores the constant on the ``_amax`` buffer so it is used by both the fake-quant forward and the exported scaling factor. For NVFP4 activations the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX)``, so ``constant_amax: 2688.0`` yields ``input_scale == 1.0``. Ships a new recipe ``modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml`` that applies this to the MoE expert activation quantizers.
 - Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization (and none has a static ``bias_calibrator``). Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it.

examples/hf_ptq/README.md

@@ -352,8 +352,9 @@ Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize
 `AutoQuantize` can be performed for Huggingface LLM models like [Qwen](https://huggingface.co/Qwen/Qwen3-8B) / [Nemotron](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) as shown below:
 
 `AutoQuantize` is driven by an **AutoQuantize recipe** passed with `--recipe`. The recipe defines the
-candidate formats, the `effective_bits` target, cost model, scoring method, search-disabled layers, and
-cost-excluded layers — see [`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in
+candidate formats, optional fixed PTQ baseline, `effective_bits` target, cost model, scoring method,
+search-disabled layers, and cost-excluded layers — see
+[`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in
 [`modelopt_recipes/general/auto_quantize/`](../../modelopt_recipes/general/auto_quantize); model-specific
 recipes (carrying architecture-specific disabled layers — e.g. VL vision towers) live under
 `modelopt_recipes/huggingface/<model>/auto_quantize/`.
@@ -382,12 +383,54 @@ 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`).
+`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`).
+
+AutoQuantize recipes support two mutually exclusive search-space styles:
+
+1. Set top-level `auto_quantize.candidate_formats` to search every unmatched quantizable module, with
+   optional `module_search_spaces` overrides.
+2. Set a normal top-level `quantize` config as the fixed PTQ baseline, omit top-level
+   `candidate_formats`, and use `module_search_spaces` to list only the modules AutoQuantize should
+   search. The fixed and searched modules still run through one integrated calibration, scoring, cost,
+   and export flow.
+
+For example, this keeps unmatched modules at the normal W4A16 NVFP4 PTQ setting while searching only
+attention between W4A16 and FP8:
+
+```yaml
+imports:
+  w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4
+  fp8: configs/ptq/presets/model/fp8
+
+quantize:
+  $import: w4a16_nvfp4
+
+auto_quantize:
+  constraints:
+    effective_bits: 6.0
+  module_search_spaces:
+    - module_name_patterns: ["*self_attn*", "*linear_attn*"]
+      candidate_formats:
+        - $import: w4a16_nvfp4
+        - $import: fp8
+      allow_no_quant: false
+```
 
-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.
+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. Use the
+top-level `quantize` baseline, rather than a one-candidate search rule, for modules that are fixed and
+not actually searched.
+
+The fixed baseline may also reuse a model-specific PTQ configuration. For example, the Qwen3.6 MoE
+AutoQuantize recipe imports the same model-specific `quant_cfg` used by
+`huggingface/qwen3_5_moe/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast`, reproduces that recipe's `quantize`
+section, and lists only shared experts, attention, and `lm_head` under `module_search_spaces`. A
+loader test asserts that the inherited fixed baseline remains equal to the original PTQ recipe while
+leaving the original recipe unchanged.
 
 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,12 +303,34 @@ def _match_candidate_to_preset(fmt) -> tuple[str | None, dict]:
     return None, fmt.model_dump()
 
 
-def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) -> dict:
+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, fixed_quantize_config=None
+) -> dict:
     """Map a resolved AutoQuantizeConfig to mtq.auto_quantize inputs.
 
-    Single, testable place where a recipe maps to mtq inputs. ``disabled_layers`` and candidate
-    cost come entirely from the recipe (no model introspection). KV cache falls back to
-    ``--kv_cache_qformat`` when the recipe omits it.
+    Single, testable place where a recipe maps to mtq inputs. ``fixed_quantize_config`` is the
+    optional normal PTQ baseline for modules outside explicit search spaces. ``disabled_layers``
+    and candidate cost come entirely from the recipe (no model introspection). KV cache falls back
+    to ``--kv_cache_qformat`` when the recipe omits it.
     """
     constraints = aq_config.constraints.model_dump(exclude_none=True)
     # cost_excluded_layers (sibling of disabled_layers) maps to the mtq cost key: these layers are
@@ -327,23 +349,25 @@ 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)
+    fixed_quantization_config = (
+        _mtq_candidate_formats([fixed_quantize_config])[0]
+        if fixed_quantize_config is not None
+        else None
+    )
+    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,
+        "fixed_quantization_config": fixed_quantization_config,
+        "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,
@@ -398,11 +422,12 @@ def auto_quantize(
     calib_dataloader: DataLoader,
     aq_config,
     full_model: torch.nn.Module | None = None,
+    fixed_quantize_config=None,
 ):
     """Recipe-driven auto_quantize, organized around an AutoQuantizeConfig.
 
-    The sole AutoQuantize entry point: it is driven entirely by the recipe's AutoQuantizeConfig
-    (candidate formats, constraints, disabled/cost-excluded layers) and wraps ``mtq.auto_quantize``.
+    The sole AutoQuantize entry point: it is driven by the recipe's AutoQuantizeConfig and optional
+    fixed PTQ config, then wraps ``mtq.auto_quantize``.
     """
     if args.calib_with_images:
         raise NotImplementedError(
@@ -413,7 +438,9 @@ def auto_quantize(
         "Auto Quantization is not supported for pipeline parallel size > 1"
     )
 
-    inputs = _mtq_inputs_from_auto_quantize_config(aq_config, args)
+    inputs = _mtq_inputs_from_auto_quantize_config(
+        aq_config, args, fixed_quantize_config=fixed_quantize_config
+    )
 
     # base-model lm_head handling (mirrors the CLI helper)
     is_base_model = (
@@ -467,6 +494,8 @@ def forward_step(model, batch):
         forward_step=forward_step,
         loss_func=loss_func,
         quantization_formats=inputs["quantization_formats"],
+        fixed_quantization_config=inputs["fixed_quantization_config"],
+        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,
@@ -1082,19 +1111,24 @@ def quantize_main(
     # --auto_quantize_* CLI flags converted on the fly. Everything downstream is recipe-driven.
     if isinstance(recipe, ModelOptAutoQuantizeRecipe):
         aq_config = recipe.auto_quantize
+        fixed_quantize_config = recipe.quantize
     elif args.recipe is None and args.auto_quantize_bits is not None:
         warnings.warn(
             "The --auto_quantize_* CLI flags are deprecated; use an AutoQuantize --recipe instead. "
             "They are converted to an AutoQuantizeConfig on the fly for now.",
             DeprecationWarning,
         )
         aq_config = _auto_quantize_config_from_cli(args)
+        fixed_quantize_config = None
     else:
         aq_config = None
+        fixed_quantize_config = None
 
     def _is_layerwise(obj):
         if isinstance(obj, ModelOptPTQRecipe):
             return _is_layerwise(obj.quantize.algorithm)
+        if isinstance(obj, ModelOptAutoQuantizeRecipe):
+            return obj.quantize is not None and _is_layerwise(obj.quantize.algorithm)
         if isinstance(obj, list):
             return any(_is_layerwise(a) for a in obj)
         layerwise = getattr(obj, "layerwise", None)
@@ -1183,6 +1217,7 @@ def _is_layerwise(obj):
             calib_dataloader,
             aq_config,
             full_model=full_model,
+            fixed_quantize_config=fixed_quantize_config,
         )
 
     else: