diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5002918175d..030c3d41fd8 100755 --- a/CHANGELOG.rst +++ b/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 `__. - 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//auto_quantize/``. **Bug Fixes** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 3e85142a5df..3a7380a9f9e 100755 --- a/examples/hf_ptq/README.md +++ b/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//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//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= -# --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= +# --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). diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 9e8dea5f107..73ccc991b00 100755 --- a/examples/hf_ptq/example_utils.py +++ b/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, diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 959316233fb..8dcc78afa27 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -27,8 +27,6 @@ from cast_mxfp4_to_nvfp4 import apply_to_model as apply_cast_mxfp4_to_nvfp4 from cast_mxfp4_to_nvfp4 import force_weight_quantizers_static from example_utils import ( - _get_auto_quantize_cost_excluded_patterns, - _get_auto_quantize_disabled_layers, _resolve_model_path, build_quant_cfg, copy_custom_model_files, @@ -58,13 +56,8 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq import modelopt.torch.sparsity as mts -from modelopt.recipe import ModelOptPTQRecipe, load_recipe -from modelopt.recipe.presets import ( - KV_CACHE_NONE, - KV_QUANT_CFG_CHOICES, - QFORMAT_ALIASES, - QUANT_CFG_CHOICES, -) +from modelopt.recipe import ModelOptAutoQuantizeRecipe, ModelOptPTQRecipe, load_recipe +from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES from modelopt.torch.export import ( export_hf_checkpoint, export_hf_vllm_fq_checkpoint, @@ -75,7 +68,6 @@ save_expert_token_count_table, ) from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model -from modelopt.torch.quantization._auto_quantize_cost import EXCLUDED_MODULE_NAME_PATTERNS_KEY from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils import is_quantized @@ -114,51 +106,6 @@ def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool: return False -# Formats supported by mtq.auto_quantize unified-checkpoint export. -# -# This stays hardcoded — and intentionally not derived from the preset directory — -# because auto_quantize compatibility is a property of the export path (the unified -# HF checkpoint writer, TRT-LLM consumer constraints, layer-wise mixing rules), not -# of the YAML itself. A preset can exist and be valid for plain PTQ while not being -# safe to mix into an auto_quantize search. Update this set when adding/removing a -# format from auto_quantize support. -# -# NOTE: auto_quantize is being refactored/reimplemented; this table and the -# _canonical_qformat helper below are expected to be removed in the near future, so -# deliberately not invested in deriving them from the presets. -_AUTO_QUANTIZE_QFORMATS: frozenset[str] = frozenset( - { - "fp8", - "int8_smoothquant", - "int8_weight_only", - "int4_awq", - "nvfp4", - "nvfp4_awq_lite", - "nvfp4_w4a4_weight_mse_fp8_sweep", - "w4a8_awq_beta", - "w4a16_nvfp4", - "fp8_2d_blockwise_weight_only", - "w4a8_mxfp4_fp8", - "nvfp4_mlp_only", - "nvfp4_experts_only", - "nvfp4_omlp_only", - "nvfp4_w4a4_weight_local_hessian", - "mxfp8", - } -) - - -def _canonical_qformat(name: str) -> str: - """Resolve a user-provided qformat token to its canonical preset basename. - - Lets membership checks (e.g. against :data:`_AUTO_QUANTIZE_QFORMATS`) accept - either the short alias (``int8_sq``) or the canonical YAML basename - (``int8_smoothquant``). Unknown tokens pass through unchanged so the existing - error paths still fire. - """ - return QFORMAT_ALIASES.get(name, name) - - mto.enable_huggingface_checkpointing() @@ -230,6 +177,7 @@ def make_calib_dataloader( tokenizer: PreTrainedTokenizerBase | None, device: torch.device, model_type: str | None, + autoquant_gradient_recipe: bool = False, ) -> tuple[DataLoader | _DeviceDataLoader, str | None]: calib_dataloader = None first_text_speech_dataset = None @@ -293,9 +241,7 @@ def make_calib_dataloader( tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast) ), "The PreTrainedTokenizer must be set" # Labels are only needed for gradient-based auto_quantize - include_labels = ( - args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient" - ) + include_labels = autoquant_gradient_recipe calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, @@ -309,45 +255,173 @@ def make_calib_dataloader( return calib_dataloader, first_text_speech_dataset +# Presets safe to mix into an AutoQuantize search *and* write via the unified HF checkpoint +# exporter. Export-compatibility is a property of the export path, not of a preset's validity for +# plain PTQ, so this is a curated set rather than something derived from QUANT_CFG_CHOICES. +# TODO: drop the partial-model presets (e.g. nvfp4_mlp_only, nvfp4_experts_only) from this set as future work. +_AUTO_QUANTIZE_QFORMATS: frozenset[str] = frozenset( + { + "fp8", + "int8_smoothquant", + "int8_weight_only", + "int4_awq", + "nvfp4", + "nvfp4_awq_lite", + "nvfp4_w4a4_weight_mse_fp8_sweep", + "w4a8_awq_beta", + "w4a16_nvfp4", + "fp8_2d_blockwise_weight_only", + "w4a8_mxfp4_fp8", + "nvfp4_mlp_only", + "nvfp4_experts_only", + "nvfp4_omlp_only", + "nvfp4_w4a4_weight_local_hessian", + "mxfp8", + } +) + + +def _match_candidate_to_preset(fmt) -> tuple[str | None, dict]: + """Match a recipe candidate against the shipped QUANT_CFG_CHOICES presets by value. + + Returns ``(preset_name, quant_cfg)``: ``preset_name`` is the matched preset (or None for a + custom config matching none), and ``quant_cfg`` is the dict passed to mtq.auto_quantize. + Passing the matched preset dict (rather than the candidate's own dump) keeps the search naming + the candidate after the preset (e.g. FP8_DEFAULT_CFG), consistent with CLI-produced checkpoints. + + ``effective_bits`` is cost-only metadata (it does not affect export), so it is excluded when + identifying the preset — otherwise a per-candidate override would make a shipped preset look + "custom" and slip past the export-compat whitelist. Any override is preserved in the return. + """ + stripped = fmt.model_dump(exclude_unset=True) + match_key = {k: v for k, v in stripped.items() if k != "effective_bits"} + for name, preset in QUANT_CFG_CHOICES.items(): + if preset == match_key: + if "effective_bits" in stripped: + return name, {**preset, "effective_bits": stripped["effective_bits"]} + return name, preset + return None, fmt.model_dump() + + +def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) -> 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. + """ + 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 + # kept out of the bit-budget denominator (cost_weight 0) — e.g. VL vision towers — distinct from + # disabled_layers, which removes them from the search. + if aq_config.cost_excluded_layers: + constraints.setdefault("cost", {})["excluded_module_name_patterns"] = ( + aq_config.cost_excluded_layers + ) + if aq_config.kv_cache is not None: + kv_cache_quant_cfg = aq_config.kv_cache.model_dump() + elif args.kv_cache_qformat == KV_CACHE_NONE: + kv_cache_quant_cfg = None + else: + kv_cache_quant_cfg = copy.deepcopy(KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]) + # 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) + return { + "constraints": constraints, + "quantization_formats": quantization_formats, + "disabled_layers": aq_config.disabled_layers, + "kv_cache_quant_cfg": kv_cache_quant_cfg, + "method": aq_config.auto_quantize_method, + "score_size": aq_config.score_size, + } + + +def _auto_quantize_config_from_cli(args: argparse.Namespace): + """Convert the deprecated ``--auto_quantize_*`` flags into an AutoQuantizeConfig on the fly. + + Backward-compat shim: old CLI invocations are turned into the same config object the recipe + path consumes, so the rest of the flow is recipe-driven. Layer patterns come from the shared + base sets loaded once in modelopt.recipe.config (no model introspection, no new CLI flags): the + base disabled set, and the base cost-excluded set — the latter is appended unconditionally + because it is harmless on non-VL models (nothing matches → cost_weight 0 is a no-op) and correct + on VL models. + """ + from modelopt.recipe.config import ( + AUTOQUANT_BASE_COST_EXCLUDED_LAYERS, + AUTOQUANT_BASE_DISABLED_LAYERS, + AutoQuantizeConfig, + AutoQuantizeConstraints, + AutoQuantizeCost, + ) + from modelopt.torch.quantization.config import QuantizeConfig + + disabled_layers = list(AUTOQUANT_BASE_DISABLED_LAYERS) + cost_excluded_layers = list(AUTOQUANT_BASE_COST_EXCLUDED_LAYERS) + + cost = ( + AutoQuantizeCost(active_moe_expert_ratio=args.auto_quantize_active_moe_expert_ratio) + if args.auto_quantize_cost_model == "active_moe" + else None + ) + return AutoQuantizeConfig( + constraints=AutoQuantizeConstraints( + effective_bits=args.auto_quantize_bits, + cost_model=args.auto_quantize_cost_model, + cost=cost, + ), + candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES[q]) for q in args.qformat.split(",")], + auto_quantize_method=args.auto_quantize_method, + score_size=args.auto_quantize_score_size, + disabled_layers=disabled_layers, + cost_excluded_layers=cost_excluded_layers, + ) + + def auto_quantize( args: argparse.Namespace, language_model: torch.nn.Module, calib_dataloader: DataLoader, - auto_quantize_method="gradient", - auto_quantize_score_size=128, - auto_quantize_checkpoint=None, + aq_config, full_model: torch.nn.Module | None = None, ): - """Auto search quantization of multiple formats.""" + """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``. + """ if args.calib_with_images: raise NotImplementedError( "AutoQuantize with image-text calibration is not supported yet. " "Please run plain PTQ (e.g., --qformat nvfp4) with --calib_with_images." ) - - assert not (args.auto_quantize_bits and args.inference_pipeline_parallel > 1), ( + assert args.inference_pipeline_parallel <= 1, ( "Auto Quantization is not supported for pipeline parallel size > 1" ) - qformat_list = args.qformat.split(",") - assert qformat_list, "No quantization formats provided" - # Check if all provided quantization formats are supported. Canonicalize first so - # callers may pass either the short alias (``int8_sq``) or the canonical YAML - # basename (``int8_smoothquant``). - assert all( - _canonical_qformat(qformat) in _AUTO_QUANTIZE_QFORMATS for qformat in qformat_list - ), "One or more quantization formats provided are not supported for unified checkpoint export" - - # When language_model is a base text model without lm_head (e.g. Gemma4TextModel), - # use full_model's lm_head to compute logits/loss from hidden states. + inputs = _mtq_inputs_from_auto_quantize_config(aq_config, args) + + # base-model lm_head handling (mirrors the CLI helper) is_base_model = ( full_model is not None and language_model is not full_model and not hasattr(language_model, "lm_head") and hasattr(full_model, "lm_head") ) - if is_base_model: assert full_model is not None lm_head = full_model.lm_head @@ -360,23 +434,22 @@ def loss_func(output, data): return torch.nn.functional.cross_entropy( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) ) - else: def loss_func(output, data): return output.loss - if auto_quantize_method == "gradient": + if inputs["method"] == "gradient": def forward_step(model, batch): - inputs = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch - return model(**inputs) + inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch + return model(**inputs_) - elif auto_quantize_method == "kl_div": + elif inputs["method"] == "kl_div": def forward_step(model, batch): - inputs = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch - output = model(**inputs) + inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch + output = model(**inputs_) if is_base_model: assert full_model is not None return full_model.lm_head(output.last_hidden_state) @@ -384,62 +457,47 @@ def forward_step(model, batch): else: raise ValueError( - f"Invalid auto_quantize_method: {auto_quantize_method}. Must be 'gradient' or 'kl_div'" + f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'" ) - auto_quantize_constraints = { - "effective_bits": args.auto_quantize_bits, - "cost_model": args.auto_quantize_cost_model, - } - auto_quantize_cost = {} - if args.auto_quantize_active_moe_expert_ratio is not None: - auto_quantize_cost["active_moe_expert_ratio"] = args.auto_quantize_active_moe_expert_ratio - cost_excluded_patterns = _get_auto_quantize_cost_excluded_patterns(language_model) - if cost_excluded_patterns: - auto_quantize_cost[EXCLUDED_MODULE_NAME_PATTERNS_KEY] = cost_excluded_patterns - if auto_quantize_cost: - auto_quantize_constraints["cost"] = auto_quantize_cost - language_model, _ = mtq.auto_quantize( language_model, - constraints=auto_quantize_constraints, + constraints=inputs["constraints"], data_loader=calib_dataloader, forward_step=forward_step, - loss_func=loss_func, # Only used for gradient-based method - # TRTLLM only support one quantization format or None (do not quantize, internally supported) - quantization_formats=[QUANT_CFG_CHOICES[format] for format in qformat_list], + loss_func=loss_func, + quantization_formats=inputs["quantization_formats"], num_calib_steps=len(calib_dataloader), - # AutoQuantize scoring is the costly phase; allow smaller sample counts than calibration. - num_score_steps=min( - len(calib_dataloader), max(auto_quantize_score_size // args.batch_size, 1) - ), + num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)), verbose=True, - disabled_layers=_get_auto_quantize_disabled_layers(language_model), - method=auto_quantize_method, - checkpoint=auto_quantize_checkpoint, + disabled_layers=inputs["disabled_layers"], + method=inputs["method"], + checkpoint=args.auto_quantize_checkpoint, ) + # KV cache quantization is uniform; applied after the LP search. + kv_cache_quant_cfg = inputs["kv_cache_quant_cfg"] calibrate_loop = create_forward_loop(dataloader=calib_dataloader) - # We need to explicitly set up KV cache quantization after auto_quantize - enable_quant_kv_cache = args.kv_cache_qformat != KV_CACHE_NONE - print(f"{'Enable' if enable_quant_kv_cache else 'Disable'} KV cache quantization") - if enable_quant_kv_cache: - kv_cache_quant_cfg = copy.deepcopy(KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]["quant_cfg"]) - kv_cache_quant_cfg = [ - e for e in kv_cache_quant_cfg if e["quantizer_name"] != "*" - ] # keep other quantizers from auto_quantize - - mtq.set_quantizer_by_cfg(language_model, quant_cfg=kv_cache_quant_cfg) - if not _kv_cfg_uses_constant_amax(kv_cache_quant_cfg): - # Calibrate only the KV cache quantizers; disable all others. + print(f"{'Enable' if kv_cache_quant_cfg is not None else 'Disable'} KV cache quantization") + if kv_cache_quant_cfg is not None: + kv_entries = [ + e for e in copy.deepcopy(kv_cache_quant_cfg["quant_cfg"]) if e["quantizer_name"] != "*" + ] + mtq.set_quantizer_by_cfg(language_model, quant_cfg=kv_entries) + if not _kv_cfg_uses_constant_amax(kv_entries): with mtq.set_quantizer_by_cfg_context( language_model, - [{"quantizer_name": "*", "enable": False}, *kv_cache_quant_cfg], + [{"quantizer_name": "*", "enable": False}, *kv_entries], ): mtq.calibrate(language_model, algorithm="max", forward_loop=calibrate_loop) return language_model +def _recipe_is_auto_quantize(recipe: str | None) -> bool: + """True if ``recipe`` resolves to an AutoQuantize recipe (peeked before model load).""" + return recipe is not None and isinstance(load_recipe(recipe), ModelOptAutoQuantizeRecipe) + + def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False @@ -489,8 +547,16 @@ def load_model(args: argparse.Namespace): is_nemotron_vl_model = is_nemotron_vl(full_model) - # Default to image-text calibration for VLM models - if is_nemotron_vl_model and not args.calib_with_images and args.auto_quantize_bits is None: + # Default to image-text calibration for VLM models. Skip for either AutoQuantize path (recipe or + # the deprecated --auto_quantize_bits CLI), whose text-only path does not support image-text + # calibration yet (auto_quantize() would raise); auto-enabling it here would make Nemotron-VL + # AutoQuantize fail unconditionally. + if ( + is_nemotron_vl_model + and not args.calib_with_images + and not _recipe_is_auto_quantize(args.recipe) + and args.auto_quantize_bits is None + ): print("Nemotron VL model detected. Enabling image-text calibration by default.") args.calib_with_images = True @@ -542,9 +608,10 @@ def load_model(args: argparse.Namespace): : len(args.dataset) ] - # Plain PTQ quantizes only the extracted language model. Recipe and - # AutoQuantize paths keep the outer CausalLM so recipes/search can see - # Qwen3.5/3.6-MoE VLM lm_head. + # Plain PTQ quantizes only the extracted language model. Recipe and AutoQuantize paths + # (incl. the deprecated --auto_quantize_bits CLI) keep the outer CausalLM so recipes / + # search can see the Qwen3.5/3.6-MoE VLM lm_head; extracting here would leave modelopt + # state on the ancestors and make auto_quantize() fail with "multiple modelopt states". if args.recipe is None and args.auto_quantize_bits is None: extracted_lm, extracted_model_type = extract_and_prepare_language_model_from_vl( full_model @@ -1002,14 +1069,29 @@ def quantize_main( ): # Load the recipe up front so we can detect layerwise calibration before batch-size probing. recipe = None - if args.recipe is not None and not args.auto_quantize_bits: + if args.recipe is not None: print(f"Use recipe {args.recipe} for quantization") recipe = load_recipe(args.recipe) - if not isinstance(recipe, ModelOptPTQRecipe): + if not isinstance(recipe, (ModelOptPTQRecipe, ModelOptAutoQuantizeRecipe)): raise TypeError( - f"Expected PTQ recipe, but got {type(recipe).__name__} from {args.recipe}" + f"Expected PTQ or AutoQuantize recipe, but got {type(recipe).__name__} " + f"from {args.recipe}" ) + # Resolve the AutoQuantizeConfig from either source: a recipe, or the deprecated + # --auto_quantize_* CLI flags converted on the fly. Everything downstream is recipe-driven. + if isinstance(recipe, ModelOptAutoQuantizeRecipe): + aq_config = recipe.auto_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) + else: + aq_config = None + def _is_layerwise(obj): if isinstance(obj, ModelOptPTQRecipe): return _is_layerwise(obj.quantize.algorithm) @@ -1059,7 +1141,7 @@ def _is_layerwise(obj): else: sample_input_single_batch = None - run_auto_quant = args.auto_quantize_bits is not None + run_auto_quant = aq_config is not None args.batch_size = get_max_batch_size( language_model, @@ -1073,7 +1155,15 @@ def _is_layerwise(obj): print(f"Use calib batch_size {args.batch_size}") calib_dataloader, first_text_speech_dataset = make_calib_dataloader( - args, language_model, processor, tokenizer, device, model_type + args, + language_model, + processor, + tokenizer, + device, + model_type, + autoquant_gradient_recipe=( + aq_config is not None and aq_config.auto_quantize_method == "gradient" + ), ) # Detect if this is a Nemotron VL model using architecture-based detection @@ -1083,25 +1173,15 @@ def _is_layerwise(obj): args, full_model, model_type, tokenizer, calib_dataloader, is_nemotron_vl_model ) - if args.auto_quantize_bits: - assert len(args.qformat.split(",")) > 1, ( - "Auto quantization needs multiple quantization format." - ) - - # For VL models, autoquant must walk submodules of the OUTER CausalLM - # (which carries lm_head and the LM-head forward path) — otherwise - # lm_head and any sibling-of-language_model modules are silently - # invisible to the search. ``forward_step`` also needs the outer model - # to produce ``CausalLMOutputWithPast`` (for ``.loss`` / ``.logits``). - # Visual tower and MTP siblings are auto-excluded inside - # ``auto_quantize()`` via *visual* / *mtp* / *vision_tower* patterns. + if aq_config is not None: + # AutoQuantize (recipe or the deprecated --auto_quantize_* CLI, converted on the fly). For + # VL models the search walks the OUTER CausalLM (which carries lm_head and the LM-head + # forward path); architecture-specific exclusions come from aq_config.disabled_layers. auto_quantize( args, full_model, calib_dataloader, - auto_quantize_method=args.auto_quantize_method, - auto_quantize_score_size=args.auto_quantize_score_size, - auto_quantize_checkpoint=args.auto_quantize_checkpoint, + aq_config, full_model=full_model, ) @@ -1217,9 +1297,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--recipe", help=( - "PTQ recipe YAML file or name without suffix (e.g. general/ptq/fp8_default-kv_fp8_cast, " - "general/ptq/nvfp4_default-kv_fp8_cast, general/ptq/nvfp4_default-kv_nvfp4_cast). " - "When set, --kv_cache_qformat is ignored; the recipe fully determines KV cache config." + "PTQ or AutoQuantize recipe YAML file or name without suffix (e.g. " + "general/ptq/nvfp4_default-kv_fp8_cast, general/auto_quantize/nvfp4_fp8_at_4p8bits). " + "KV cache source depends on the recipe type: PTQ recipes bake KV cache into quant_cfg " + "and --kv_cache_qformat is ignored; AutoQuantize recipes fall back to --kv_cache_qformat " + "unless the recipe sets an explicit kv_cache field." ), default=None, ) @@ -1227,10 +1309,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--device", default="cuda") parser.add_argument( "--qformat", - help=( - "Quantization format. If --auto_quantize_bits is set, this argument specifies the quantization " - "format for optimal per-layer auto_quantize search." - ), + help="Quantization format for single-format PTQ. For mixed-precision search, use an " + "AutoQuantize recipe via --recipe.", default="fp8", ) parser.add_argument( @@ -1290,15 +1370,6 @@ def parse_args() -> argparse.Namespace: default="dense", choices=["dense", "sparsegpt"], ) - parser.add_argument( - "--auto_quantize_bits", - default=None, - type=float, - help=( - "Effective bits constraint for auto_quantize. If not set, " - "regular quantization without auto_quantize search will be applied." - ), - ) parser.add_argument( "--kv_cache_qformat", required=False, @@ -1309,8 +1380,9 @@ def parse_args() -> argparse.Namespace: "Formats whose preset pins use_constant_amax on the KV bmm quantizer " "(e.g. fp8_cast, nvfp4_cast) set the amax to FP8 range without data-driven " "calibration; all other formats (fp8, nvfp4, ...) use data-driven calibration. " - "Ignored when --recipe is given: the recipe YAML is authoritative for KV " - "cache config (use the *_cast_kv.yaml recipes for the cast variants)." + "With --recipe, the source depends on the recipe type: a PTQ recipe is " + "authoritative for KV cache and ignores this flag; an AutoQuantize recipe " + "falls back to this flag unless it sets an explicit kv_cache field." ), ) parser.add_argument( @@ -1380,58 +1452,52 @@ def parse_args() -> argparse.Namespace: default=None, type=str, ) + parser.add_argument( + "--auto_quantize_checkpoint", + type=str, + default=None, + help=( + "Path to checkpoint file for saving/restoring auto_quantize search state " + "(sensitivity scores, costs, etc.). Used with an AutoQuantize --recipe or the " + "deprecated --auto_quantize_bits CLI path." + ), + ) + # Deprecated AutoQuantize CLI flags: kept as a backward-compat shim that converts them into an + # AutoQuantizeConfig on the fly (see _auto_quantize_config_from_cli). Prefer --recipe. The old + # CLI also lives on the 0.45 branch. + parser.add_argument( + "--auto_quantize_bits", + type=float, + default=None, + help="[Deprecated: use an AutoQuantize --recipe] Effective-bits target; also enables the " + "AutoQuantize CLI path. Candidate formats are taken from --qformat (comma-separated).", + ) parser.add_argument( "--auto_quantize_method", type=str, default="gradient", choices=["gradient", "kl_div"], - help=( - "Method for auto_quantize sensitivity analysis. 'gradient' uses gradient-based method " - "(requires labels in dataset). 'kl_div' uses KL divergence between original and " - "quantized model outputs (no labels required). Default: 'gradient'" - ), + help="[Deprecated: use an AutoQuantize --recipe] Sensitivity scoring method.", ) parser.add_argument( "--auto_quantize_score_size", type=int, default=128, - help=( - "Number of samples to use for auto_quantize scoring. Most of auto_quantize time is spent on " - "sensitivity score estimation, so reducing this speeds it up while only minimally affecting " - "final model accuracy compared to lowering --calib_size (the number of samples used for calibration)." - ), - ) - parser.add_argument( - "--auto_quantize_checkpoint", - type=str, - default=None, - help=( - "Path to checkpoint file for saving/restoring auto_quantize search state " - "(sensitivity scores, costs, etc.). Only used when auto_quantize_bits is specified." - ), + help="[Deprecated: use an AutoQuantize --recipe] Number of samples for sensitivity scoring.", ) parser.add_argument( "--auto_quantize_cost_model", type=str, default="weight", choices=["weight", "active_moe"], - help=( - "Cost model for auto_quantize effective-bits accounting. 'weight' counts all " - "quantizable weights equally. 'active_moe' scales routed MoE expert weights by " - "--auto_quantize_active_moe_expert_ratio, or infers top_k/num_experts from model config." - ), + help="[Deprecated: use an AutoQuantize --recipe] Cost model for the effective-bits search.", ) parser.add_argument( "--auto_quantize_active_moe_expert_ratio", type=float, default=None, - help=( - "Routed MoE expert active ratio for --auto_quantize_cost_model active_moe. " - "For top-k MoE this is top_k / num_experts. If omitted, common model config " - "fields such as num_experts_per_tok and num_experts are used when available. " - "This only affects AutoQuant cost accounting and does not change calibration " - "routing; use --moe_calib_experts_ratio to control calibration expert coverage." - ), + help="[Deprecated: use an AutoQuantize --recipe] Routed-expert active ratio for the " + "'active_moe' cost model.", ) parser.add_argument( "--moe_calib_experts_ratio", @@ -1466,20 +1532,6 @@ def parse_args() -> argparse.Namespace: args = parser.parse_args() if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0): parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].") - if args.auto_quantize_bits is not None and args.calib_with_images: - parser.error("--calib_with_images is not supported with --auto_quantize_bits.") - if args.auto_quantize_active_moe_expert_ratio is not None and not ( - 0.0 < args.auto_quantize_active_moe_expert_ratio <= 1.0 - ): - parser.error("--auto_quantize_active_moe_expert_ratio must be in the range (0.0, 1.0].") - if ( - args.auto_quantize_cost_model == "weight" - and args.auto_quantize_active_moe_expert_ratio is not None - ): - parser.error( - "--auto_quantize_active_moe_expert_ratio requires " - "--auto_quantize_cost_model active_moe." - ) if args.specdec_offline_dataset is not None and args.sparsity_fmt != "dense": parser.error("--specdec_offline_dataset is only supported with --sparsity_fmt dense (PTQ).") @@ -1491,10 +1543,10 @@ def parse_args() -> argparse.Namespace: # via init_quantized_weights(), so it cannot honor a --recipe (which is authoritative # for the quant layout in quantize_main). Reject the combination rather than silently # instrumenting a layout that diverges from the recipe. - if args.low_memory_mode and args.recipe is not None: + if args.low_memory_mode and (args.recipe is not None or args.auto_quantize_bits is not None): parser.error( - "--low_memory_mode does not yet support --recipe; the low-memory loader still " - "initializes quantizers from --qformat/--kv_cache_qformat." + "--low_memory_mode does not support --recipe or AutoQuantize (--auto_quantize_bits); " + "the low-memory loader initializes quantizers from --qformat/--kv_cache_qformat." ) return args @@ -1565,10 +1617,5 @@ def main(args: argparse.Namespace): "--cast_mxfp4_to_nvfp4 requires NVFP4-family --qformat values " f"(got {args.qformat!r}). Use e.g. --qformat nvfp4 or nvfp4_mlp_only." ) - if args.auto_quantize_bits is not None: - raise ValueError( - "--cast_mxfp4_to_nvfp4 is not supported with --auto_quantize_bits " - "(multi-format auto-quantize)." - ) main(args) diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index a073fb7fecb..84057e468c9 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -94,29 +94,27 @@ if [ "$LOW_MEMORY_MODE" = "true" ]; then PTQ_ARGS+=" --low_memory_mode " fi -if [ -n "$AUTO_QUANTIZE_BITS" ]; then - PTQ_ARGS+=" --auto_quantize_bits=$AUTO_QUANTIZE_BITS " -fi - -if [ -n "$AUTO_QUANTIZE_METHOD" ]; then - PTQ_ARGS+=" --auto_quantize_method=$AUTO_QUANTIZE_METHOD " -fi - -if [ -n "$AUTO_QUANTIZE_SCORE_SIZE" ]; then - PTQ_ARGS+=" --auto_quantize_score_size=$AUTO_QUANTIZE_SCORE_SIZE " -fi - -# Automatically generate auto_quantize checkpoint path if not provided -if [ -n "$AUTO_QUANTIZE_BITS" ] && [ -z "$AUTO_QUANTIZE_CHECKPOINT" ]; then - # Create a descriptive checkpoint name based on model and quantization settings - AQ_METHOD=${AUTO_QUANTIZE_METHOD:-gradient} - AUTO_QUANTIZE_CHECKPOINT="${ROOT_SAVE_PATH}/auto_quantize_checkpoints/${MODEL_NAME}_${AQ_METHOD}.pth" - mkdir -p $(dirname $AUTO_QUANTIZE_CHECKPOINT) +# AutoQuantize runs via an AutoQuantize --recipe or the deprecated --auto_quantize_bits CLI path. +# Auto-generate a checkpoint path (to save/restore the search state) when the user didn't supply one. +if [ -z "$AUTO_QUANTIZE_CHECKPOINT" ] && { [[ "$RECIPE" == *auto_quantize* ]] || [ -n "$AUTO_QUANTIZE_BITS" ]; }; then + AUTO_QUANTIZE_CHECKPOINT="${ROOT_SAVE_PATH}/auto_quantize_checkpoints/${MODEL_NAME}.pth" + mkdir -p "$(dirname "$AUTO_QUANTIZE_CHECKPOINT")" echo "Auto-generated auto_quantize checkpoint path: $AUTO_QUANTIZE_CHECKPOINT" fi +if [ -n "$AUTO_QUANTIZE_CHECKPOINT" ]; then + PTQ_ARGS+=" --auto_quantize_checkpoint=$AUTO_QUANTIZE_CHECKPOINT " +fi +# Deprecated AutoQuantize CLI flags: passed through to hf_ptq.py, which converts them into an +# AutoQuantizeConfig on the fly. Prefer an AutoQuantize --recipe. if [ -n "$AUTO_QUANTIZE_BITS" ]; then - PTQ_ARGS+=" --auto_quantize_checkpoint=$AUTO_QUANTIZE_CHECKPOINT " + PTQ_ARGS+=" --auto_quantize_bits=$AUTO_QUANTIZE_BITS " + PTQ_ARGS+=" --auto_quantize_method=${AUTO_QUANTIZE_METHOD:-gradient} " + PTQ_ARGS+=" --auto_quantize_score_size=${AUTO_QUANTIZE_SCORE_SIZE:-128} " + PTQ_ARGS+=" --auto_quantize_cost_model=${AUTO_QUANTIZE_COST_MODEL:-weight} " + if [ -n "$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO" ]; then + PTQ_ARGS+=" --auto_quantize_active_moe_expert_ratio=$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO " + fi fi if [ -n "$CALIB_DATASET" ]; then diff --git a/examples/hf_ptq/scripts/parser.sh b/examples/hf_ptq/scripts/parser.sh index 06b440e5731..03ed3a57631 100644 --- a/examples/hf_ptq/scripts/parser.sh +++ b/examples/hf_ptq/scripts/parser.sh @@ -41,7 +41,7 @@ parse_options() { CALIB_WITH_IMAGES=false # Parse command-line options - ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,auto_quantize_bits:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_checkpoint:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@") + ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@") eval set -- "$ARGS" while true; do @@ -56,7 +56,6 @@ parse_options() { --awq_block_size ) AWQ_BLOCK_SIZE="$2"; shift 2;; --calib ) CALIB_SIZE="$2"; shift 2;; --calib_batch_size ) CALIB_BATCH_SIZE="$2"; shift 2;; - --auto_quantize_bits ) AUTO_QUANTIZE_BITS="$2"; shift 2;; --output ) BUILD_MAX_OUTPUT_LEN="$2"; shift 2;; --batch ) BUILD_MAX_BATCH_SIZE="$2"; shift 2;; --tasks ) TASKS="$2"; shift 2;; @@ -73,9 +72,12 @@ parse_options() { --low_memory_mode ) LOW_MEMORY_MODE=true; shift;; --calib_dataset ) CALIB_DATASET="$2"; shift 2;; --calib_seq ) CALIB_SEQ="$2"; shift 2;; + --auto_quantize_checkpoint ) AUTO_QUANTIZE_CHECKPOINT="$2"; shift 2;; + --auto_quantize_bits ) AUTO_QUANTIZE_BITS="$2"; shift 2;; --auto_quantize_method ) AUTO_QUANTIZE_METHOD="$2"; shift 2;; --auto_quantize_score_size ) AUTO_QUANTIZE_SCORE_SIZE="$2"; shift 2;; - --auto_quantize_checkpoint ) AUTO_QUANTIZE_CHECKPOINT="$2"; shift 2;; + --auto_quantize_cost_model ) AUTO_QUANTIZE_COST_MODEL="$2"; shift 2;; + --auto_quantize_active_moe_expert_ratio ) AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO="$2"; shift 2;; --moe_calib_experts_ratio ) MOE_CALIB_EXPERTS_RATIO="$2"; shift 2;; --cast_mxfp4_to_nvfp4 ) CAST_MXFP4_TO_NVFP4=true; shift;; --vlm ) VLM=true; shift;; @@ -158,7 +160,6 @@ parse_options() { echo "awq_block_size: $AWQ_BLOCK_SIZE" echo "calib: $CALIB_SIZE" echo "calib_batch_size: $CALIB_BATCH_SIZE" - echo "auto_quantize_bits: $AUTO_QUANTIZE_BITS" echo "input: $BUILD_MAX_INPUT_LEN" echo "output: $BUILD_MAX_OUTPUT_LEN" echo "batch: $BUILD_MAX_BATCH_SIZE" @@ -175,9 +176,12 @@ parse_options() { echo "low_memory_mode: $LOW_MEMORY_MODE" echo "calib_dataset: $CALIB_DATASET" echo "calib_seq: $CALIB_SEQ" + echo "auto_quantize_checkpoint: $AUTO_QUANTIZE_CHECKPOINT" + echo "auto_quantize_bits: $AUTO_QUANTIZE_BITS" echo "auto_quantize_method: $AUTO_QUANTIZE_METHOD" echo "auto_quantize_score_size: $AUTO_QUANTIZE_SCORE_SIZE" - echo "auto_quantize_checkpoint: $AUTO_QUANTIZE_CHECKPOINT" + echo "auto_quantize_cost_model: $AUTO_QUANTIZE_COST_MODEL" + echo "auto_quantize_active_moe_expert_ratio: $AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO" echo "moe_calib_experts_ratio: $MOE_CALIB_EXPERTS_RATIO" echo "cast_mxfp4_to_nvfp4: $CAST_MXFP4_TO_NVFP4" echo "vlm: $VLM" diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index ea72efdc7c7..2cf5a2f0cfb 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -19,10 +19,12 @@ import warnings from enum import Enum +from typing import Literal -from pydantic import Field, model_validator +from pydantic import Field, field_validator, model_validator from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField +from modelopt.torch.opt.config_loader import load_config from modelopt.torch.quantization.config import QuantizeConfig # noqa: TC001 from modelopt.torch.speculative.config import DFlashConfig, EagleConfig, MedusaConfig from modelopt.torch.speculative.plugins.hf_training_args import DataArguments as SpecDataArgs @@ -33,6 +35,10 @@ __all__ = [ "RECIPE_TYPE_TO_CLASS", + "AutoQuantizeConfig", + "AutoQuantizeConstraints", + "AutoQuantizeCost", + "ModelOptAutoQuantizeRecipe", "ModelOptDFlashRecipe", "ModelOptEagleRecipe", "ModelOptMedusaRecipe", @@ -48,6 +54,7 @@ class RecipeType(str, Enum): """List of recipe types. See ``RECIPE_TYPE_TO_CLASS`` at the bottom for the schema mapping.""" PTQ = "ptq" + AUTO_QUANTIZE = "auto_quantize" SPECULATIVE_EAGLE = "speculative_eagle" SPECULATIVE_DFLASH = "speculative_dflash" SPECULATIVE_MEDUSA = "speculative_medusa" @@ -116,6 +123,143 @@ class ModelOptPTQRecipe(ModelOptRecipeBase): ) +# Named alias so a shared layer-pattern unit (e.g. configs/auto_quantize/units/base_disabled_layers) +# can declare ``modelopt-schema: modelopt.recipe.config.LayerPatternList`` and be spliced into a +# ``list[str]`` field — mirrors how base_disable_all is imported into a PTQ quant_cfg list. +LayerPatternList = list[str] + + +def _load_layer_pattern_list(config_path: str) -> list[str]: + """Load a ``list[str]`` layer-pattern unit (e.g. AutoQuantize base disabled/cost-excluded). + + Relies on the unit's ``modelopt-schema: ...LayerPatternList`` comment (like + _load_quantizer_cfg_dict_list) rather than an explicit ``list[str]`` schema_type. + """ + return list(load_config(config_path)) + + +# Base AutoQuantize layer-pattern sets, loaded once (used by the deprecated --auto_quantize_* CLI shim). +AUTOQUANT_BASE_DISABLED_LAYERS: list[str] = _load_layer_pattern_list( + "configs/auto_quantize/units/base_disabled_layers" +) +AUTOQUANT_BASE_COST_EXCLUDED_LAYERS: list[str] = _load_layer_pattern_list( + "configs/auto_quantize/units/base_cost_excluded_layers" +) + + +class AutoQuantizeCost(ModeloptBaseConfig): + """Cost-model parameters (the ``cost`` sub-dict of ``mtq.auto_quantize`` constraints).""" + + active_moe_expert_ratio: float | None = ModeloptField( + default=None, + title="Active MoE expert ratio", + description="Routed experts active per token, in (0, 1]. Used by the 'active_moe' cost model.", + ) + + @field_validator("active_moe_expert_ratio") + @classmethod + def _validate_active_moe_expert_ratio(cls, v: float | None) -> float | None: + if v is not None and not (0 < v <= 1): + raise ValueError(f"active_moe_expert_ratio must be in (0, 1], got {v}") + return v + + +class AutoQuantizeConstraints(ModeloptBaseConfig): + """LP search constraints + cost model; matches the ``mtq.auto_quantize`` constraints dict.""" + + effective_bits: float = ModeloptField( + default=4.8, + title="Effective bits per weight", + description="Average weight-storage bits target for the LP, in (0, 16].", + ) + cost_model: Literal["weight", "active_moe"] = ModeloptField( + default="weight", + title="Cost model", + description="'weight' counts all weights equally; 'active_moe' scales routed-expert weights.", + ) + cost: AutoQuantizeCost | None = ModeloptField( + default=None, + title="Cost-model parameters", + description="Extra cost-model parameters; omit for the 'weight' cost model.", + ) + + @field_validator("effective_bits") + @classmethod + def _validate_effective_bits(cls, v: float) -> float: + if not (0 < v <= 16): + raise ValueError(f"effective_bits must be in (0, 16], got {v}") + return v + + +class AutoQuantizeConfig(ModeloptBaseConfig): + """Schema for the ``auto_quantize`` block of an AutoQuantize recipe.""" + + constraints: AutoQuantizeConstraints = Field( + title="Search constraints + cost model", + description="LP budget and cost model.", + ) + candidate_formats: list[QuantizeConfig] = ModeloptField( + default=[], + title="Candidate quantization formats", + description="Per-layer search space; each entry is a full QuantizeConfig. At least 1 " + "required — bf16/no-quant is always an implicit additional choice, so a single format " + "(e.g. [fp8]) yields a {fp8, bf16} per-layer search.", + validate_default=True, + ) + auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField( + default="gradient", + title="Sensitivity scoring method", + description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).", + ) + score_size: int = ModeloptField( + default=128, + title="Scoring sample count", + description="Number of samples used for sensitivity scoring (divided by batch_size to get " + "the number of mtq scoring steps). Matches the former --auto_quantize_score_size.", + ) + disabled_layers: LayerPatternList = ModeloptField( + default=[], + title="Search-excluded layer patterns", + description="Glob patterns; matching layers are excluded from the search (kept full precision).", + ) + cost_excluded_layers: LayerPatternList = ModeloptField( + default=[], + title="Cost-excluded layer patterns", + description="Glob patterns excluded from the bit-budget accounting (cost_weight 0) — e.g. VL " + "vision towers. Distinct from disabled_layers: those are removed from the search; these still " + "get searched but don't count toward effective_bits. The two roles overlap but are independent.", + ) + kv_cache: QuantizeConfig | None = ModeloptField( + default=None, + title="KV cache config (optional)", + description="QuantizeConfig applied as a uniform post-step; falls back to " + "the --kv_cache_qformat CLI flag when omitted.", + ) + + @field_validator("candidate_formats") + @classmethod + def _at_least_one_candidate(cls, v: list[QuantizeConfig]) -> list[QuantizeConfig]: + # mtq.auto_quantize always adds an implicit bf16/no-quant choice per layer, so a single + # explicit format already gives a real {format, bf16} search; only an empty list is invalid. + if not v: + raise ValueError( + "auto_quantize requires at least 1 candidate_format (bf16/no-quant is always an " + "implicit additional choice). For uniform quantization, use a PTQ recipe instead." + ) + return v + + +class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase): + """Our config class for AutoQuantize recipes.""" + + metadata: RecipeMetadataConfig = _metadata_field(RecipeType.AUTO_QUANTIZE) + + auto_quantize: AutoQuantizeConfig = Field( + title="AutoQuantize config", + description="AutoQuantize search configuration. Required.", + ) + + class ModelOptSpeculativeRecipeBase(ModelOptRecipeBase): """Base class for speculative-decoding recipes. @@ -215,6 +359,7 @@ class ModelOptMedusaRecipe(ModelOptSpeculativeRecipeBase): # uses this for typed-list ``$import`` resolution; add a new entry when introducing a recipe. RECIPE_TYPE_TO_CLASS: dict[RecipeType, type[ModelOptRecipeBase]] = { RecipeType.PTQ: ModelOptPTQRecipe, + RecipeType.AUTO_QUANTIZE: ModelOptAutoQuantizeRecipe, RecipeType.SPECULATIVE_EAGLE: ModelOptEagleRecipe, RecipeType.SPECULATIVE_DFLASH: ModelOptDFlashRecipe, RecipeType.SPECULATIVE_MEDUSA: ModelOptMedusaRecipe, diff --git a/modelopt/recipe/loader.py b/modelopt/recipe/loader.py index 0a9218ff7d0..6af6d0a8a7a 100644 --- a/modelopt/recipe/loader.py +++ b/modelopt/recipe/loader.py @@ -42,6 +42,7 @@ # must contain 'quantize'" instead of pydantic's generic missing-field error. _REQUIRED_SECTION_PER_RECIPE_TYPE: dict[RecipeType, str] = { RecipeType.PTQ: "quantize", + RecipeType.AUTO_QUANTIZE: "auto_quantize", RecipeType.SPECULATIVE_EAGLE: "eagle", RecipeType.SPECULATIVE_DFLASH: "dflash", RecipeType.SPECULATIVE_MEDUSA: "medusa", @@ -171,9 +172,9 @@ def _load_recipe_from_file( raw = yaml.safe_load(recipe_file.read_text()) or {} if not isinstance(raw, dict) or required_section not in raw: - kind = ( - rtype.value.split("_", 1)[-1].upper() if "_" in rtype.value else rtype.value.upper() - ) + # Strip only the ``speculative_`` prefix so multi-word non-speculative types + # (e.g. ``auto_quantize``) keep their full name: AUTO_QUANTIZE, not QUANTIZE. + kind = rtype.value.removeprefix("speculative_").upper() raise ValueError(f"{kind} recipe file {recipe_file} must contain {required_section!r}.") # Passing ``schema_type=schema_class`` to ``load_config`` enables typed-list diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index a48655e1fb2..c7803ecf838 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -127,9 +127,11 @@ def _make_fresh_quantizer_for_attr(module: nn.Module, attr_name: str) -> nn.Modu def estimate_quant_compression(quant_cfg: QuantizeConfig) -> float: """Estimate the compression ratio of a quantization configuration. - Right now, we find the minimum compression ratio across all quantizer attribute configs. - This is not perfect but is a good proxy for the overall compression ratio. We will improve - this in future releases. + Effective bits per element resolve in priority order: (1) recipe-level + ``quant_cfg.effective_bits``; (2) per-entry ``cfg.effective_bits`` (library default, + e.g. NVFP4 = 4.5); (3) the ``num_bits`` heuristic (``num_bits / 16`` for ints, + ``(E + M + 1) / 16`` for FP tuples). Per-entry values are aggregated via ``min``, which + still under-counts activation cost for mixed weight+activation formats. Args: quant_cfg: The quantization configuration to estimate compression for. @@ -137,6 +139,8 @@ def estimate_quant_compression(quant_cfg: QuantizeConfig) -> float: Returns: float: The estimated compression ratio (0.0 to 1.0). """ + if quant_cfg.effective_bits is not None: + return quant_cfg.effective_bits / 16.0 def estimate_quant_compression_for_quantizer(quantizer_attr_cfg): if isinstance(quantizer_attr_cfg, list): @@ -147,6 +151,9 @@ def estimate_quant_compression_for_quantizer(quantizer_attr_cfg): # Handle raw quantizer cfg dicts (e.g. {"num_bits": (4, 3), "axis": None}) if not quantizer_attr_cfg.get("enable", True): return 1.0 + effective_bits = quantizer_attr_cfg.get("effective_bits") + if effective_bits is not None: + return effective_bits / 16 num_bits = quantizer_attr_cfg.get("num_bits") if num_bits is None: return 1.0 @@ -160,6 +167,8 @@ def estimate_quant_compression_for_quantizer(quantizer_attr_cfg): if isinstance(quantizer_attr_cfg, QuantizerAttributeConfig): if not quantizer_attr_cfg.enable: return 1.0 + if quantizer_attr_cfg.effective_bits is not None: + return quantizer_attr_cfg.effective_bits / 16 if not hasattr(quantizer_attr_cfg, "num_bits"): return 1.0 if isinstance(quantizer_attr_cfg.num_bits, tuple): diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 344292264fa..f4a935239a8 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -323,6 +323,22 @@ class QuantizerAttributeConfig(ModeloptBaseConfig): #. String specifying the quantization format. This is current used only for custom backends.""", ) + effective_bits: float | None = ModeloptField( + default=None, + title="Effective bits per element (autoquant cost).", + description=( + "Per-format effective bits for the autoquant cost model; overrides the " + "``num_bits`` heuristic for this entry (e.g. NVFP4 = 4.5). Must be in (0, 16]." + ), + ) + + @field_validator("effective_bits") + @classmethod + def _validate_effective_bits(cls, v: float | None) -> float | None: + if v is not None and not (0 < v <= 16): + raise ValueError(f"effective_bits must be in (0, 16], got {v}") + return v + @model_validator(mode="before") @classmethod def validate_config(cls, values): @@ -1317,6 +1333,22 @@ class QuantizeConfig(ModeloptBaseConfig): validate_default=True, ) + effective_bits: float | None = ModeloptField( + default=None, + title="Effective bits per element (autoquant cost override)", + description=( + "Recipe-level override for the autoquant cost model; replaces the per-entry " + "``num_bits`` heuristic for the whole config. Must be in (0, 16]." + ), + ) + + @field_validator("effective_bits") + @classmethod + def _validate_effective_bits(cls, v: float | None) -> float | None: + if v is not None and not (0 < v <= 16): + raise ValueError(f"effective_bits must be in (0, 16], got {v}") + return v + @field_validator("quant_cfg", mode="before") @classmethod def normalize_quant_cfg( diff --git a/modelopt_recipes/configs/auto_quantize/units/base_cost_excluded_layers.yaml b/modelopt_recipes/configs/auto_quantize/units/base_cost_excluded_layers.yaml new file mode 100644 index 00000000000..15437e25fbb --- /dev/null +++ b/modelopt_recipes/configs/auto_quantize/units/base_cost_excluded_layers.yaml @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Base cost-excluded layer patterns for AutoQuantize (spliced into a recipe's cost_excluded_layers +# via ``$import``, and appended by the deprecated-CLI shim). These are VL patterns; appending them +# unconditionally is safe: on a non-VL model nothing matches, so cost_weight 0 is a no-op, while on +# a VL model it keeps the vision tower / MTP out of the bit-budget denominator. This avoids the old +# model-introspection path while preserving VL cost behavior. + +# modelopt-schema: modelopt.recipe.config.LayerPatternList + - "*visual*" + - "*mtp*" + - "*vision_tower*" diff --git a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml new file mode 100644 index 00000000000..fd27046d608 --- /dev/null +++ b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Base set of non-quantizable layer patterns excluded from the AutoQuantize search (gates, +# routers, conv/mixer projections, output/embedding heads, vision towers). Spliced into a +# recipe's ``disabled_layers`` via ``$import``; recipes append architecture-specific patterns. + +# modelopt-schema: modelopt.recipe.config.LayerPatternList + - "*block_sparse_moe.gate*" + - "*linear_attn.conv1d*" + - "*linear_attn.in_proj_a*" + - "*linear_attn.in_proj_b*" + - "*mixer.conv1d*" + - "*mlp.gate.*" + - "*mlp.shared_expert_gate.*" + - "*output_layer*" + - "*proj_out.*" + - "*router*" + - "output.*" + - "*embed_vision*" + - "*vision_tower*" + - "*visual*" + # Qwen3.6 MoE (e.g. Qwen/Qwen3.6-35B-A3B): must be disabled or export fails at linear fusion. + # Kept in the base set because the deprecated --auto_quantize_* CLI can't inject arch patterns. + - "*shared_expert_gate*" diff --git a/modelopt_recipes/configs/numerics/nvfp4.yaml b/modelopt_recipes/configs/numerics/nvfp4.yaml index 88598e36e85..6ef99602f4e 100644 --- a/modelopt_recipes/configs/numerics/nvfp4.yaml +++ b/modelopt_recipes/configs/numerics/nvfp4.yaml @@ -21,3 +21,6 @@ block_sizes: -1: 16 type: dynamic scale_bits: e4m3 +# Autoquant LP cost: 4 value bits + an FP8 (8-bit) scale per 16-element block = 4.5 bits/element +# (the num_bits heuristic under-counts this as 4.0). Read only by autoquant. +effective_bits: 4.5 diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_5p4bits.yaml new file mode 100644 index 00000000000..7f4f337100d --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_5p4bits.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# AutoQuantize: per-layer search over {NVFP4 (W4A4), FP8 (W8A8)} at 5.4 effective bits. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + nvfp4: configs/ptq/presets/model/nvfp4 + fp8: configs/ptq/presets/model/fp8 + +metadata: + recipe_type: auto_quantize + description: Mixed NVFP4 + FP8 per-layer search at 5.4 effective bits. + +auto_quantize: + constraints: + effective_bits: 5.4 + + candidate_formats: + - $import: nvfp4 + - $import: fp8 + + auto_quantize_method: gradient + score_size: 128 + + # Base (model-agnostic) non-quantizable layers, spliced from the shared unit. Arch-specific + # models use a recipe under huggingface//auto_quantize/ that appends to this set. + disabled_layers: + - $import: base_disabled_layers diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits.yaml new file mode 100644 index 00000000000..aac822ed0c0 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits.yaml @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# AutoQuantize: per-layer search over {NVFP4 (W4A4), FP8 (W8A8)} at 5.4 effective bits, using the +# kl_div scoring method. kl_div needs no labels/backprop, so it suits models without backprop +# support (e.g. Llama-4) where the default gradient method cannot run. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + nvfp4: configs/ptq/presets/model/nvfp4 + fp8: configs/ptq/presets/model/fp8 + +metadata: + recipe_type: auto_quantize + description: Mixed NVFP4 + FP8 per-layer search at 5.4 effective bits, kl_div scoring (no backprop). + +auto_quantize: + constraints: + effective_bits: 5.4 + + candidate_formats: + - $import: nvfp4 + - $import: fp8 + + auto_quantize_method: kl_div + score_size: 128 + + # Base (model-agnostic) non-quantizable layers, spliced from the shared unit. Arch-specific + # models use a recipe under huggingface//auto_quantize/ that appends to this set. + disabled_layers: + - $import: base_disabled_layers diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml new file mode 100644 index 00000000000..ac9546d049d --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# AutoQuantize: per-layer search over {NVFP4 W4A4 (weight-MSE + FP8 sweep), FP8} at 6.0 bits. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + nvfp4_mse: configs/ptq/presets/model/nvfp4_w4a4_weight_mse_fp8_sweep + fp8: configs/ptq/presets/model/fp8 + +metadata: + recipe_type: auto_quantize + description: Mixed NVFP4 (weight-MSE + FP8 sweep) + FP8 per-layer search at 6.0 effective bits. + +auto_quantize: + constraints: + effective_bits: 6.0 + + candidate_formats: + - $import: nvfp4_mse + - $import: fp8 + + auto_quantize_method: gradient + score_size: 128 + + # Base (model-agnostic) non-quantizable layers, spliced from the shared unit. Arch-specific + # models use a recipe under huggingface//auto_quantize/ that appends to this set. + disabled_layers: + - $import: base_disabled_layers diff --git a/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml b/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml new file mode 100644 index 00000000000..56fc5fd789a --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# AutoQuantize recipe: per-layer search over {FP8 (W8A8), NVFP4 weight-only (W4A16)} +# at 6.0 effective bits, active-MoE cost model. Recipe form of the CLI command: +# --qformat fp8,w4a16_nvfp4 --auto_quantize_bits 6.0 --auto_quantize_method gradient +# --auto_quantize_cost_model active_moe --auto_quantize_active_moe_expert_ratio 0.03125 + +# 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: >- + Mixed FP8 + NVFP4-weight-only per-layer search at 6.0 effective bits with the + active-MoE cost model (expert ratio 0.03125). + +auto_quantize: + constraints: + effective_bits: 6.0 + cost_model: active_moe + cost: + active_moe_expert_ratio: 0.03125 + + candidate_formats: + # LP cost comes from the presets' numerics (NVFP4 weight = 4.5 via configs/numerics/nvfp4, + # FP8 = 8); no per-candidate effective_bits override needed. + - $import: fp8 + - $import: w4a16_nvfp4 + + auto_quantize_method: gradient + score_size: 128 + # kv_cache omitted -> falls back to --kv_cache_qformat (none in the reference command). + + # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under + # huggingface//auto_quantize/ that extends this set. + disabled_layers: + - $import: base_disabled_layers diff --git a/modelopt_recipes/general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits.yaml b/modelopt_recipes/general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits.yaml new file mode 100644 index 00000000000..d0135f52a4d --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# AutoQuantize: per-layer search over {W4A8 AWQ-beta, FP8 (W8A8)} at 6.0 effective bits. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + base_disabled_layers: configs/auto_quantize/units/base_disabled_layers + w4a8_awq_beta: configs/ptq/presets/model/w4a8_awq_beta + fp8: configs/ptq/presets/model/fp8 + +metadata: + recipe_type: auto_quantize + description: Mixed W4A8 AWQ-beta + FP8 per-layer search at 6.0 effective bits. + +auto_quantize: + constraints: + effective_bits: 6.0 + + candidate_formats: + - $import: w4a8_awq_beta + - $import: fp8 + + auto_quantize_method: gradient + score_size: 128 + + # Base (model-agnostic) non-quantizable layers, spliced from the shared unit. Arch-specific + # models use a recipe under huggingface//auto_quantize/ that appends to this set. + disabled_layers: + - $import: base_disabled_layers diff --git a/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml new file mode 100644 index 00000000000..201a70614eb --- /dev/null +++ b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml @@ -0,0 +1,58 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Qwen3.6 MoE AutoQuantize: mixed FP8 + NVFP4 weight-only at 6.0 effective bits, active-MoE +# cost model. Carries the architecture's disabled-layer patterns explicitly in the recipe +# (the set kept in sync with example_utils._get_auto_quantize_disabled_layers for a Qwen model; +# pinned by tests/examples/llm_ptq/test_hf_ptq_args.py). + +# 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: FP8 + NVFP4-weight-only per-layer search at 6.0 effective bits, active-MoE + cost model (expert ratio 0.03125), with architecture-specific disabled layers. + +auto_quantize: + constraints: + effective_bits: 6.0 + cost_model: active_moe + cost: + active_moe_expert_ratio: 0.03125 + + candidate_formats: + - $import: fp8 + - $import: w4a16_nvfp4 + + auto_quantize_method: gradient + score_size: 128 + + # Shared base patterns spliced in; this architecture adds the Qwen MoE shared-expert gate. + disabled_layers: + - $import: base_disabled_layers + - "*shared_expert_gate*" + + # VL model: kept out of the bit-budget denominator (cost_weight 0) so they don't inflate the + # budget. Same patterns are also in disabled_layers above (excluded from search) — cost-exclusion + # is a separate, independent role. + cost_excluded_layers: + - "*visual*" + - "*mtp*" + - "*vision_tower*" diff --git a/tests/_test_utils/examples/hf_ptq_utils.py b/tests/_test_utils/examples/hf_ptq_utils.py index 1742158ae07..16c5952ba98 100644 --- a/tests/_test_utils/examples/hf_ptq_utils.py +++ b/tests/_test_utils/examples/hf_ptq_utils.py @@ -24,7 +24,8 @@ @dataclass class PTQCommand: - quant: str + quant: str | None = None + recipe: str | None = None tasks: str = "quant" calib: int = 16 sparsity: str | None = None @@ -32,7 +33,6 @@ class PTQCommand: trust_remote_code: bool = False calib_dataset: str = "cnn_dailymail" calib_batch_size: int | None = None - auto_quantize_bits: float | None = None tp: int | None = None pp: int | None = None min_sm: int | None = None @@ -40,6 +40,10 @@ class PTQCommand: min_gpu: int | None = None batch: int | None = None + def __post_init__(self): + if (self.quant is None) == (self.recipe is None): + raise ValueError("Exactly one of `quant` or `recipe` must be set.") + def run(self, model_path: str): if self.min_sm and torch.cuda.get_device_capability() < ( self.min_sm // 10, diff --git a/tests/_test_utils/examples/run_command.py b/tests/_test_utils/examples/run_command.py index 0cccd97c644..dcef541b71b 100644 --- a/tests/_test_utils/examples/run_command.py +++ b/tests/_test_utils/examples/run_command.py @@ -62,7 +62,7 @@ def run_command_in_background( return process -def run_hf_ptq_command(*, model: str, quant: str, vlm: bool = False, **kwargs): +def run_hf_ptq_command(*, model: str, quant: str | None = None, vlm: bool = False, **kwargs): kwargs.update({"model": model, "quant": quant}) kwargs.setdefault("tasks", "quant") kwargs.setdefault("calib", 16) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index b06c1357411..7160dd38cb4 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -16,10 +16,14 @@ import importlib import sys from pathlib import Path -from types import SimpleNamespace import pytest +from modelopt.recipe import load_recipe +from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints +from modelopt.recipe.presets import QUANT_CFG_CHOICES +from modelopt.torch.quantization.config import QuantizeConfig + _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -28,11 +32,6 @@ def _import_hf_ptq(monkeypatch): return importlib.import_module("hf_ptq") -def _import_example_utils(monkeypatch): - monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) - return importlib.import_module("example_utils") - - def _parse_hf_ptq_args(monkeypatch, *args): hf_ptq = _import_hf_ptq(monkeypatch) monkeypatch.setattr(sys, "argv", ["hf_ptq.py", *args]) @@ -46,63 +45,124 @@ def _parse_hf_ptq_args(monkeypatch, *args): return hf_ptq, parsed_args -def test_parse_args_rejects_autoquant_image_calibration(monkeypatch): - hf_ptq = _import_hf_ptq(monkeypatch) - monkeypatch.setattr( - sys, - "argv", - [ - "hf_ptq.py", - "--pyt_ckpt_path", - "nemotron-vl", - "--auto_quantize_bits", - "5.0", - "--calib_with_images", - ], +def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): + """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" ) - - with pytest.raises(SystemExit) as error: - hf_ptq.parse_args() - - assert error.value.code == 2 - - -def test_load_model_keeps_nemotron_vl_text_calibration_for_autoquant(monkeypatch): + aq = load_recipe("general/auto_quantize/nvfp4_fp8_at_5p4bits").auto_quantize + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) + + assert inputs["constraints"] == {"effective_bits": 5.4, "cost_model": "weight"} + assert inputs["kv_cache_quant_cfg"] is None + assert inputs["method"] == "gradient" + assert inputs["score_size"] == 128 + # disabled_layers come straight from the recipe (no model introspection). + assert inputs["disabled_layers"] == aq.disabled_layers + assert "*output_layer*" in inputs["disabled_layers"] + # Candidates resolve to the exact preset dicts mtq expects (preset identity preserved). + assert inputs["quantization_formats"][0] == QUANT_CFG_CHOICES["nvfp4"] + assert inputs["quantization_formats"][1] == QUANT_CFG_CHOICES["fp8"] + + +def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch): + """Top-level cost_excluded_layers maps to the mtq constraints.cost.excluded_module_name_patterns + key (distinct from disabled_layers), so a cost-exclusion recipe matches the nested mtq dict.""" hf_ptq, args = _parse_hf_ptq_args( - monkeypatch, - "--pyt_ckpt_path", - "nemotron-vl", - "--auto_quantize_bits", - "5.0", + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" ) - fake_model = SimpleNamespace(device="cpu") - fake_tokenizer = SimpleNamespace(padding_side="right", pad_token="") + aq = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe" + ).auto_quantize + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) + + # cost-exclusion is hoisted to a sibling of disabled_layers but still reaches the mtq cost dict. + assert aq.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"] + assert inputs["constraints"]["cost"] == { + "active_moe_expert_ratio": 0.03125, + "excluded_module_name_patterns": ["*visual*", "*mtp*", "*vision_tower*"], + } + # The two exclusions are independent: cost-excluded patterns are also disabled here, but the + # roles (cost-accounting vs search) are tracked separately. + assert "*visual*" in inputs["disabled_layers"] - monkeypatch.setattr(hf_ptq, "get_model", lambda *args, **kwargs: fake_model) - monkeypatch.setattr(hf_ptq, "get_model_type", lambda model: "qwen2") - monkeypatch.setattr(hf_ptq, "get_tokenizer", lambda *args, **kwargs: fake_tokenizer) - monkeypatch.setattr(hf_ptq, "is_nemotron_vl", lambda model: True) - full_model, language_model, _, _, _, tokenizer, _, _, _ = hf_ptq.load_model(args) +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( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" + ) + non_safe = next(k for k in QUANT_CFG_CHOICES if k not in hf_ptq._AUTO_QUANTIZE_QFORMATS) + aq = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=4.8), + candidate_formats=[ + QuantizeConfig(**QUANT_CFG_CHOICES["fp8"]), + QuantizeConfig(**QUANT_CFG_CHOICES[non_safe]), + ], + ) + with pytest.raises(ValueError, match="not supported for unified checkpoint export"): + hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) - assert args.calib_with_images is False - assert full_model is fake_model - assert language_model is fake_model - assert tokenizer is fake_tokenizer +def test_autoquant_warns_on_custom_candidate(monkeypatch): + """A candidate matching no shipped preset can't be export-verified, so it warns (not blocks).""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" + ) + custom = QuantizeConfig(quant_cfg=[{"quantizer_name": "*", "enable": False}]) + aq = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=4.8), + candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES["fp8"]), custom], + ) + with pytest.warns(UserWarning, match="export compatibility cannot be verified"): + hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) -def test_qwen_autoquant_disabled_layers_are_scoped_to_qwen_models(monkeypatch): - example_utils = _import_example_utils(monkeypatch) - qwen_model = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_moe")) - llama_model = SimpleNamespace(config=SimpleNamespace(model_type="llama")) - qwen_only_patterns = { - "*shared_expert_gate*", - } - monkeypatch.setattr(example_utils, "is_multimodal_model", lambda model: False) +def test_autoquant_export_guard_not_bypassed_by_effective_bits(monkeypatch): + """A non-export-safe preset can't dodge the guard by adding a cost-only effective_bits override.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" + ) + non_safe = next(k for k in QUANT_CFG_CHOICES if k not in hf_ptq._AUTO_QUANTIZE_QFORMATS) + tampered = QuantizeConfig(**{**QUANT_CFG_CHOICES[non_safe], "effective_bits": 4.5}) + aq = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=5.4), + candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES["fp8"]), tampered], + ) + with pytest.raises(ValueError, match="not supported for unified checkpoint export"): + hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) - qwen_disabled_layers = set(example_utils._get_auto_quantize_disabled_layers(qwen_model)) - llama_disabled_layers = set(example_utils._get_auto_quantize_disabled_layers(llama_model)) - assert qwen_only_patterns <= qwen_disabled_layers - assert qwen_only_patterns.isdisjoint(llama_disabled_layers) +def test_autoquant_config_from_deprecated_cli_flags(monkeypatch): + """The deprecated --auto_quantize_* flags convert to an AutoQuantizeConfig with the shared + base disabled + cost-excluded patterns appended (no new flags, no model introspection).""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--qformat", + "fp8,nvfp4", + "--auto_quantize_bits", + "5.4", + "--auto_quantize_cost_model", + "active_moe", + "--auto_quantize_active_moe_expert_ratio", + "0.03125", + "--kv_cache_qformat", + "none", + ) + aq = hf_ptq._auto_quantize_config_from_cli(args) + + assert aq.constraints.effective_bits == 5.4 + assert aq.constraints.cost_model == "active_moe" + assert aq.constraints.cost.active_moe_expert_ratio == 0.03125 + assert aq.auto_quantize_method == "gradient" + assert aq.score_size == 128 + # candidates come from --qformat and resolve to their shipped presets. + assert [hf_ptq._match_candidate_to_preset(f)[0] for f in aq.candidate_formats] == [ + "fp8", + "nvfp4", + ] + # base disabled + base cost-excluded appended from the shared units (no introspection). + assert "*output_layer*" in aq.disabled_layers + assert aq.cost_excluded_layers == ["*visual*", "*mtp*", "*vision_tower*"] diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 7242b2234f9..4b66bad254e 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -78,28 +78,25 @@ def test_ptq_whisper(command): PTQCommand(quant="w4a8_awq", kv_cache_quant="none"), PTQCommand(quant="nvfp4"), PTQCommand(quant="nvfp4_awq"), - # autoquant + # autoquant (recipe-driven) PTQCommand( - quant="int4_awq,nvfp4,fp8,w4a8_awq", + recipe="general/auto_quantize/nvfp4_fp8_at_5p4bits", calib_batch_size=4, - auto_quantize_bits=6.4, kv_cache_quant="none", ), # kv_cache PTQCommand(quant="nvfp4_awq", kv_cache_quant="nvfp4"), PTQCommand(quant="fp8", kv_cache_quant="fp8_cast", min_sm=89), - # autoquant_kv_cache + # autoquant_kv_cache (recipe-driven; KV via --kv_cache_quant fallback) PTQCommand( - quant="nvfp4,fp8", + recipe="general/auto_quantize/nvfp4_fp8_at_5p4bits", kv_cache_quant="fp8", calib_batch_size=4, - auto_quantize_bits=6.4, ), PTQCommand( - quant="nvfp4,fp8", + recipe="general/auto_quantize/nvfp4_fp8_at_5p4bits", kv_cache_quant="nvfp4", calib_batch_size=4, - auto_quantize_bits=6.4, ), # sm89 PTQCommand(quant="fp8", min_sm=89), diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index f4c27f74b2a..3aaacaa3e0e 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -27,6 +27,7 @@ import modelopt.torch.quantization.config as qcfg from modelopt.recipe.config import ( + ModelOptAutoQuantizeRecipe, ModelOptDFlashRecipe, ModelOptEagleRecipe, ModelOptPTQRecipe, @@ -1689,3 +1690,141 @@ def test_import_imports_not_a_dict_raises(tmp_path): config_file.write_text("imports:\n - some/path\nkey: value\n") with pytest.raises(ValueError, match="must be a dict"): load_config(config_file) + + +# --------------------------------------------------------------------------- +# load_recipe — AutoQuantize recipes +# --------------------------------------------------------------------------- + +_AQ_MINIMAL_BODY = ( + "metadata:\n" + " recipe_type: auto_quantize\n" + "auto_quantize:\n" + " constraints:\n" + " effective_bits: 4.8\n" + " candidate_formats:\n" + " - algorithm: max\n" + " quant_cfg: []\n" + " - algorithm: max\n" + " quant_cfg: []\n" +) + + +def test_load_recipe_autoquantize_minimal(tmp_path): + """Minimal AutoQuantize recipe loads with the right type and field defaults.""" + recipe_file = tmp_path / "aq.yml" + recipe_file.write_text(_AQ_MINIMAL_BODY) + recipe = load_recipe(recipe_file) + + assert recipe.recipe_type == RecipeType.AUTO_QUANTIZE + assert isinstance(recipe, ModelOptAutoQuantizeRecipe) + aq = recipe.auto_quantize + assert aq.auto_quantize_method == "gradient" + assert aq.score_size == 128 + assert aq.kv_cache is None + assert aq.constraints.effective_bits == 4.8 + assert aq.constraints.cost_model == "weight" + assert aq.constraints.cost is None + assert len(aq.candidate_formats) == 2 + + +def test_load_recipe_autoquantize_active_moe_cost_roundtrip(tmp_path): + """cost_model + cost.active_moe_expert_ratio parse and dump to the mtq constraints dict shape.""" + recipe_file = tmp_path / "aq.yml" + recipe_file.write_text( + "metadata:\n" + " recipe_type: auto_quantize\n" + "auto_quantize:\n" + " constraints:\n" + " effective_bits: 6.0\n" + " cost_model: active_moe\n" + " cost:\n" + " active_moe_expert_ratio: 0.03125\n" + " candidate_formats:\n" + " - algorithm: max\n" + " quant_cfg: []\n" + " - algorithm: max\n" + " quant_cfg: []\n" + ) + constraints = load_recipe(recipe_file).auto_quantize.constraints + assert constraints.cost_model == "active_moe" + assert constraints.cost.active_moe_expert_ratio == 0.03125 + assert constraints.model_dump(exclude_none=True) == { + "effective_bits": 6.0, + "cost_model": "active_moe", + "cost": {"active_moe_expert_ratio": 0.03125}, + } + + +def test_load_recipe_autoquantize_missing_section_raises(tmp_path): + """Missing auto_quantize section gives the clean loader-level error.""" + bad = tmp_path / "bad.yml" + bad.write_text("metadata:\n recipe_type: auto_quantize\n") + with pytest.raises( + ValueError, match=r"AUTO_QUANTIZE recipe file .* must contain 'auto_quantize'" + ): + load_recipe(bad) + + +def test_load_recipe_autoquantize_empty_candidates_raises(tmp_path): + """Empty candidate_formats is rejected (a single format is valid — bf16 is implicit).""" + bad = tmp_path / "bad.yml" + bad.write_text( + "metadata:\n recipe_type: auto_quantize\n" + "auto_quantize:\n constraints:\n effective_bits: 4.8\n" + " candidate_formats: []\n" + ) + with pytest.raises(ValueError, match="at least 1"): + load_recipe(bad) + + +def test_load_recipe_autoquantize_single_candidate_ok(tmp_path): + """A single candidate format is valid: the {format, bf16} per-layer search (bf16 implicit).""" + recipe_file = tmp_path / "single.yml" + recipe_file.write_text( + "metadata:\n recipe_type: auto_quantize\n" + "auto_quantize:\n constraints:\n effective_bits: 6.0\n" + " candidate_formats:\n - algorithm: max\n quant_cfg: []\n" + ) + aq = load_recipe(recipe_file).auto_quantize + assert len(aq.candidate_formats) == 1 + + +def test_load_recipe_autoquantize_effective_bits_out_of_range_raises(tmp_path): + """effective_bits outside (0, 16] is rejected.""" + bad = tmp_path / "bad.yml" + bad.write_text(_AQ_MINIMAL_BODY.replace("effective_bits: 4.8", "effective_bits: 20")) + with pytest.raises(ValueError, match="effective_bits"): + load_recipe(bad) + + +def test_load_recipe_autoquantize_builtin_active_moe(): + """The shipped active-MoE AutoQuantize recipe resolves to the expected values.""" + recipe = load_recipe("general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe") + assert isinstance(recipe, ModelOptAutoQuantizeRecipe) + aq = recipe.auto_quantize + assert aq.constraints.effective_bits == 6.0 + assert aq.constraints.cost_model == "active_moe" + assert aq.constraints.cost.active_moe_expert_ratio == 0.03125 + assert aq.auto_quantize_method == "gradient" + assert aq.kv_cache is None + # No per-candidate override; NVFP4 cost (4.5) comes from configs/numerics/nvfp4. + assert all(c.effective_bits is None for c in aq.candidate_formats) + + +@pytest.mark.parametrize( + "recipe_path", + [ + "general/auto_quantize/nvfp4_fp8_at_5p4bits", + "general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits", + "general/auto_quantize/nvfp4_mse_fp8_at_6p0bits", + "general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits", + "general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe", + ], +) +def test_load_recipe_autoquantize_builtin_general(recipe_path): + """Every shipped general AutoQuantize recipe loads and has >= 2 candidate formats.""" + recipe = load_recipe(recipe_path) + assert isinstance(recipe, ModelOptAutoQuantizeRecipe) + assert len(recipe.auto_quantize.candidate_formats) >= 2 + assert recipe.auto_quantize.auto_quantize_method in ("gradient", "kl_div") diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 1978f389069..b85feb32649 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -173,6 +173,27 @@ def test_quant_recipe_hparam_zero_cost_weight(): assert hparam.get_cost(QuantRecipe(mtq.INT8_DEFAULT_CFG)) == pytest.approx(0.0) +def test_quant_recipe_hparam_cost_weight_and_effective_bits_compose(): + """cost_weight (active_moe) and effective_bits stack multiplicatively in get_cost.""" + model_test = mtq.quantize(torch.nn.Linear(4, 16), mtq.NVFP4_DEFAULT_CFG) + numel = model_test.weight.numel() + hparam = QuantRecipeHparam( + [QuantRecipe(mtq.NVFP4_DEFAULT_CFG)], + quant_modules=[model_test], + quant_module_names=["layers.0.mlp.experts.0.down_proj"], + cost_weight=0.03125, + ) + + # NVFP4's library-default effective_bits (4.5, from configs/numerics/nvfp4) stacks with cost_weight: + # numel * cost_weight * (effective_bits / 16). + assert hparam.get_cost(QuantRecipe(mtq.NVFP4_DEFAULT_CFG)) == pytest.approx( + numel * 0.03125 * (4.5 / 16) + ) + # A recipe-level effective_bits override (8.0) wins over the library default and still stacks. + override = QuantRecipe({**mtq.NVFP4_DEFAULT_CFG, "effective_bits": 8.0}, name="NVFP4_8B") + assert hparam.get_cost(override) == pytest.approx(numel * 0.03125 * 0.5) + + def test_auto_quantize_cost_model_excludes_module_name_patterns(): cost_model = get_auto_quantize_cost_model("weight") cost_constraints = {EXCLUDED_MODULE_NAME_PATTERNS_KEY: ["*visual*", "*vision_tower*", "*mtp*"]} @@ -507,29 +528,30 @@ def _raise_local_total_weight_size(modules): def test_estimate_quant_compression(): + # NVFP4 weight/input carry effective_bits=4.5 from configs/numerics/nvfp4 -> 4.5/16 = 0.28125. nvfp4_affine_kv_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_AFFINE_KV_CFG) - assert estimate_quant_compression(nvfp4_affine_kv_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_affine_kv_cfg) == 0.28125 nvfp4_awq_clip_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_AWQ_CLIP_CFG) - assert estimate_quant_compression(nvfp4_awq_clip_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_awq_clip_cfg) == 0.28125 nvfp4_awq_full_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_AWQ_FULL_CFG) - assert estimate_quant_compression(nvfp4_awq_full_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_awq_full_cfg) == 0.28125 nvfp4_awq_lite_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_AWQ_LITE_CFG) - assert estimate_quant_compression(nvfp4_awq_lite_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_awq_lite_cfg) == 0.28125 nvfp4_default_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG) - assert estimate_quant_compression(nvfp4_default_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_default_cfg) == 0.28125 nvfp4_kv_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_KV_CFG) - assert estimate_quant_compression(nvfp4_kv_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_kv_cfg) == 0.28125 nvfp4_kv_rotate_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_KV_ROTATE_CFG) - assert estimate_quant_compression(nvfp4_kv_rotate_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_kv_rotate_cfg) == 0.28125 nvfp4_svdquant_default_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_SVDQUANT_DEFAULT_CFG) - assert estimate_quant_compression(nvfp4_svdquant_default_cfg) == 0.25 + assert estimate_quant_compression(nvfp4_svdquant_default_cfg) == 0.28125 int8_default_cfg = mtq.config.QuantizeConfig(**mtq.INT8_DEFAULT_CFG) assert estimate_quant_compression(int8_default_cfg) == 0.5 @@ -576,6 +598,82 @@ def test_estimate_quant_compression(): assert estimate_quant_compression(fp8_affine_kv_cfg) == 0.5 +def test_estimate_quant_compression_effective_bits_override(): + """Recipe-level ``QuantizeConfig.effective_bits`` overrides the per-entry library default.""" + # NVFP4 weight carries effective_bits=4.5 from configs/numerics/nvfp4 (per-entry default). + nvfp4_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG) + assert nvfp4_cfg.effective_bits is None # no recipe-level override + assert estimate_quant_compression(nvfp4_cfg) == 4.5 / 16.0 # per-entry library default + + # A recipe-level QuantizeConfig.effective_bits override wins over the per-entry default. + nvfp4_cfg_overridden = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG, effective_bits=8.0) + assert estimate_quant_compression(nvfp4_cfg_overridden) == 8.0 / 16.0 + + # Override can also represent a higher cost (e.g., conservative for a sensitive recipe). + nvfp4_cfg_high = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG, effective_bits=16.0) + assert estimate_quant_compression(nvfp4_cfg_high) == 1.0 + + # Out-of-range values are rejected by the Pydantic validator. + with pytest.raises(ValueError, match="effective_bits must be in"): + mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG, effective_bits=0.0) + with pytest.raises(ValueError, match="effective_bits must be in"): + mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG, effective_bits=17.0) + + +def test_estimate_quant_compression_per_entry_effective_bits(): + """Per-entry ``effective_bits`` overrides the heuristic; recipe-level wins over it; min across entries.""" + # num_bits=(2,1) -> heuristic 0.25, but per-entry library default is 4.5. + cfg = mtq.config.QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": (2, 1), "effective_bits": 4.5}, + }, + ], + algorithm="max", + ) + assert cfg.effective_bits is None + assert estimate_quant_compression(cfg) == 4.5 / 16.0 + + # Recipe-level override (layer 1) wins over the per-entry value (layer 2). + cfg_recipe_override = mtq.config.QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": (2, 1), "effective_bits": 4.5}, + }, + ], + algorithm="max", + effective_bits=8.0, + ) + assert estimate_quant_compression(cfg_recipe_override) == 8.0 / 16.0 + + # min across entries: weight 4.5/16 = 0.28125 vs heuristic fp8 input 0.5 -> 0.28125. + cfg_mixed = mtq.config.QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": (2, 1), "effective_bits": 4.5}, + }, + {"quantizer_name": "*input_quantizer", "cfg": {"num_bits": (4, 3)}}, + ], + algorithm="max", + ) + assert estimate_quant_compression(cfg_mixed) == 4.5 / 16.0 + + # Per-entry out-of-range is rejected by the QuantizerAttributeConfig validator. + with pytest.raises(ValueError, match="effective_bits must be in"): + mtq.config.QuantizeConfig( + quant_cfg=[ + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": (2, 1), "effective_bits": 20.0}, + }, + ], + algorithm="max", + ) + + @pytest.mark.parametrize("method", ["gradient", "kl_div"]) def test_auto_quantize_checkpoint_resume(method, tmp_path, capsys): """Test that checkpoint can be used to resume an interrupted search."""