Add AutoQuantize recipe support [OMNIML-5116]

CHANGELOG.rst

@@ -11,6 +11,8 @@ Changelog
 
 **Deprecations**
 
+- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``). The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags are **deprecated** but still work: they are converted into an ``AutoQuantizeConfig`` on the fly (emitting a ``DeprecationWarning``) and will be removed in a future release. Prefer a recipe under ``modelopt_recipes/general/auto_quantize/``. See ``examples/hf_ptq/README.md``.
+
 - Renamed ``examples/llm_ptq`` to ``examples/hf_ptq`` to reflect that it covers Hugging Face LLM **and** VLM PTQ. A relative symlink ``examples/llm_ptq`` -> ``hf_ptq`` keeps existing paths and commands working; it will be removed in a future release. Please update references to the new ``examples/hf_ptq`` path.
 - Consolidated ``examples/vlm_ptq`` into ``examples/hf_ptq``. Vision-language model PTQ now shares the ``hf_ptq.py`` entry point and ``scripts/huggingface_example.sh``; pass ``--vlm`` to run the TensorRT-LLM multimodal quickstart smoke test. The ``examples/vlm_ptq/scripts/huggingface_example.sh`` entry point is deprecated: it now prints a warning and forwards to the ``hf_ptq`` script with ``--vlm``, and will be removed in a future release. See `examples/hf_ptq/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/hf_ptq#vlm-quantization>`__.
 - Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed.
@@ -39,6 +41,7 @@ Changelog
   - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change.
 - 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/``.
 
 **Bug Fixes**
 

examples/hf_ptq/README.md

@@ -198,7 +198,7 @@ python hf_ptq.py \
 
 Built-in recipes are located in `modelopt_recipes/general/ptq/` for model-agnostic recipes and in `modelopt_recipes/huggingface/<model_type>/ptq/` for recipes tuned to a specific Hugging Face `model_type` (see [`modelopt_recipes/huggingface/README.md`](../../modelopt_recipes/huggingface/README.md)). You can also provide a path to your own custom YAML recipe file or directory. See the [recipe documentation](https://nvidia.github.io/Model-Optimizer) for details on the YAML schema and available recipes.
 
-> *When `--recipe` is specified, `--qformat` and `--kv_cache_qformat` are ignored. The recipe fully defines the quantization configuration.*
+> *When `--recipe` is specified, `--qformat` is ignored. KV cache handling depends on the recipe type: a **PTQ** recipe bakes KV cache into its config and ignores `--kv_cache_qformat`; an **AutoQuantize** recipe falls back to `--kv_cache_qformat` unless it sets an explicit `kv_cache` field.*
 
 #### KV Cache Quantization
 
@@ -250,7 +250,7 @@ python hf_ptq.py \
 
 The cast pins each NVFP4 block's `scale_2 = 2^(k_max - 8)` and `_amax = 6 * 2^k_j`, both derived from the source MXFP4 E8M0 scales. For blocks whose `k_j` lands in E4M3's representable window (`k_max - k_j ≤ 17`), NVFP4 dequant matches MXFP4 dequant bit-for-bit; out-of-range blocks fall back to a data-derived per-block amax.
 
-> *`--cast_mxfp4_to_nvfp4` requires an NVFP4-family `--qformat` (e.g. `nvfp4_mlp_only`, `nvfp4_experts_only`, `nvfp4`) and is incompatible with `--auto_quantize_bits`.*
+> *`--cast_mxfp4_to_nvfp4` requires an NVFP4-family `--qformat` (e.g. `nvfp4_mlp_only`, `nvfp4_experts_only`, `nvfp4`) and is incompatible with AutoQuantize recipes (multi-format search).*
 
 #### Deepseek R1
 
@@ -305,12 +305,12 @@ Megatron-LM framework PTQ and TensorRT-LLM deployment examples are maintained in
 
 [AutoQuantize (`mtq.auto_quantize`)](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) is a PTQ algorithm which quantizes a model by searching for the best quantization format per-layer while meeting performance constraints specified by the user. `AutoQuantize` streamlines the trade-off of model accuracy and performance.
 
-Currently `AutoQuantize` supports only `auto_quantize_bits` as the performance constraint (for both weight-only
-quantization and weight & activation quantization). `auto_quantize_bits` constraint specifies the effective number of bits for the quantized model.
+`AutoQuantize` uses an effective-bits target (`effective_bits`) as the performance constraint (for both
+weight-only and weight & activation quantization) — the effective number of bits for the quantized model.
 
-You may specify an `auto_quantize_bits` constraint such as 4.8 for mixed precision quantization using `NVFP4_DEFAULT_CFG` & `FP8_DEFAULT_CFG`.
+You may specify an `effective_bits` target such as 5.4 for mixed precision quantization using `NVFP4_DEFAULT_CFG` & `FP8_DEFAULT_CFG`.
 `AutoQuantize` will automatically quantize highly sensitive layers in `FP8_DEFAULT_CFG` while keeping less sensitive layers in `NVFP4_DEFAULT_CFG` (and even skip quantization for any extremely sensitive layers) so that
-the the final mixed precision quantized model has an effective quantized bits of 4.8. This model would give a better accuracy than the model quantized with vanilla `NVFP4_DEFAULT_CFG` configuration since the more aggressive `NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers.
+the the final mixed precision quantized model has an effective quantized bits of 5.4. This model would give a better accuracy than the model quantized with vanilla `NVFP4_DEFAULT_CFG` configuration since the more aggressive `NVFP4_DEFAULT_CFG` quantization was not applied for the highly sensitive layers.
 
 Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize](https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.quantization.model_quant.html#modelopt.torch.quantization.model_quant.auto_quantize) API for more details):
 
@@ -337,7 +337,7 @@ Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize
     # Perform AutoQuantize
     model, search_state_dict = mtq.auto_quantize(
         model,
-        constraints = {"auto_quantize_bits": 4.8},
+        constraints = {"effective_bits": 5.4},
         # supported quantization formats are listed in `modelopt.torch.quantization.config.choices`
         quantization_formats = ["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"]
         data_loader = calib_dataloader,
@@ -351,31 +351,56 @@ 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
+[`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/`.
+
+> *Migration: prefer an AutoQuantize `--recipe`. The `--auto_quantize_bits`, `--auto_quantize_method`,
+> `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and `--auto_quantize_active_moe_expert_ratio`
+> CLI flags are **deprecated but still work** — they are converted into an `AutoQuantizeConfig` on the fly
+> (with a `DeprecationWarning`) and will be removed in a future release. They map to recipe fields:
+> `--auto_quantize_bits` → `constraints.effective_bits`, `--auto_quantize_method` → `auto_quantize_method`,
+> `--auto_quantize_score_size` → `score_size`, `--auto_quantize_cost_model` → `constraints.cost_model`,
+> `--auto_quantize_active_moe_expert_ratio` → `constraints.cost.active_moe_expert_ratio`, and the
+> `--qformat fp8,nvfp4` candidate list → `candidate_formats`. When converted, the shared base
+> `disabled_layers` and `cost_excluded_layers` patterns are appended automatically. `--auto_quantize_checkpoint`
+> is unchanged. Start from a shipped recipe under `modelopt_recipes/general/auto_quantize/`.*
+
 [Script](./scripts/huggingface_example.sh)
 
 ```bash
-export HF_PATH=<the downloaded LLaMA checkpoint from the Hugging Face hub, or simply the model card>
-# --auto_quantize_bits specifies the constraint for `AutoQuantize`
-# --quant specifies the formats to be searched for `AutoQuantize`
-# NOTE: auto_quantize_bits cannot be lower than the number of bits for the smallest quantization format in --quant
-scripts/huggingface_example.sh --model $HF_PATH --quant nvfp4_mse,fp8 --auto_quantize_bits 4.75 --calib_batch_size 4
+export HF_PATH=<the downloaded checkpoint from the Hugging Face hub, or simply the model card>
+# --recipe selects an AutoQuantize recipe; the recipe defines the candidate formats and the
+# effective-bits target (here NVFP4 + FP8 at 5.4 effective bits).
+scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits --calib_batch_size 4
 ```
 
-The above example perform `AutoQuantize` where the less quantization accuracy sensitive layers are quantized with `nvfp4_mse` (specified by `--quant nvfp4_mse`) and the more sensitive layers
-are kept un-quantized such that the effective bits is 4.75 (specified by `--auto_quantize_bits 4.75`).
+The recipe quantizes the less accuracy-sensitive layers with the more aggressive format (e.g. NVFP4) and
+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.
+
+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.
 
-#### AutoQuantize Advanced Options
+KV cache is applied as a uniform post-step, not part of the per-layer search. An AutoQuantize recipe
+falls back to `--kv_cache_qformat` (default `fp8_cast`) unless it sets an explicit `kv_cache` field.
 
-| Flag | Default | Description |
-| :--- | :---: | :--- |
-| `--auto_quantize_method` | `gradient` | Sensitivity analysis method. `gradient` uses gradient-based scoring (requires labels). `kl_div` uses KL divergence between original and quantized outputs (no labels required). |
-| `--auto_quantize_score_size` | `128` | Number of samples for sensitivity scoring. Reducing this speeds up the search while only minimally affecting accuracy (compared to reducing `--calib_size`). |
-| `--auto_quantize_checkpoint` | auto-generated | Path to save/restore search state (sensitivity scores, costs). Useful for resuming interrupted searches. |
+The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an
+interrupted search (skips re-scoring):
 
 ```bash
-# Use KL divergence method with smaller scoring set for faster search
-scripts/huggingface_example.sh --model $HF_PATH --quant nvfp4_mse,fp8 \
-  --auto_quantize_bits 4.75 --auto_quantize_method kl_div --auto_quantize_score_size 64
+scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits \
+  --auto_quantize_checkpoint /path/to/auto_quantize.pth --calib_batch_size 4
 ```
 
 The example scripts above also have an additional flag `--tasks`, where the actual tasks run in the script can be customized. The allowed tasks are `quant,mmlu,lm_eval,livecodebench,simple_eval` specified in the script [parser](./scripts/parser.sh). The tasks combo can be specified with a comma-separated task list. Some tasks like mmlu can take a long time to run. To run lm_eval tasks, please also specify the `--lm_eval_tasks` flag with comma separated lm_eval tasks [here](https://github.com/EleutherAI/lm-evaluation-harness/tree/main/lm_eval/tasks).

examples/hf_ptq/example_utils.py

@@ -42,7 +42,6 @@
 )
 
 from modelopt.torch.export.model_utils import is_multimodal_model
-from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg
 
 try:
     from huggingface_hub import snapshot_download
@@ -53,54 +52,6 @@
 
 SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"]
 
-# TODO: Refactor into the config system.
-_QWEN36_AUTOQ_DISABLED_LAYERS = ("*shared_expert_gate*",)
-_VLM_AUTOQ_DISABLED_LAYERS = ("*visual*", "*mtp*", "*vision_tower*")
-
-
-def _is_qwen_model(model) -> bool:
-    """Return True when model/config identifiers indicate a Qwen-family model."""
-    candidates = [type(model).__name__]
-    config = getattr(model, "config", None)
-    configs = [
-        config,
-        getattr(config, "text_config", None),
-        getattr(config, "language_config", None),
-    ]
-    for cfg in configs:
-        if cfg is None:
-            continue
-        candidates.append(type(cfg).__name__)
-        model_type = getattr(cfg, "model_type", None)
-        if model_type is not None:
-            candidates.append(str(model_type))
-        architectures = getattr(cfg, "architectures", ()) or ()
-        if isinstance(architectures, str):
-            architectures = (architectures,)
-        candidates.extend(str(architecture) for architecture in architectures)
-    return any("qwen" in candidate.lower() for candidate in candidates)
-
-
-def _get_auto_quantize_disabled_layers(model) -> list[str]:
-    """Return layer patterns that should be excluded from AutoQuantize search."""
-    disabled_layers = [
-        entry["quantizer_name"]
-        for entry in _default_disabled_quantizer_cfg
-        if "parent_class" not in entry and entry["quantizer_name"] != "*lm_head*"
-    ]
-    if _is_qwen_model(model):
-        disabled_layers.extend(p for p in _QWEN36_AUTOQ_DISABLED_LAYERS if p not in disabled_layers)
-    if is_multimodal_model(model):
-        disabled_layers.extend(p for p in _VLM_AUTOQ_DISABLED_LAYERS if p not in disabled_layers)
-    return disabled_layers
-
-
-def _get_auto_quantize_cost_excluded_patterns(model) -> list[str]:
-    """Return layer patterns excluded only from AutoQuantize cost accounting."""
-    if is_multimodal_model(model):
-        return list(_VLM_AUTOQ_DISABLED_LAYERS)
-    return []
-
 
 def run_nemotron_vl_preview(
     full_model,