From 5512ac98d598f0350b34f7e7772d32630f9cffc4 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 01:55:02 +0000 Subject: [PATCH 01/17] autoquant: add effective_bits override to the cost model Add an effective_bits field at two levels for the autoquant LP cost model: QuantizeConfig (recipe-level override) and QuantizerAttributeConfig (per-format library default). estimate_quant_compression resolves in priority order: recipe-level > per-entry > num_bits heuristic, fixing the heuristic's undercount of block-scaled formats (e.g. NVFP4 = 4.5 vs 4.0). Per-entry values are aggregated via min. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- modelopt/torch/quantization/algorithms.py | 15 +++- modelopt/torch/quantization/config.py | 32 ++++++++ .../unit/torch/quantization/test_autoquant.py | 75 +++++++++++++++++++ 3 files changed, 119 insertions(+), 3 deletions(-) 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/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 1978f389069..deb8f8dcf5e 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -576,6 +576,81 @@ 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 num_bits heuristic; unset falls back to it.""" + # NVFP4 — heuristic returns 4.0 bits / 16 = 0.25, but true effective bits is 4.5. + nvfp4_cfg = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG) + assert nvfp4_cfg.effective_bits is None + assert estimate_quant_compression(nvfp4_cfg) == 0.25 # heuristic baseline + + nvfp4_cfg_overridden = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG, effective_bits=4.5) + assert estimate_quant_compression(nvfp4_cfg_overridden) == 4.5 / 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.""" From 77cc7b10d6b392789bbbf307c823a95d04efd107 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 01:55:17 +0000 Subject: [PATCH 02/17] recipe: add AutoQuantize recipe schema and loader support Add the auto_quantize recipe type: AutoQuantizeConfig (candidate_formats, constraints, auto_quantize_method, num_score_steps, disabled_layers, kv_cache), AutoQuantizeConstraints (effective_bits, cost_model, cost) mirroring the mtq.auto_quantize constraints dict, and AutoQuantizeCost (active_moe_expert_ratio). Register RecipeType.AUTO_QUANTIZE in RECIPE_TYPE_TO_CLASS and the loader required-section map, and fix kind-extraction so multi-word non-speculative names stay intact (AUTO_QUANTIZE, not QUANTIZE). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- modelopt/recipe/config.py | 102 +++++++++++++++++++++++++- modelopt/recipe/loader.py | 7 +- tests/unit/recipe/test_loader.py | 118 +++++++++++++++++++++++++++++++ 3 files changed, 223 insertions(+), 4 deletions(-) diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index ea72efdc7c7..bc32e4de53f 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -19,8 +19,9 @@ 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.quantization.config import QuantizeConfig # noqa: TC001 @@ -33,6 +34,10 @@ __all__ = [ "RECIPE_TYPE_TO_CLASS", + "AutoQuantizeConfig", + "AutoQuantizeConstraints", + "AutoQuantizeCost", + "ModelOptAutoQuantizeRecipe", "ModelOptDFlashRecipe", "ModelOptEagleRecipe", "ModelOptMedusaRecipe", @@ -48,6 +53,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 +122,99 @@ class ModelOptPTQRecipe(ModelOptRecipeBase): ) +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.", + ) + + +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 2 required.", + ) + 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).", + ) + num_score_steps: int = ModeloptField( + default=128, + title="Scoring sample count", + description="Number of batches used for sensitivity scoring.", + ) + disabled_layers: list[str] = ModeloptField( + default=[], + title="Excluded layer patterns", + description="Glob patterns; matching layers are excluded from the search.", + ) + 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_two_candidates(cls, v: list[QuantizeConfig]) -> list[QuantizeConfig]: + if len(v) < 2: + raise ValueError( + "auto_quantize requires at least 2 candidate_formats. " + "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 +314,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/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index f4c27f74b2a..ecde540de0d 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,120 @@ 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.num_score_steps == 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_too_few_candidates_raises(tmp_path): + """candidate_formats with fewer than 2 entries is rejected.""" + 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 - algorithm: max\n quant_cfg: []\n" + ) + with pytest.raises(ValueError, match="at least 2"): + load_recipe(bad) + + +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 + # Inline effective_bits overrides: fp8 = 8, w4a16_nvfp4 = 4.5. + assert {c.effective_bits for c in aq.candidate_formats} == {8.0, 4.5} + + +def test_load_recipe_autoquantize_builtin_active_moe_heuristic(): + """The heuristic equivalence-test recipe loads with no effective_bits overrides.""" + aq = load_recipe( + "general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic" + ).auto_quantize + assert aq.constraints.cost_model == "active_moe" + assert all(c.effective_bits is None for c in aq.candidate_formats) From 698569bf7fc479bb3efa71229fd8b1c08f79e3a6 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 01:55:31 +0000 Subject: [PATCH 03/17] examples/llm_ptq: add recipe-driven auto_quantize path, example recipes, and equivalence tests Add auto_quantize_recipe (organized around AutoQuantizeConfig) and _mtq_inputs_from_auto_quantize_config, which maps a recipe to mtq.auto_quantize inputs mirroring the CLI defaults; recipe candidates that match a known preset are passed as the preset dict (_canonical_candidate_dict) so the search names them identically to the CLI and checkpoints stay compatible. The existing CLI auto_quantize helper is left untouched as the equivalence baseline; shared-flow edits are additive and inert when no recipe is used. Ship the active_moe example recipe plus a -heuristic variant for the CLI-equivalence smoke. Add GPU-free tests: per-config recipe-vs-CLI input equivalence and a flag-coverage guard. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/hf_ptq.py | 185 +++++++++++++++++- ...4_fp8_at_6p0bits-active_moe-heuristic.yaml | 45 +++++ ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 48 +++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 182 +++++++++++++++++ 4 files changed, 454 insertions(+), 6 deletions(-) create mode 100644 modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml create mode 100644 modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 959316233fb..0d134893646 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -58,7 +58,7 @@ 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 import ModelOptAutoQuantizeRecipe, ModelOptPTQRecipe, load_recipe from modelopt.recipe.presets import ( KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, @@ -230,6 +230,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 @@ -295,7 +296,7 @@ def make_calib_dataloader( # Labels are only needed for gradient-based auto_quantize include_labels = ( args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient" - ) + ) or autoquant_gradient_recipe calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, @@ -440,6 +441,155 @@ def forward_step(model, batch): return language_model +def _canonical_candidate_dict(fmt) -> dict: + """Return a candidate as a known preset dict when it matches one, else its full dump. + + Mirrors the CLI (which passes ``QUANT_CFG_CHOICES[name]`` directly): matching the preset + makes the search name the candidate after the preset (e.g. FP8_DEFAULT_CFG) instead of + CUSTOM_N, keeping format identity consistent with CLI-produced auto_quantize checkpoints. + """ + stripped = fmt.model_dump(exclude_unset=True) + for preset in QUANT_CFG_CHOICES.values(): + if preset == stripped: + return preset + return fmt.model_dump() + + +def _mtq_inputs_from_auto_quantize_config( + aq_config, args: argparse.Namespace, search_model: torch.nn.Module +) -> dict: + """Map a resolved AutoQuantizeConfig to mtq.auto_quantize inputs. + + Single, testable place where a recipe maps to mtq inputs; mirrors the CLI defaults + (model-derived disabled layers, cost exclusions, KV fallback, preset candidate identity) + so the recipe path stays equivalent to the CLI path. + """ + constraints = aq_config.constraints.model_dump(exclude_none=True) + excluded = _get_auto_quantize_cost_excluded_patterns(search_model) + if excluded: + constraints.setdefault("cost", {})[EXCLUDED_MODULE_NAME_PATTERNS_KEY] = excluded + 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]) + return { + "constraints": constraints, + "quantization_formats": [ + _canonical_candidate_dict(fmt) for fmt in aq_config.candidate_formats + ], + "disabled_layers": aq_config.disabled_layers + or _get_auto_quantize_disabled_layers(search_model), + "kv_cache_quant_cfg": kv_cache_quant_cfg, + "method": aq_config.auto_quantize_method, + "num_score_steps": aq_config.num_score_steps, + } + + +def auto_quantize_recipe( + args: argparse.Namespace, + language_model: torch.nn.Module, + calib_dataloader: DataLoader, + aq_config, + full_model: torch.nn.Module | None = None, +): + """Recipe-driven auto_quantize, organized around an AutoQuantizeConfig. + + Forward-looking (recipe-only) entry point. The CLI ``auto_quantize`` helper is left + untouched as the equivalence baseline and will be retired once the recipe path is verified. + """ + 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 args.inference_pipeline_parallel <= 1, ( + "Auto Quantization is not supported for pipeline parallel size > 1" + ) + + inputs = _mtq_inputs_from_auto_quantize_config(aq_config, args, full_model or language_model) + + # 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 + + def loss_func(output, data): + logits = lm_head(output.last_hidden_state) + labels = data["labels"] + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + 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 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_) + + 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_) + if is_base_model: + assert full_model is not None + return full_model.lm_head(output.last_hidden_state) + return output.logits + + else: + raise ValueError( + f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'" + ) + + language_model, _ = mtq.auto_quantize( + language_model, + constraints=inputs["constraints"], + data_loader=calib_dataloader, + forward_step=forward_step, + loss_func=loss_func, + quantization_formats=inputs["quantization_formats"], + num_calib_steps=len(calib_dataloader), + num_score_steps=min( + len(calib_dataloader), max(inputs["num_score_steps"] // args.batch_size, 1) + ), + verbose=True, + 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) + 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_entries], + ): + mtq.calibrate(language_model, algorithm="max", forward_loop=calibrate_loop) + return language_model + + def load_model(args: argparse.Namespace): # If low memory mode is enabled, we compress the model while loading the HF checkpoint. calibration_only = False @@ -1005,9 +1155,10 @@ def quantize_main( if args.recipe is not None and not args.auto_quantize_bits: 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}" ) def _is_layerwise(obj): @@ -1059,7 +1210,9 @@ def _is_layerwise(obj): else: sample_input_single_batch = None - run_auto_quant = args.auto_quantize_bits is not None + run_auto_quant = args.auto_quantize_bits is not None or isinstance( + recipe, ModelOptAutoQuantizeRecipe + ) args.batch_size = get_max_batch_size( language_model, @@ -1073,7 +1226,16 @@ 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=( + isinstance(recipe, ModelOptAutoQuantizeRecipe) + and recipe.auto_quantize.auto_quantize_method == "gradient" + ), ) # Detect if this is a Nemotron VL model using architecture-based detection @@ -1105,6 +1267,17 @@ def _is_layerwise(obj): full_model=full_model, ) + elif isinstance(recipe, ModelOptAutoQuantizeRecipe): + # Recipe-driven auto_quantize (forward-looking path; the CLI branch above stays the + # untouched equivalence baseline). + auto_quantize_recipe( + args, + full_model, + calib_dataloader, + recipe.auto_quantize, + full_model=full_model, + ) + else: # mono quantization diff --git a/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml b/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml new file mode 100644 index 00000000000..561e27d441f --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml @@ -0,0 +1,45 @@ +# 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. + +# Equivalence-test variant: identical to w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml but +# WITHOUT per-candidate effective_bits, so the LP uses the num_bits heuristic (NVFP4 = 4.0, +# FP8 = 8) — matching the bare CLI command, which has no effective_bits. Use this to verify +# the recipe path is byte-identical to the CLI; use the non-heuristic recipe for corrected cost. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + fp8: configs/ptq/presets/model/fp8 + w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4 + +metadata: + recipe_type: auto_quantize + description: >- + Equivalence baseline: mixed FP8 + NVFP4-weight-only at 6.0 effective bits, active-MoE + cost model (expert ratio 0.03125), num_bits heuristic (no effective_bits override). + +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 + num_score_steps: 128 + # kv_cache omitted -> falls back to --kv_cache_qformat (none in the reference command). 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..d1dd2e49c03 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml @@ -0,0 +1,48 @@ +# 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: + 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: + # effective_bits overrides the LP cost (num_bits heuristic): FP8 = 8, NVFP4 weight = 4.5. + - $import: fp8 + effective_bits: 8 + - $import: w4a16_nvfp4 + effective_bits: 4.5 + + auto_quantize_method: gradient + num_score_steps: 128 + # kv_cache omitted -> falls back to --kv_cache_qformat (none in the reference command). diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index b06c1357411..6bef16dee3d 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -106,3 +106,185 @@ def test_qwen_autoquant_disabled_layers_are_scoped_to_qwen_models(monkeypatch): assert qwen_only_patterns <= qwen_disabled_layers assert qwen_only_patterns.isdisjoint(llama_disabled_layers) + + +def test_autoquant_recipe_builds_canonical_mtq_inputs(monkeypatch): + """Recipe input-building matches the CLI defaults it must stay equivalent to.""" + from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints + from modelopt.recipe.presets import QUANT_CFG_CHOICES + from modelopt.torch.quantization.config import QuantizeConfig + + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--kv_cache_qformat", + "none", + ) + # Isolate the model-derived pieces so the test targets the recipe input-building. + monkeypatch.setattr(hf_ptq, "_get_auto_quantize_disabled_layers", lambda m: ["*lm_head*"]) + monkeypatch.setattr(hf_ptq, "_get_auto_quantize_cost_excluded_patterns", lambda m: []) + fake_model = SimpleNamespace() + + aq_config = AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=6.0), + candidate_formats=[ + QuantizeConfig(**QUANT_CFG_CHOICES["nvfp4"]), + QuantizeConfig(**QUANT_CFG_CHOICES["fp8"]), + ], + ) + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq_config, args, fake_model) + + assert inputs["constraints"] == {"effective_bits": 6.0, "cost_model": "weight"} + assert inputs["disabled_layers"] == ["*lm_head*"] + assert inputs["kv_cache_quant_cfg"] is None + assert inputs["method"] == "gradient" + assert inputs["num_score_steps"] == 128 + # Candidates resolve to the exact preset dicts the CLI feeds mtq, so the search names + # them identically (FP8_DEFAULT_CFG / NVFP4_DEFAULT_CFG) and checkpoints stay compatible. + assert inputs["quantization_formats"][0] == QUANT_CFG_CHOICES["nvfp4"] + assert inputs["quantization_formats"][1] == QUANT_CFG_CHOICES["fp8"] + + +def _recipe_config_from_cli_args(args): + """Build the AutoQuantizeConfig a user would write to mirror the given CLI args.""" + from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints, AutoQuantizeCost + from modelopt.recipe.presets import QUANT_CFG_CHOICES + from modelopt.torch.quantization.config import QuantizeConfig + + cost = None + if args.auto_quantize_active_moe_expert_ratio is not None: + cost = AutoQuantizeCost(active_moe_expert_ratio=args.auto_quantize_active_moe_expert_ratio) + 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[f]) for f in args.qformat.split(",")], + auto_quantize_method=args.auto_quantize_method, + num_score_steps=args.auto_quantize_score_size, + # kv_cache omitted -> recipe path falls back to --kv_cache_qformat, like the CLI. + ) + + +def _cli_expected_mtq_inputs(hf_ptq, args, model): + """Reconstruct the mtq.auto_quantize inputs the CLI helper builds from args. + + Uses the same building blocks the CLI helper uses (QUANT_CFG_CHOICES, the disabled/excluded + helpers, KV presets), so it is the reference the recipe path must match field-for-field. + """ + import copy + + from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES + from modelopt.torch.quantization._auto_quantize_cost import EXCLUDED_MODULE_NAME_PATTERNS_KEY + + constraints = { + "effective_bits": args.auto_quantize_bits, + "cost_model": args.auto_quantize_cost_model, + } + cost = {} + if args.auto_quantize_active_moe_expert_ratio is not None: + cost["active_moe_expert_ratio"] = args.auto_quantize_active_moe_expert_ratio + excluded = hf_ptq._get_auto_quantize_cost_excluded_patterns(model) + if excluded: + cost[EXCLUDED_MODULE_NAME_PATTERNS_KEY] = excluded + if cost: + constraints["cost"] = cost + + if args.kv_cache_qformat == KV_CACHE_NONE: + kv = None + else: + kv = copy.deepcopy(KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]) + + return { + "constraints": constraints, + "quantization_formats": [QUANT_CFG_CHOICES[f] for f in args.qformat.split(",")], + "disabled_layers": hf_ptq._get_auto_quantize_disabled_layers(model), + "kv_cache_quant_cfg": kv, + "method": args.auto_quantize_method, + "num_score_steps": args.auto_quantize_score_size, + } + + +@pytest.mark.parametrize( + "cli_flags", + [ + ["--qformat", "fp8,nvfp4", "--auto_quantize_bits", "6.0"], + [ + "--qformat", + "fp8,w4a16_nvfp4", + "--auto_quantize_bits", + "6.0", + "--auto_quantize_cost_model", + "active_moe", + "--auto_quantize_active_moe_expert_ratio", + "0.03125", + ], + ["--qformat", "fp8,nvfp4", "--auto_quantize_bits", "4.8", "--kv_cache_qformat", "fp8"], + [ + "--qformat", + "fp8,nvfp4", + "--auto_quantize_bits", + "5.0", + "--auto_quantize_method", + "kl_div", + ], + ], +) +def test_recipe_inputs_match_cli_inputs(monkeypatch, cli_flags): + """Across the supported matrix, the recipe path feeds mtq the same inputs as the CLI.""" + hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "dummy", *cli_flags) + monkeypatch.setattr(hf_ptq, "_get_auto_quantize_disabled_layers", lambda m: ["*lm_head*"]) + monkeypatch.setattr(hf_ptq, "_get_auto_quantize_cost_excluded_patterns", lambda m: []) + model = SimpleNamespace() + + recipe_inputs = hf_ptq._mtq_inputs_from_auto_quantize_config( + _recipe_config_from_cli_args(args), args, model + ) + cli_inputs = _cli_expected_mtq_inputs(hf_ptq, args, model) + assert recipe_inputs == cli_inputs + + +def test_autoquant_cli_flags_have_recipe_mapping(monkeypatch): + """Every autoquant spec CLI flag maps to a recipe field (or is intentionally runtime-only). + + Introspects the parsed args, so a newly added ``--auto_quantize_*`` flag that isn't mapped + fails here — flagging that the recipe schema/dispatch needs updating. + """ + from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints, AutoQuantizeCost + + _, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--qformat", + "fp8,nvfp4", + "--auto_quantize_bits", + "6.0", + ) + spec_flags = {k for k in vars(args) if k.startswith("auto_quantize_")} | { + "qformat", + "kv_cache_qformat", + } + + aq_fields = set(AutoQuantizeConfig.model_fields) + constraint_fields = set(AutoQuantizeConstraints.model_fields) + cost_fields = set(AutoQuantizeCost.model_fields) + + # CLI flag (args dest) -> True if covered by the recipe schema (or runtime-only by design). + covered = { + "auto_quantize_bits": "effective_bits" in constraint_fields, + "auto_quantize_method": "auto_quantize_method" in aq_fields, + "auto_quantize_score_size": "num_score_steps" in aq_fields, + "auto_quantize_cost_model": "cost_model" in constraint_fields, + "auto_quantize_active_moe_expert_ratio": "active_moe_expert_ratio" in cost_fields, + "auto_quantize_checkpoint": True, # runtime filesystem path, intentionally CLI-only + "qformat": "candidate_formats" in aq_fields, + "kv_cache_qformat": "kv_cache" in aq_fields, + } + unmapped = spec_flags - set(covered) + assert not unmapped, ( + f"Unmapped autoquant CLI flags (add to recipe schema + mapping): {unmapped}" + ) + assert all(covered.values()), f"A mapped recipe field is missing from the schema: {covered}" From 75901a51f68cfc87ee1bd7c35aaa1535a92bfda0 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 18:33:55 +0000 Subject: [PATCH 04/17] tests: pin cost_weight x effective_bits composition in get_cost Verify the two autoquant cost multipliers stack multiplicatively: a routed NVFP4 expert in active-MoE mode (cost_weight=0.03125) with an effective_bits=4.5 override costs numel * cost_weight * (4.5/16), and falls back to the num_bits heuristic (0.25) without the override. Guards the Phase-A effective_bits / PR-#1497 cost_weight interaction against future cost-model changes. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- .../unit/torch/quantization/test_autoquant.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index deb8f8dcf5e..c1690f616f6 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -173,6 +173,28 @@ 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 (recipe override) stack multiplicatively.""" + model_test = mtq.quantize(torch.nn.Linear(4, 16), mtq.NVFP4_DEFAULT_CFG) + numel = model_test.weight.numel() + + # NVFP4 routed expert in active-MoE mode: cost_weight = expert ratio, effective_bits = 4.5. + override = QuantRecipe({**mtq.NVFP4_DEFAULT_CFG, "effective_bits": 4.5}, name="NVFP4_4P5") + hparam = QuantRecipeHparam( + [override], + quant_modules=[model_test], + quant_module_names=["layers.0.mlp.experts.0.down_proj"], + cost_weight=0.03125, + ) + + # Both factors apply: numel * cost_weight * (effective_bits / 16). + assert hparam.get_cost(override) == pytest.approx(numel * 0.03125 * (4.5 / 16)) + # Without the override the same recipe would use the num_bits heuristic (4.0 / 16 = 0.25). + assert hparam.get_cost(QuantRecipe(mtq.NVFP4_DEFAULT_CFG)) == pytest.approx( + numel * 0.03125 * 0.25 + ) + + 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*"]} From ddd2f273d0b8f9387ac7704110882c55bdd9d17e Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 18:49:07 +0000 Subject: [PATCH 05/17] recipe: set NVFP4 effective_bits=4.5 library default (Phase D, nvfp4 only) Add effective_bits: 4.5 to configs/numerics/nvfp4.yaml so every NVFP4 weight/input/KV entry carries the block-scale-accurate cost (4 value bits + an FP8 scale per 16-element block) as the library default. Recipes and the CLI inherit it via $import, so estimate_quant_compression returns 0.28125 for NVFP4 configs instead of the 4.0/16=0.25 num_bits heuristic. Read only by autoquant; other quantization paths ignore effective_bits. Cost-estimation tests updated to the new baseline. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- modelopt_recipes/configs/numerics/nvfp4.yaml | 3 ++ .../unit/torch/quantization/test_autoquant.py | 47 ++++++++++--------- 2 files changed, 27 insertions(+), 23 deletions(-) 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/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index c1690f616f6..b85feb32649 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -174,25 +174,24 @@ def test_quant_recipe_hparam_zero_cost_weight(): def test_quant_recipe_hparam_cost_weight_and_effective_bits_compose(): - """cost_weight (active_moe) and effective_bits (recipe override) stack multiplicatively.""" + """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() - - # NVFP4 routed expert in active-MoE mode: cost_weight = expert ratio, effective_bits = 4.5. - override = QuantRecipe({**mtq.NVFP4_DEFAULT_CFG, "effective_bits": 4.5}, name="NVFP4_4P5") hparam = QuantRecipeHparam( - [override], + [QuantRecipe(mtq.NVFP4_DEFAULT_CFG)], quant_modules=[model_test], quant_module_names=["layers.0.mlp.experts.0.down_proj"], cost_weight=0.03125, ) - # Both factors apply: numel * cost_weight * (effective_bits / 16). - assert hparam.get_cost(override) == pytest.approx(numel * 0.03125 * (4.5 / 16)) - # Without the override the same recipe would use the num_bits heuristic (4.0 / 16 = 0.25). + # 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 * 0.25 + 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(): @@ -529,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 @@ -599,14 +599,15 @@ def test_estimate_quant_compression(): def test_estimate_quant_compression_effective_bits_override(): - """Recipe-level ``QuantizeConfig.effective_bits`` overrides the num_bits heuristic; unset falls back to it.""" - # NVFP4 — heuristic returns 4.0 bits / 16 = 0.25, but true effective bits is 4.5. + """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 - assert estimate_quant_compression(nvfp4_cfg) == 0.25 # heuristic baseline + 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 - nvfp4_cfg_overridden = mtq.config.QuantizeConfig(**mtq.NVFP4_DEFAULT_CFG, effective_bits=4.5) - assert estimate_quant_compression(nvfp4_cfg_overridden) == 4.5 / 16.0 + # 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) From 25ad6dd706cfd56f95de6b148b6b90f4f3e7f0fe Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 19:07:43 +0000 Subject: [PATCH 06/17] recipe: add Qwen3.6 model-specific autoquant recipe with arch disabled-layers Ship a model-specific autoquant recipe under huggingface/qwen3_6_moe/auto_quantize/ that carries the architecture disabled-layer patterns explicitly in disabled_layers, mirroring the PTQ recipe directory structure (per Wei-Ming, PR #1381). The CLI introspection (_get_auto_quantize_disabled_layers) is kept intact as the equivalence baseline; full removal pairs with the CLI-flag deprecation. Tests: an exact-match guard that the recipe's disabled_layers set equals the CLI introspection for a Qwen model (drift detector), plus an input-equivalence case for a recipe with explicit disabled_layers. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 62 +++++++++++++++++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 48 ++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml 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..abdc31f3fa3 --- /dev/null +++ b/modelopt_recipes/huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml @@ -0,0 +1,62 @@ +# 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: + 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 + num_score_steps: 128 + + # Architecture-specific exclusions (base non-quantizable patterns + Qwen MoE gates). + disabled_layers: + - "*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*" + - "*shared_expert_gate*" diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 6bef16dee3d..702b2b26e1e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -288,3 +288,51 @@ def test_autoquant_cli_flags_have_recipe_mapping(monkeypatch): f"Unmapped autoquant CLI flags (add to recipe schema + mapping): {unmapped}" ) assert all(covered.values()), f"A mapped recipe field is missing from the schema: {covered}" + + +def test_qwen36_recipe_disabled_layers_match_cli_introspection(monkeypatch): + """The Qwen3.6 model recipe's disabled_layers equal the CLI's introspected set. + + F-equivalence guard: the recipe carries arch-disabled patterns explicitly, and they must + match example_utils._get_auto_quantize_disabled_layers for a Qwen model. If the introspection + changes without updating the recipe, this fails (drift detector). Compared as sets since + disabled_layers is order-independent (fnmatch membership). + """ + from modelopt.recipe import load_recipe + + example_utils = _import_example_utils(monkeypatch) + monkeypatch.setattr(example_utils, "is_multimodal_model", lambda model: False) + qwen_model = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_moe")) + + recipe = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe" + ) + assert set(recipe.auto_quantize.disabled_layers) == set( + example_utils._get_auto_quantize_disabled_layers(qwen_model) + ) + + +def test_recipe_with_explicit_disabled_layers_matches_cli(monkeypatch): + """A recipe that sets disabled_layers explicitly feeds mtq the same inputs as the CLI.""" + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, + "--pyt_ckpt_path", + "dummy", + "--qformat", + "fp8,w4a16_nvfp4", + "--auto_quantize_bits", + "6.0", + ) + introspected = ["*shared_expert_gate*", "*mlp.gate.*", "*router*"] + monkeypatch.setattr(hf_ptq, "_get_auto_quantize_disabled_layers", lambda m: list(introspected)) + monkeypatch.setattr(hf_ptq, "_get_auto_quantize_cost_excluded_patterns", lambda m: []) + model = SimpleNamespace() + + # Recipe carries the SAME disabled set explicitly (replace semantics). + recipe_cfg = _recipe_config_from_cli_args(args).model_copy( + update={"disabled_layers": list(introspected)} + ) + recipe_inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(recipe_cfg, args, model) + cli_inputs = _cli_expected_mtq_inputs(hf_ptq, args, model) + assert recipe_inputs["disabled_layers"] == cli_inputs["disabled_layers"] + assert recipe_inputs == cli_inputs From afdfb03d19c6d53432efbddd794992884744096a Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 21:18:52 +0000 Subject: [PATCH 07/17] recipe: add general AutoQuantize example recipes; drop redundant heuristic variant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ship general example recipes (per review): NVFP4+FP8 @ 4.8, NVFP4-W4A4-MSE+FP8 @ 6.0, W4A8-AWQ-beta+FP8 @ 6.0. Remove the now-redundant inline effective_bits from the active_moe recipe (NVFP4 cost 4.5 comes from configs/numerics/nvfp4 after Phase D), and drop the -heuristic variant — post-D it is identical to the cleaned recipe and its name was misleading. Loader test now parametrizes over all shipped general recipes. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- .../auto_quantize/nvfp4_fp8_at_4p8bits.yaml | 36 +++++++++++++++++++ ...tic.yaml => nvfp4_mse_fp8_at_6p0bits.yaml} | 17 +++------ ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 5 ++- .../w4a8_awq_beta_fp8_at_6p0bits.yaml | 36 +++++++++++++++++++ tests/unit/recipe/test_loader.py | 26 +++++++++----- 5 files changed, 95 insertions(+), 25 deletions(-) create mode 100644 modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml rename modelopt_recipes/general/auto_quantize/{w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml => nvfp4_mse_fp8_at_6p0bits.yaml} (54%) create mode 100644 modelopt_recipes/general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits.yaml diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml new file mode 100644 index 00000000000..c131a4c3697 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml @@ -0,0 +1,36 @@ +# 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 4.8 effective bits. + +# modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe +imports: + 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 4.8 effective bits. + +auto_quantize: + constraints: + effective_bits: 4.8 + + candidate_formats: + - $import: nvfp4 + - $import: fp8 + + auto_quantize_method: gradient + num_score_steps: 128 diff --git a/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml similarity index 54% rename from modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml rename to modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml index 561e27d441f..073a8d06880 100644 --- a/modelopt_recipes/general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml @@ -13,33 +13,24 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Equivalence-test variant: identical to w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml but -# WITHOUT per-candidate effective_bits, so the LP uses the num_bits heuristic (NVFP4 = 4.0, -# FP8 = 8) — matching the bare CLI command, which has no effective_bits. Use this to verify -# the recipe path is byte-identical to the CLI; use the non-heuristic recipe for corrected cost. +# AutoQuantize: per-layer search over {NVFP4 W4A4 (weight-MSE + FP8 sweep), FP8} at 6.0 bits. # modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe imports: + nvfp4_mse: configs/ptq/presets/model/nvfp4_w4a4_weight_mse_fp8_sweep fp8: configs/ptq/presets/model/fp8 - w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4 metadata: recipe_type: auto_quantize - description: >- - Equivalence baseline: mixed FP8 + NVFP4-weight-only at 6.0 effective bits, active-MoE - cost model (expert ratio 0.03125), num_bits heuristic (no effective_bits override). + description: Mixed NVFP4 (weight-MSE + FP8 sweep) + FP8 per-layer search at 6.0 effective bits. auto_quantize: constraints: effective_bits: 6.0 - cost_model: active_moe - cost: - active_moe_expert_ratio: 0.03125 candidate_formats: + - $import: nvfp4_mse - $import: fp8 - - $import: w4a16_nvfp4 auto_quantize_method: gradient num_score_steps: 128 - # kv_cache omitted -> falls back to --kv_cache_qformat (none in the reference command). 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 index d1dd2e49c03..42edb66b67d 100644 --- 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 @@ -37,11 +37,10 @@ auto_quantize: active_moe_expert_ratio: 0.03125 candidate_formats: - # effective_bits overrides the LP cost (num_bits heuristic): FP8 = 8, NVFP4 weight = 4.5. + # 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 - effective_bits: 8 - $import: w4a16_nvfp4 - effective_bits: 4.5 auto_quantize_method: gradient num_score_steps: 128 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..2eaf433cb09 --- /dev/null +++ b/modelopt_recipes/general/auto_quantize/w4a8_awq_beta_fp8_at_6p0bits.yaml @@ -0,0 +1,36 @@ +# 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: + 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 + num_score_steps: 128 diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index ecde540de0d..a98238b2e4f 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -1796,14 +1796,22 @@ def test_load_recipe_autoquantize_builtin_active_moe(): assert aq.constraints.cost.active_moe_expert_ratio == 0.03125 assert aq.auto_quantize_method == "gradient" assert aq.kv_cache is None - # Inline effective_bits overrides: fp8 = 8, w4a16_nvfp4 = 4.5. - assert {c.effective_bits for c in aq.candidate_formats} == {8.0, 4.5} + # 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) -def test_load_recipe_autoquantize_builtin_active_moe_heuristic(): - """The heuristic equivalence-test recipe loads with no effective_bits overrides.""" - aq = load_recipe( - "general/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe-heuristic" - ).auto_quantize - assert aq.constraints.cost_model == "active_moe" - assert all(c.effective_bits is None for c in aq.candidate_formats) +@pytest.mark.parametrize( + "recipe_path", + [ + "general/auto_quantize/nvfp4_fp8_at_4p8bits", + "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") From f3dfdfd2b6616df172b7c1ef57f001995fcd950f Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 21:42:47 +0000 Subject: [PATCH 08/17] recipe: make general AutoQuantize recipes self-contained with base disabled_layers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the base (model-agnostic) non-quantizable disabled_layers to every general recipe so they no longer depend on the CLI's _get_auto_quantize_disabled_layers introspection fallback — prep for dropping the CLI in the next commit. Arch-specific models use a huggingface//auto_quantize recipe that extends this set (Qwen3.6 already does). Sharing the base list via $import is a follow-up (needs loader support for schema-less list snippets). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- .../auto_quantize/nvfp4_fp8_at_4p8bits.yaml | 19 +++++++++++++++++++ .../nvfp4_mse_fp8_at_6p0bits.yaml | 19 +++++++++++++++++++ ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 18 ++++++++++++++++++ .../w4a8_awq_beta_fp8_at_6p0bits.yaml | 19 +++++++++++++++++++ 4 files changed, 75 insertions(+) diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml index c131a4c3697..a6987106059 100644 --- a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml @@ -34,3 +34,22 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 + + # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under + # huggingface//auto_quantize/ that extends this set. (TODO: share via $import once + # the loader supports schema-less list snippets — mirrors PTQ's disabled-layer units.) + disabled_layers: + - "*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*" 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 index 073a8d06880..848c8e7841c 100644 --- a/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml @@ -34,3 +34,22 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 + + # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under + # huggingface//auto_quantize/ that extends this set. (TODO: share via $import once + # the loader supports schema-less list snippets — mirrors PTQ's disabled-layer units.) + disabled_layers: + - "*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*" 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 index 42edb66b67d..56316d347ce 100644 --- 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 @@ -45,3 +45,21 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 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: + - "*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*" 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 index 2eaf433cb09..bf956dcf848 100644 --- 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 @@ -34,3 +34,22 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 + + # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under + # huggingface//auto_quantize/ that extends this set. (TODO: share via $import once + # the loader supports schema-less list snippets — mirrors PTQ's disabled-layer units.) + disabled_layers: + - "*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*" From 4cd454756d1401a477ad237773d342495fc999fb Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Mon, 29 Jun 2026 23:21:01 +0000 Subject: [PATCH 09/17] examples/llm_ptq: deprecate AutoQuantize CLI flags; recipe-only (Phase G) AutoQuantize is now driven only by an AutoQuantize --recipe. Remove the --auto_quantize_{bits,method,score_size,cost_model,active_moe_expert_ratio} CLI flags + the CLI auto_quantize() helper + the example-script (parser.sh / huggingface_example.sh) plumbing; --auto_quantize_checkpoint stays as a runtime save/restore path. Remove the model-introspection helpers (_get_auto_quantize_disabled_layers / _get_auto_quantize_cost_excluded_patterns) from example_utils; recipes now carry disabled_layers and a new cost.excluded_module_name_patterns on AutoQuantizeCost, so VL models can exclude vision-tower weights from the cost denominator (disabled-from-search and excluded-from-cost are independent roles). General recipes carry the base disabled set; model-specific recipes extend it. Integration tests (test_llm_ptq.py) and the example script switch to --recipe; README + CHANGELOG updated. Verified: recipe path byte-identical pre/post-G via shared-checkpoint smoke on Qwen3.6-VL; 260 unit tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- CHANGELOG.rst | 2 + examples/hf_ptq/README.md | 47 +-- examples/hf_ptq/example_utils.py | 49 --- examples/hf_ptq/hf_ptq.py | 330 ++---------------- .../hf_ptq/scripts/huggingface_example.sh | 25 +- examples/hf_ptq/scripts/parser.sh | 8 +- modelopt/recipe/config.py | 7 + ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 7 + tests/_test_utils/examples/hf_ptq_utils.py | 4 +- tests/_test_utils/examples/run_command.py | 2 +- tests/examples/hf_ptq/test_hf_ptq_args.py | 299 +--------------- tests/examples/hf_ptq/test_llm_ptq.py | 13 +- 12 files changed, 86 insertions(+), 707 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5002918175d..3d3b967d84e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. +- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``) instead of CLI flags. The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags (and their ``scripts/huggingface_example.sh`` / ``parser.sh`` equivalents) are removed; ``--auto_quantize_checkpoint`` remains as a runtime save/restore path. See ``examples/hf_ptq/README.md`` and ``modelopt_recipes/general/auto_quantize/``. **Deprecations** @@ -39,6 +40,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..763a61e1e39 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -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,10 +305,10 @@ 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 4.8 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. @@ -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": 4.8}, # 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,34 @@ 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, and disabled 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/`. + [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 4.8 effective bits). +scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_4p8bits --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`). - -#### AutoQuantize Advanced Options +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`), `num_score_steps`, and +`disabled_layers`. -| 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_4p8bits \ + --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 0d134893646..ad6dd530ce0 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, @@ -59,12 +57,7 @@ import modelopt.torch.quantization as mtq import modelopt.torch.sparsity as mts from modelopt.recipe import ModelOptAutoQuantizeRecipe, ModelOptPTQRecipe, load_recipe -from modelopt.recipe.presets import ( - KV_CACHE_NONE, - KV_QUANT_CFG_CHOICES, - QFORMAT_ALIASES, - QUANT_CFG_CHOICES, -) +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() @@ -294,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" - ) or autoquant_gradient_recipe + include_labels = autoquant_gradient_recipe calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, @@ -310,137 +255,6 @@ def make_calib_dataloader( return calib_dataloader, first_text_speech_dataset -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, - full_model: torch.nn.Module | None = None, -): - """Auto search quantization of multiple formats.""" - - 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), ( - "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. - 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 - - def loss_func(output, data): - logits = lm_head(output.last_hidden_state) - labels = data["labels"] - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - 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": - - 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) - - elif auto_quantize_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) - if is_base_model: - assert full_model is not None - return full_model.lm_head(output.last_hidden_state) - return output.logits - - else: - raise ValueError( - f"Invalid auto_quantize_method: {auto_quantize_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, - 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], - 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) - ), - verbose=True, - disabled_layers=_get_auto_quantize_disabled_layers(language_model), - method=auto_quantize_method, - checkpoint=auto_quantize_checkpoint, - ) - - 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. - with mtq.set_quantizer_by_cfg_context( - language_model, - [{"quantizer_name": "*", "enable": False}, *kv_cache_quant_cfg], - ): - mtq.calibrate(language_model, algorithm="max", forward_loop=calibrate_loop) - return language_model - - def _canonical_candidate_dict(fmt) -> dict: """Return a candidate as a known preset dict when it matches one, else its full dump. @@ -455,19 +269,17 @@ def _canonical_candidate_dict(fmt) -> dict: return fmt.model_dump() -def _mtq_inputs_from_auto_quantize_config( - aq_config, args: argparse.Namespace, search_model: torch.nn.Module -) -> dict: +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; mirrors the CLI defaults - (model-derived disabled layers, cost exclusions, KV fallback, preset candidate identity) - so the recipe path stays equivalent to the CLI path. + 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) - excluded = _get_auto_quantize_cost_excluded_patterns(search_model) - if excluded: - constraints.setdefault("cost", {})[EXCLUDED_MODULE_NAME_PATTERNS_KEY] = excluded + # NOTE: model-derived cost exclusions (formerly VLM patterns such as *visual*, *vision_tower* + # via model introspection) are not applied here. A future VLM AutoQuantize recipe should carry + # them explicitly (e.g. a cost.excluded_module_name_patterns field). 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: @@ -479,8 +291,7 @@ def _mtq_inputs_from_auto_quantize_config( "quantization_formats": [ _canonical_candidate_dict(fmt) for fmt in aq_config.candidate_formats ], - "disabled_layers": aq_config.disabled_layers - or _get_auto_quantize_disabled_layers(search_model), + "disabled_layers": aq_config.disabled_layers, "kv_cache_quant_cfg": kv_cache_quant_cfg, "method": aq_config.auto_quantize_method, "num_score_steps": aq_config.num_score_steps, @@ -508,7 +319,7 @@ def auto_quantize_recipe( "Auto Quantization is not supported for pipeline parallel size > 1" ) - inputs = _mtq_inputs_from_auto_quantize_config(aq_config, args, full_model or language_model) + inputs = _mtq_inputs_from_auto_quantize_config(aq_config, args) # base-model lm_head handling (mirrors the CLI helper) is_base_model = ( @@ -640,7 +451,7 @@ 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: + if is_nemotron_vl_model and not args.calib_with_images: print("Nemotron VL model detected. Enabling image-text calibration by default.") args.calib_with_images = True @@ -695,7 +506,7 @@ def load_model(args: argparse.Namespace): # 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. - if args.recipe is None and args.auto_quantize_bits is None: + if args.recipe is None: extracted_lm, extracted_model_type = extract_and_prepare_language_model_from_vl( full_model ) @@ -1152,7 +963,7 @@ 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, ModelOptAutoQuantizeRecipe)): @@ -1210,9 +1021,7 @@ def _is_layerwise(obj): else: sample_input_single_batch = None - run_auto_quant = args.auto_quantize_bits is not None or isinstance( - recipe, ModelOptAutoQuantizeRecipe - ) + run_auto_quant = isinstance(recipe, ModelOptAutoQuantizeRecipe) args.batch_size = get_max_batch_size( language_model, @@ -1245,31 +1054,10 @@ 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. - 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, - full_model=full_model, - ) - - elif isinstance(recipe, ModelOptAutoQuantizeRecipe): - # Recipe-driven auto_quantize (forward-looking path; the CLI branch above stays the - # untouched equivalence baseline). + if isinstance(recipe, ModelOptAutoQuantizeRecipe): + # Recipe-driven auto_quantize. For VL models the search walks the OUTER CausalLM + # (which carries lm_head and the LM-head forward path); architecture-specific + # exclusions come from the recipe's disabled_layers. auto_quantize_recipe( args, full_model, @@ -1400,10 +1188,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( @@ -1463,15 +1249,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, @@ -1553,57 +1330,13 @@ def parse_args() -> argparse.Namespace: default=None, type=str, ) - 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'" - ), - ) - 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." - ), - ) - 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." - ), - ) - 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." + "(sensitivity scores, costs, etc.). Used with --recipe ." ), ) parser.add_argument( @@ -1639,20 +1372,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).") @@ -1738,10 +1457,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..29608fa8966 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -94,28 +94,9 @@ 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) - echo "Auto-generated auto_quantize checkpoint path: $AUTO_QUANTIZE_CHECKPOINT" -fi - -if [ -n "$AUTO_QUANTIZE_BITS" ]; then +# AutoQuantize is driven by an AutoQuantize --recipe (see modelopt_recipes/general/auto_quantize/). +# Optional checkpoint passthrough for saving/restoring the search state. +if [ -n "$AUTO_QUANTIZE_CHECKPOINT" ]; then PTQ_ARGS+=" --auto_quantize_checkpoint=$AUTO_QUANTIZE_CHECKPOINT " fi diff --git a/examples/hf_ptq/scripts/parser.sh b/examples/hf_ptq/scripts/parser.sh index 06b440e5731..18a2c7d9746 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:,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,8 +72,6 @@ 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_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;; --moe_calib_experts_ratio ) MOE_CALIB_EXPERTS_RATIO="$2"; shift 2;; --cast_mxfp4_to_nvfp4 ) CAST_MXFP4_TO_NVFP4=true; shift;; @@ -158,7 +155,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,8 +171,6 @@ parse_options() { echo "low_memory_mode: $LOW_MEMORY_MODE" echo "calib_dataset: $CALIB_DATASET" echo "calib_seq: $CALIB_SEQ" - echo "auto_quantize_method: $AUTO_QUANTIZE_METHOD" - echo "auto_quantize_score_size: $AUTO_QUANTIZE_SCORE_SIZE" echo "auto_quantize_checkpoint: $AUTO_QUANTIZE_CHECKPOINT" echo "moe_calib_experts_ratio: $MOE_CALIB_EXPERTS_RATIO" echo "cast_mxfp4_to_nvfp4: $CAST_MXFP4_TO_NVFP4" diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index bc32e4de53f..6e21bb2c71a 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -130,6 +130,13 @@ class AutoQuantizeCost(ModeloptBaseConfig): title="Active MoE expert ratio", description="Routed experts active per token, in (0, 1]. Used by the 'active_moe' cost model.", ) + excluded_module_name_patterns: list[str] | None = ModeloptField( + default=None, + title="Cost-excluded module patterns", + description="Module-name glob patterns excluded from the cost denominator (cost_weight 0) so " + "they don't count toward the bit budget — e.g. vision-tower layers in a VL model. These are " + "typically also in disabled_layers (excluded from search); the two roles are independent.", + ) class AutoQuantizeConstraints(ModeloptBaseConfig): 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 index abdc31f3fa3..1972d0d6650 100644 --- 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 @@ -35,6 +35,13 @@ auto_quantize: cost_model: active_moe cost: active_moe_expert_ratio: 0.03125 + # VL model: exclude vision-tower / MTP weights from the cost denominator (cost_weight 0) + # so they don't count toward the bit budget. (Also in disabled_layers below — excluded from + # search — but cost-exclusion is a separate role.) + excluded_module_name_patterns: + - "*visual*" + - "*mtp*" + - "*vision_tower*" candidate_formats: - $import: fp8 diff --git a/tests/_test_utils/examples/hf_ptq_utils.py b/tests/_test_utils/examples/hf_ptq_utils.py index 1742158ae07..3d2ccc9caba 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 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 702b2b26e1e..b1488f234d8 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -16,9 +16,6 @@ import importlib import sys from pathlib import Path -from types import SimpleNamespace - -import pytest _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -28,11 +25,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,293 +38,24 @@ 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", - ], - ) - - 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): - hf_ptq, args = _parse_hf_ptq_args( - monkeypatch, - "--pyt_ckpt_path", - "nemotron-vl", - "--auto_quantize_bits", - "5.0", - ) - fake_model = SimpleNamespace(device="cpu") - fake_tokenizer = SimpleNamespace(padding_side="right", pad_token="") - - 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) - - 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_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) - - 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_recipe_builds_canonical_mtq_inputs(monkeypatch): - """Recipe input-building matches the CLI defaults it must stay equivalent to.""" - from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints +def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): + """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs.""" + from modelopt.recipe import load_recipe from modelopt.recipe.presets import QUANT_CFG_CHOICES - from modelopt.torch.quantization.config import QuantizeConfig hf_ptq, args = _parse_hf_ptq_args( - monkeypatch, - "--pyt_ckpt_path", - "dummy", - "--kv_cache_qformat", - "none", - ) - # Isolate the model-derived pieces so the test targets the recipe input-building. - monkeypatch.setattr(hf_ptq, "_get_auto_quantize_disabled_layers", lambda m: ["*lm_head*"]) - monkeypatch.setattr(hf_ptq, "_get_auto_quantize_cost_excluded_patterns", lambda m: []) - fake_model = SimpleNamespace() - - aq_config = AutoQuantizeConfig( - constraints=AutoQuantizeConstraints(effective_bits=6.0), - candidate_formats=[ - QuantizeConfig(**QUANT_CFG_CHOICES["nvfp4"]), - QuantizeConfig(**QUANT_CFG_CHOICES["fp8"]), - ], + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" ) - inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq_config, args, fake_model) + aq = load_recipe("general/auto_quantize/nvfp4_fp8_at_4p8bits").auto_quantize + inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) - assert inputs["constraints"] == {"effective_bits": 6.0, "cost_model": "weight"} - assert inputs["disabled_layers"] == ["*lm_head*"] + assert inputs["constraints"] == {"effective_bits": 4.8, "cost_model": "weight"} assert inputs["kv_cache_quant_cfg"] is None assert inputs["method"] == "gradient" assert inputs["num_score_steps"] == 128 - # Candidates resolve to the exact preset dicts the CLI feeds mtq, so the search names - # them identically (FP8_DEFAULT_CFG / NVFP4_DEFAULT_CFG) and checkpoints stay compatible. + # 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 _recipe_config_from_cli_args(args): - """Build the AutoQuantizeConfig a user would write to mirror the given CLI args.""" - from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints, AutoQuantizeCost - from modelopt.recipe.presets import QUANT_CFG_CHOICES - from modelopt.torch.quantization.config import QuantizeConfig - - cost = None - if args.auto_quantize_active_moe_expert_ratio is not None: - cost = AutoQuantizeCost(active_moe_expert_ratio=args.auto_quantize_active_moe_expert_ratio) - 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[f]) for f in args.qformat.split(",")], - auto_quantize_method=args.auto_quantize_method, - num_score_steps=args.auto_quantize_score_size, - # kv_cache omitted -> recipe path falls back to --kv_cache_qformat, like the CLI. - ) - - -def _cli_expected_mtq_inputs(hf_ptq, args, model): - """Reconstruct the mtq.auto_quantize inputs the CLI helper builds from args. - - Uses the same building blocks the CLI helper uses (QUANT_CFG_CHOICES, the disabled/excluded - helpers, KV presets), so it is the reference the recipe path must match field-for-field. - """ - import copy - - from modelopt.recipe.presets import KV_CACHE_NONE, KV_QUANT_CFG_CHOICES, QUANT_CFG_CHOICES - from modelopt.torch.quantization._auto_quantize_cost import EXCLUDED_MODULE_NAME_PATTERNS_KEY - - constraints = { - "effective_bits": args.auto_quantize_bits, - "cost_model": args.auto_quantize_cost_model, - } - cost = {} - if args.auto_quantize_active_moe_expert_ratio is not None: - cost["active_moe_expert_ratio"] = args.auto_quantize_active_moe_expert_ratio - excluded = hf_ptq._get_auto_quantize_cost_excluded_patterns(model) - if excluded: - cost[EXCLUDED_MODULE_NAME_PATTERNS_KEY] = excluded - if cost: - constraints["cost"] = cost - - if args.kv_cache_qformat == KV_CACHE_NONE: - kv = None - else: - kv = copy.deepcopy(KV_QUANT_CFG_CHOICES[args.kv_cache_qformat]) - - return { - "constraints": constraints, - "quantization_formats": [QUANT_CFG_CHOICES[f] for f in args.qformat.split(",")], - "disabled_layers": hf_ptq._get_auto_quantize_disabled_layers(model), - "kv_cache_quant_cfg": kv, - "method": args.auto_quantize_method, - "num_score_steps": args.auto_quantize_score_size, - } - - -@pytest.mark.parametrize( - "cli_flags", - [ - ["--qformat", "fp8,nvfp4", "--auto_quantize_bits", "6.0"], - [ - "--qformat", - "fp8,w4a16_nvfp4", - "--auto_quantize_bits", - "6.0", - "--auto_quantize_cost_model", - "active_moe", - "--auto_quantize_active_moe_expert_ratio", - "0.03125", - ], - ["--qformat", "fp8,nvfp4", "--auto_quantize_bits", "4.8", "--kv_cache_qformat", "fp8"], - [ - "--qformat", - "fp8,nvfp4", - "--auto_quantize_bits", - "5.0", - "--auto_quantize_method", - "kl_div", - ], - ], -) -def test_recipe_inputs_match_cli_inputs(monkeypatch, cli_flags): - """Across the supported matrix, the recipe path feeds mtq the same inputs as the CLI.""" - hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", "dummy", *cli_flags) - monkeypatch.setattr(hf_ptq, "_get_auto_quantize_disabled_layers", lambda m: ["*lm_head*"]) - monkeypatch.setattr(hf_ptq, "_get_auto_quantize_cost_excluded_patterns", lambda m: []) - model = SimpleNamespace() - - recipe_inputs = hf_ptq._mtq_inputs_from_auto_quantize_config( - _recipe_config_from_cli_args(args), args, model - ) - cli_inputs = _cli_expected_mtq_inputs(hf_ptq, args, model) - assert recipe_inputs == cli_inputs - - -def test_autoquant_cli_flags_have_recipe_mapping(monkeypatch): - """Every autoquant spec CLI flag maps to a recipe field (or is intentionally runtime-only). - - Introspects the parsed args, so a newly added ``--auto_quantize_*`` flag that isn't mapped - fails here — flagging that the recipe schema/dispatch needs updating. - """ - from modelopt.recipe.config import AutoQuantizeConfig, AutoQuantizeConstraints, AutoQuantizeCost - - _, args = _parse_hf_ptq_args( - monkeypatch, - "--pyt_ckpt_path", - "dummy", - "--qformat", - "fp8,nvfp4", - "--auto_quantize_bits", - "6.0", - ) - spec_flags = {k for k in vars(args) if k.startswith("auto_quantize_")} | { - "qformat", - "kv_cache_qformat", - } - - aq_fields = set(AutoQuantizeConfig.model_fields) - constraint_fields = set(AutoQuantizeConstraints.model_fields) - cost_fields = set(AutoQuantizeCost.model_fields) - - # CLI flag (args dest) -> True if covered by the recipe schema (or runtime-only by design). - covered = { - "auto_quantize_bits": "effective_bits" in constraint_fields, - "auto_quantize_method": "auto_quantize_method" in aq_fields, - "auto_quantize_score_size": "num_score_steps" in aq_fields, - "auto_quantize_cost_model": "cost_model" in constraint_fields, - "auto_quantize_active_moe_expert_ratio": "active_moe_expert_ratio" in cost_fields, - "auto_quantize_checkpoint": True, # runtime filesystem path, intentionally CLI-only - "qformat": "candidate_formats" in aq_fields, - "kv_cache_qformat": "kv_cache" in aq_fields, - } - unmapped = spec_flags - set(covered) - assert not unmapped, ( - f"Unmapped autoquant CLI flags (add to recipe schema + mapping): {unmapped}" - ) - assert all(covered.values()), f"A mapped recipe field is missing from the schema: {covered}" - - -def test_qwen36_recipe_disabled_layers_match_cli_introspection(monkeypatch): - """The Qwen3.6 model recipe's disabled_layers equal the CLI's introspected set. - - F-equivalence guard: the recipe carries arch-disabled patterns explicitly, and they must - match example_utils._get_auto_quantize_disabled_layers for a Qwen model. If the introspection - changes without updating the recipe, this fails (drift detector). Compared as sets since - disabled_layers is order-independent (fnmatch membership). - """ - from modelopt.recipe import load_recipe - - example_utils = _import_example_utils(monkeypatch) - monkeypatch.setattr(example_utils, "is_multimodal_model", lambda model: False) - qwen_model = SimpleNamespace(config=SimpleNamespace(model_type="qwen3_moe")) - - recipe = load_recipe( - "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_at_6p0bits-active_moe" - ) - assert set(recipe.auto_quantize.disabled_layers) == set( - example_utils._get_auto_quantize_disabled_layers(qwen_model) - ) - - -def test_recipe_with_explicit_disabled_layers_matches_cli(monkeypatch): - """A recipe that sets disabled_layers explicitly feeds mtq the same inputs as the CLI.""" - hf_ptq, args = _parse_hf_ptq_args( - monkeypatch, - "--pyt_ckpt_path", - "dummy", - "--qformat", - "fp8,w4a16_nvfp4", - "--auto_quantize_bits", - "6.0", - ) - introspected = ["*shared_expert_gate*", "*mlp.gate.*", "*router*"] - monkeypatch.setattr(hf_ptq, "_get_auto_quantize_disabled_layers", lambda m: list(introspected)) - monkeypatch.setattr(hf_ptq, "_get_auto_quantize_cost_excluded_patterns", lambda m: []) - model = SimpleNamespace() - - # Recipe carries the SAME disabled set explicitly (replace semantics). - recipe_cfg = _recipe_config_from_cli_args(args).model_copy( - update={"disabled_layers": list(introspected)} - ) - recipe_inputs = hf_ptq._mtq_inputs_from_auto_quantize_config(recipe_cfg, args, model) - cli_inputs = _cli_expected_mtq_inputs(hf_ptq, args, model) - assert recipe_inputs["disabled_layers"] == cli_inputs["disabled_layers"] - assert recipe_inputs == cli_inputs diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 7242b2234f9..cbeb51ac979 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_4p8bits", 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_4p8bits", kv_cache_quant="fp8", calib_batch_size=4, - auto_quantize_bits=6.4, ), PTQCommand( - quant="nvfp4,fp8", + recipe="general/auto_quantize/nvfp4_fp8_at_4p8bits", kv_cache_quant="nvfp4", calib_batch_size=4, - auto_quantize_bits=6.4, ), # sm89 PTQCommand(quant="fp8", min_sm=89), From 040a8b597be951542cd5ae47f6c515e5a98d2af4 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Tue, 30 Jun 2026 18:23:31 +0000 Subject: [PATCH 10/17] AutoQuantize recipe schema: hoist cost_excluded_layers, share disabled_layers via $import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two recipe-author-facing readability cleanups (mtq inputs unchanged — recipe path verified byte-identical to the prior reference, version-string metadata aside): - Hoist excluded_module_name_patterns out of constraints.cost up to a top-level cost_excluded_layers, sibling of disabled_layers. The two 'exclusion' lists (search vs cost-budget) now sit at the same level; the dispatch re-merges cost_excluded_layers into the mtq constraints.cost dict. - Factor the shared 14-pattern base disabled_layers list into a reusable unit (configs/auto_quantize/units/base_disabled_layers) spliced via $import, mirroring PTQ's base_disable_all. Needs a named list[str] schema (LayerPatternList) since the modelopt-schema resolver only accepts modelopt.* dotted paths and str/list[str] have no such name (PTQ reused the existing QuantizerCfgListConfig alias). Adds test_autoquant_recipe_cost_excluded_layers_map_into_cost. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/hf_ptq.py | 10 ++++-- modelopt/recipe/config.py | 26 ++++++++------ .../units/base_disabled_layers.yaml | 34 +++++++++++++++++++ .../auto_quantize/nvfp4_fp8_at_4p8bits.yaml | 21 +++--------- .../nvfp4_mse_fp8_at_6p0bits.yaml | 21 +++--------- ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 16 ++------- .../w4a8_awq_beta_fp8_at_6p0bits.yaml | 21 +++--------- ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 33 ++++++------------ tests/examples/hf_ptq/test_hf_ptq_args.py | 24 +++++++++++++ 9 files changed, 106 insertions(+), 100 deletions(-) create mode 100644 modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index ad6dd530ce0..ab6c591068a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -277,9 +277,13 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) - ``--kv_cache_qformat`` when the recipe omits it. """ constraints = aq_config.constraints.model_dump(exclude_none=True) - # NOTE: model-derived cost exclusions (formerly VLM patterns such as *visual*, *vision_tower* - # via model introspection) are not applied here. A future VLM AutoQuantize recipe should carry - # them explicitly (e.g. a cost.excluded_module_name_patterns field). + # 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: diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index 6e21bb2c71a..e7dd8725c4e 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -122,6 +122,12 @@ 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] + + class AutoQuantizeCost(ModeloptBaseConfig): """Cost-model parameters (the ``cost`` sub-dict of ``mtq.auto_quantize`` constraints).""" @@ -130,13 +136,6 @@ class AutoQuantizeCost(ModeloptBaseConfig): title="Active MoE expert ratio", description="Routed experts active per token, in (0, 1]. Used by the 'active_moe' cost model.", ) - excluded_module_name_patterns: list[str] | None = ModeloptField( - default=None, - title="Cost-excluded module patterns", - description="Module-name glob patterns excluded from the cost denominator (cost_weight 0) so " - "they don't count toward the bit budget — e.g. vision-tower layers in a VL model. These are " - "typically also in disabled_layers (excluded from search); the two roles are independent.", - ) class AutoQuantizeConstraints(ModeloptBaseConfig): @@ -188,10 +187,17 @@ class AutoQuantizeConfig(ModeloptBaseConfig): title="Scoring sample count", description="Number of batches used for sensitivity scoring.", ) - disabled_layers: list[str] = ModeloptField( + 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="Excluded layer patterns", - description="Glob patterns; matching layers are excluded from the search.", + 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, 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..65deb997293 --- /dev/null +++ b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml @@ -0,0 +1,34 @@ +# 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*" diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml index a6987106059..387c2e741f7 100644 --- a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml @@ -17,6 +17,7 @@ # 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 @@ -35,21 +36,7 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 - # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under - # huggingface//auto_quantize/ that extends this set. (TODO: share via $import once - # the loader supports schema-less list snippets — mirrors PTQ's disabled-layer units.) + # 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: - - "*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*" + - $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 index 848c8e7841c..f216c401566 100644 --- a/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml @@ -17,6 +17,7 @@ # 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 @@ -35,21 +36,7 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 - # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under - # huggingface//auto_quantize/ that extends this set. (TODO: share via $import once - # the loader supports schema-less list snippets — mirrors PTQ's disabled-layer units.) + # 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: - - "*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*" + - $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 index 56316d347ce..bed5767385e 100644 --- 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 @@ -20,6 +20,7 @@ # 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 @@ -49,17 +50,4 @@ auto_quantize: # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under # huggingface//auto_quantize/ that extends this set. disabled_layers: - - "*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*" + - $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 index bf956dcf848..97f44c3c444 100644 --- 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 @@ -17,6 +17,7 @@ # 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 @@ -35,21 +36,7 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 - # Base (model-agnostic) non-quantizable layers. Arch-specific models use a recipe under - # huggingface//auto_quantize/ that extends this set. (TODO: share via $import once - # the loader supports schema-less list snippets — mirrors PTQ's disabled-layer units.) + # 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: - - "*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*" + - $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 index 1972d0d6650..75ee5618a36 100644 --- 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 @@ -20,6 +20,7 @@ # 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 @@ -35,13 +36,6 @@ auto_quantize: cost_model: active_moe cost: active_moe_expert_ratio: 0.03125 - # VL model: exclude vision-tower / MTP weights from the cost denominator (cost_weight 0) - # so they don't count toward the bit budget. (Also in disabled_layers below — excluded from - # search — but cost-exclusion is a separate role.) - excluded_module_name_patterns: - - "*visual*" - - "*mtp*" - - "*vision_tower*" candidate_formats: - $import: fp8 @@ -50,20 +44,15 @@ auto_quantize: auto_quantize_method: gradient num_score_steps: 128 - # Architecture-specific exclusions (base non-quantizable patterns + Qwen MoE gates). + # Shared base patterns spliced in; this architecture adds the Qwen MoE shared-expert gate. disabled_layers: - - "*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*" + - $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/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index b1488f234d8..66de35601e9 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -59,3 +59,27 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): # 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.""" + from modelopt.recipe import load_recipe + + hf_ptq, args = _parse_hf_ptq_args( + monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" + ) + aq = load_recipe( + "huggingface/qwen3_6_moe/auto_quantize/w4a16_nvfp4_fp8_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"] From b2c8f4379f2ec98c4631485835e062f563b37275 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Tue, 30 Jun 2026 18:31:35 +0000 Subject: [PATCH 11/17] examples/llm_ptq: rename auto_quantize_recipe->auto_quantize + AutoQuantize recipe docs Rename: with the CLI auto_quantize() helper removed in Phase G, the recipe-driven function is the sole AutoQuantize entry point, so the _recipe suffix is redundant. Rename auto_quantize_recipe -> auto_quantize (def + call site) and refresh the now-stale docstring (it still referred to the removed CLI helper as an 'equivalence baseline'). Pure rename, no behavior change; no name clash with the namespaced mtq.auto_quantize. Docs (no behavior change): - The --recipe / --kv_cache_qformat help and README claimed --kv_cache_qformat is ignored and the recipe 'fully defines' the config under --recipe. True for PTQ recipes (KV baked into quant_cfg) but not AutoQuantize recipes, which fall back to --kv_cache_qformat (default fp8_cast) unless they set an explicit kv_cache field. Clarify the recipe-type split in both help strings and the README; note KV cache is a uniform post-step. - Document cost_excluded_layers (cost-budget exclusion, distinct from disabled_layers) and the shared base_disabled_layers $import unit. - Add a migration note: the --auto_quantize_* CLI flags are removed (AutoQuantize is recipe-only) and how each maps to a recipe field (per Asma's review). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/README.md | 24 +++++++++++++++++++----- examples/hf_ptq/hf_ptq.py | 21 ++++++++++++--------- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 763a61e1e39..12e602c6112 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 @@ -352,12 +352,21 @@ 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, and disabled layers — see -[`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in +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: AutoQuantize is now recipe-only. The former `--auto_quantize_bits`, `--auto_quantize_method`, +> `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and `--auto_quantize_active_moe_expert_ratio` +> CLI flags are removed and map to recipe fields: `--auto_quantize_bits` → `constraints.effective_bits`, +> `--auto_quantize_method` → `auto_quantize_method`, `--auto_quantize_score_size` → `num_score_steps`, +> `--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`. `--auto_quantize_checkpoint` is unchanged. Start from a shipped recipe under +> `modelopt_recipes/general/auto_quantize/` and adjust as needed.* + [Script](./scripts/huggingface_example.sh) ```bash @@ -370,8 +379,13 @@ scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/n 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`), `num_score_steps`, and -`disabled_layers`. +`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div`), `num_score_steps`, +`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`). + +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. The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search state to resume an interrupted search (skips re-scoring): diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index ab6c591068a..b9596aa4d28 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -302,7 +302,7 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) - } -def auto_quantize_recipe( +def auto_quantize( args: argparse.Namespace, language_model: torch.nn.Module, calib_dataloader: DataLoader, @@ -311,8 +311,8 @@ def auto_quantize_recipe( ): """Recipe-driven auto_quantize, organized around an AutoQuantizeConfig. - Forward-looking (recipe-only) entry point. The CLI ``auto_quantize`` helper is left - untouched as the equivalence baseline and will be retired once the recipe path is verified. + 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( @@ -1062,7 +1062,7 @@ def _is_layerwise(obj): # Recipe-driven auto_quantize. For VL models the search walks the OUTER CausalLM # (which carries lm_head and the LM-head forward path); architecture-specific # exclusions come from the recipe's disabled_layers. - auto_quantize_recipe( + auto_quantize( args, full_model, calib_dataloader, @@ -1182,9 +1182,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, ) @@ -1263,8 +1265,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( From 5a4e405b0e6a64e0ec90b131403a49dca5708b32 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Tue, 30 Jun 2026 20:41:28 +0000 Subject: [PATCH 12/17] examples/hf_ptq: address review feedback (CodeRabbit) - VL/AutoQuantize control-flow bug (functional): load_model auto-enables image-text calibration for Nemotron-VL models, which auto_quantize() rejects -> AutoQuantize on a Nemotron-VL model raised NotImplementedError unconditionally. Skip the image-calib default when the run is an AutoQuantize recipe (peek via _recipe_is_auto_quantize). - Validate active_moe_expert_ratio in (0, 1] at the schema boundary (field_validator). - candidate_formats: validate_default=True so an omitted/empty list fails the >=2 check at parse time instead of slipping through. - test_hf_ptq_args: move load_recipe / QUANT_CFG_CHOICES imports to module scope. - PTQCommand: enforce exactly one of quant/recipe via __post_init__. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/hf_ptq.py | 15 +++++++++++++-- modelopt/recipe/config.py | 8 ++++++++ tests/_test_utils/examples/hf_ptq_utils.py | 4 ++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 8 +++----- 4 files changed, 28 insertions(+), 7 deletions(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index b9596aa4d28..a3bec51e277 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -405,6 +405,11 @@ def forward_step(model, batch): 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 @@ -454,8 +459,14 @@ 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: + # Default to image-text calibration for VLM models. Skip for AutoQuantize recipes, 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) + ): print("Nemotron VL model detected. Enabling image-text calibration by default.") args.calib_with_images = True diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index e7dd8725c4e..df7a0eb8a50 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -137,6 +137,13 @@ class AutoQuantizeCost(ModeloptBaseConfig): 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.""" @@ -176,6 +183,7 @@ class AutoQuantizeConfig(ModeloptBaseConfig): default=[], title="Candidate quantization formats", description="Per-layer search space; each entry is a full QuantizeConfig. At least 2 required.", + validate_default=True, ) auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField( default="gradient", diff --git a/tests/_test_utils/examples/hf_ptq_utils.py b/tests/_test_utils/examples/hf_ptq_utils.py index 3d2ccc9caba..16c5952ba98 100644 --- a/tests/_test_utils/examples/hf_ptq_utils.py +++ b/tests/_test_utils/examples/hf_ptq_utils.py @@ -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/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 66de35601e9..c1580f82ed3 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -17,6 +17,9 @@ import sys from pathlib import Path +from modelopt.recipe import load_recipe +from modelopt.recipe.presets import QUANT_CFG_CHOICES + _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -40,9 +43,6 @@ def _parse_hf_ptq_args(monkeypatch, *args): def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): """The recipe path maps an AutoQuantizeConfig to the expected mtq.auto_quantize inputs.""" - from modelopt.recipe import load_recipe - from modelopt.recipe.presets import QUANT_CFG_CHOICES - hf_ptq, args = _parse_hf_ptq_args( monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" ) @@ -64,8 +64,6 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): 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.""" - from modelopt.recipe import load_recipe - hf_ptq, args = _parse_hf_ptq_args( monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" ) From 2afa594382ac26a9c5d3ca9b00c84b189c5ca7d4 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Wed, 1 Jul 2026 21:05:47 +0000 Subject: [PATCH 13/17] examples/hf_ptq: address review round 2 (Edwardf0t1 + Asma) - Export-compat guard (Edwardf0t1): re-add _AUTO_QUANTIZE_QFORMATS and fold an export check into the recipe->mtq translation. _canonical_candidate_dict becomes _match_candidate_to_preset (returns preset name + dict); raise on a non-export-safe candidate, warn on a custom (no-preset) one. Fails fast, before the search. (+tests) - num_score_steps -> score_size (Edwardf0t1): the field is a sample count (divided by batch_size to get mtq steps), so name/describe it honestly and match the former --auto_quantize_score_size. Behavior unchanged (the // batch_size math and 128 default are untouched); disambiguates from mtq's batches-based num_score_steps kwarg. - Auto-generate --auto_quantize_checkpoint (Asma): re-add in huggingface_example.sh, now gated on an AutoQuantize recipe instead of the removed --auto_quantize_bits. - Default effective_bits 4.8 -> 5.4 (Asma): FP4 cost is now 4.5, so 4.8 is too aggressive; rename nvfp4_fp8_at_4p8bits -> nvfp4_fp8_at_5p4bits and update refs/docs. - Add a kl_div example recipe (Asma): nvfp4_fp8_kl_div_at_5p4bits (no backprop; e.g. Llama-4), plus a one-line README pointer. - Note the old AutoQuantize CLI remains on the 0.45 branch (README migration + CHANGELOG). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- CHANGELOG.rst | 2 +- examples/hf_ptq/README.md | 22 +++--- examples/hf_ptq/hf_ptq.py | 70 +++++++++++++++---- .../hf_ptq/scripts/huggingface_example.sh | 8 ++- modelopt/recipe/config.py | 5 +- ...4p8bits.yaml => nvfp4_fp8_at_5p4bits.yaml} | 8 +-- .../nvfp4_fp8_kl_div_at_5p4bits.yaml | 44 ++++++++++++ .../nvfp4_mse_fp8_at_6p0bits.yaml | 2 +- ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 2 +- .../w4a8_awq_beta_fp8_at_6p0bits.yaml | 2 +- ...w4a16_nvfp4_fp8_at_6p0bits-active_moe.yaml | 2 +- tests/examples/hf_ptq/test_hf_ptq_args.py | 41 ++++++++++- tests/examples/hf_ptq/test_llm_ptq.py | 6 +- tests/unit/recipe/test_loader.py | 5 +- 14 files changed, 175 insertions(+), 44 deletions(-) rename modelopt_recipes/general/auto_quantize/{nvfp4_fp8_at_4p8bits.yaml => nvfp4_fp8_at_5p4bits.yaml} (91%) create mode 100644 modelopt_recipes/general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3d3b967d84e..460404fcbaa 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,7 +8,7 @@ Changelog - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. -- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``) instead of CLI flags. The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags (and their ``scripts/huggingface_example.sh`` / ``parser.sh`` equivalents) are removed; ``--auto_quantize_checkpoint`` remains as a runtime save/restore path. See ``examples/hf_ptq/README.md`` and ``modelopt_recipes/general/auto_quantize/``. +- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``) instead of CLI flags. The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags (and their ``scripts/huggingface_example.sh`` / ``parser.sh`` equivalents) are removed; ``--auto_quantize_checkpoint`` remains as a runtime save/restore path. The old AutoQuantize CLI remains available on the 0.45 release branch. See ``examples/hf_ptq/README.md`` and ``modelopt_recipes/general/auto_quantize/``. **Deprecations** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 12e602c6112..cb784889a32 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -308,9 +308,9 @@ Megatron-LM framework PTQ and TensorRT-LLM deployment examples are maintained in `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 `effective_bits` target 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 = {"effective_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, @@ -361,29 +361,33 @@ recipes (carrying architecture-specific disabled layers — e.g. VL vision tower > *Migration: AutoQuantize is now recipe-only. The former `--auto_quantize_bits`, `--auto_quantize_method`, > `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and `--auto_quantize_active_moe_expert_ratio` > CLI flags are removed and map to recipe fields: `--auto_quantize_bits` → `constraints.effective_bits`, -> `--auto_quantize_method` → `auto_quantize_method`, `--auto_quantize_score_size` → `num_score_steps`, +> `--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`. `--auto_quantize_checkpoint` is unchanged. Start from a shipped recipe under -> `modelopt_recipes/general/auto_quantize/` and adjust as needed.* +> `modelopt_recipes/general/auto_quantize/` and adjust as needed. The removed AutoQuantize CLI +> remains available on the 0.45 release branch for anyone who needs the old flags.* [Script](./scripts/huggingface_example.sh) ```bash export HF_PATH= # --recipe selects an AutoQuantize recipe; the recipe defines the candidate formats and the -# effective-bits target (here NVFP4 + FP8 at 4.8 effective bits). -scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_4p8bits --calib_batch_size 4 +# 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 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`), `num_score_steps`, +`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`). +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. + 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. @@ -391,7 +395,7 @@ The one runtime flag is `--auto_quantize_checkpoint` — save/restore the search interrupted search (skips re-scoring): ```bash -scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_4p8bits \ +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 ``` diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a3bec51e277..2d01f615d48 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -255,18 +255,45 @@ def make_calib_dataloader( return calib_dataloader, first_text_speech_dataset -def _canonical_candidate_dict(fmt) -> dict: - """Return a candidate as a known preset dict when it matches one, else its full dump. +# 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", + } +) + - Mirrors the CLI (which passes ``QUANT_CFG_CHOICES[name]`` directly): matching the preset - makes the search name the candidate after the preset (e.g. FP8_DEFAULT_CFG) instead of - CUSTOM_N, keeping format identity consistent with CLI-produced auto_quantize checkpoints. +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. """ stripped = fmt.model_dump(exclude_unset=True) - for preset in QUANT_CFG_CHOICES.values(): + for name, preset in QUANT_CFG_CHOICES.items(): if preset == stripped: - return preset - return fmt.model_dump() + return name, preset + return None, fmt.model_dump() def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) -> dict: @@ -290,15 +317,30 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) - 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": [ - _canonical_candidate_dict(fmt) for fmt in aq_config.candidate_formats - ], + "quantization_formats": quantization_formats, "disabled_layers": aq_config.disabled_layers, "kv_cache_quant_cfg": kv_cache_quant_cfg, "method": aq_config.auto_quantize_method, - "num_score_steps": aq_config.num_score_steps, + "score_size": aq_config.score_size, } @@ -378,9 +420,7 @@ def forward_step(model, batch): loss_func=loss_func, quantization_formats=inputs["quantization_formats"], num_calib_steps=len(calib_dataloader), - num_score_steps=min( - len(calib_dataloader), max(inputs["num_score_steps"] // args.batch_size, 1) - ), + num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)), verbose=True, disabled_layers=inputs["disabled_layers"], method=inputs["method"], diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index 29608fa8966..2cc4ce5865e 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -95,7 +95,13 @@ if [ "$LOW_MEMORY_MODE" = "true" ]; then fi # AutoQuantize is driven by an AutoQuantize --recipe (see modelopt_recipes/general/auto_quantize/). -# Optional checkpoint passthrough for saving/restoring the search state. +# For an AutoQuantize recipe, auto-generate a checkpoint path (to save/restore the search state) +# when the user didn't supply one. Detected by the recipe path living under an auto_quantize/ dir. +if [ -z "$AUTO_QUANTIZE_CHECKPOINT" ] && [[ "$RECIPE" == *auto_quantize* ]]; 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 diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index df7a0eb8a50..acc70b8d574 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -190,10 +190,11 @@ class AutoQuantizeConfig(ModeloptBaseConfig): title="Sensitivity scoring method", description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).", ) - num_score_steps: int = ModeloptField( + score_size: int = ModeloptField( default=128, title="Scoring sample count", - description="Number of batches used for sensitivity scoring.", + 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=[], diff --git a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_5p4bits.yaml similarity index 91% rename from modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml rename to modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_5p4bits.yaml index 387c2e741f7..7f4f337100d 100644 --- a/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_4p8bits.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_fp8_at_5p4bits.yaml @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# AutoQuantize: per-layer search over {NVFP4 (W4A4), FP8 (W8A8)} at 4.8 effective bits. +# AutoQuantize: per-layer search over {NVFP4 (W4A4), FP8 (W8A8)} at 5.4 effective bits. # modelopt-schema: modelopt.recipe.config.ModelOptAutoQuantizeRecipe imports: @@ -23,18 +23,18 @@ imports: metadata: recipe_type: auto_quantize - description: Mixed NVFP4 + FP8 per-layer search at 4.8 effective bits. + description: Mixed NVFP4 + FP8 per-layer search at 5.4 effective bits. auto_quantize: constraints: - effective_bits: 4.8 + effective_bits: 5.4 candidate_formats: - $import: nvfp4 - $import: fp8 auto_quantize_method: gradient - num_score_steps: 128 + 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. 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 index f216c401566..ac9546d049d 100644 --- a/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml +++ b/modelopt_recipes/general/auto_quantize/nvfp4_mse_fp8_at_6p0bits.yaml @@ -34,7 +34,7 @@ auto_quantize: - $import: fp8 auto_quantize_method: gradient - num_score_steps: 128 + 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. 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 index bed5767385e..56fc5fd789a 100644 --- 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 @@ -44,7 +44,7 @@ auto_quantize: - $import: w4a16_nvfp4 auto_quantize_method: gradient - num_score_steps: 128 + 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 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 index 97f44c3c444..d0135f52a4d 100644 --- 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 @@ -34,7 +34,7 @@ auto_quantize: - $import: fp8 auto_quantize_method: gradient - num_score_steps: 128 + 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. 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 index 75ee5618a36..201a70614eb 100644 --- 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 @@ -42,7 +42,7 @@ auto_quantize: - $import: w4a16_nvfp4 auto_quantize_method: gradient - num_score_steps: 128 + score_size: 128 # Shared base patterns spliced in; this architecture adds the Qwen MoE shared-expert gate. disabled_layers: diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index c1580f82ed3..2fed74dfc1d 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -17,8 +17,12 @@ import sys from pathlib import Path +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" @@ -46,13 +50,13 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): hf_ptq, args = _parse_hf_ptq_args( monkeypatch, "--pyt_ckpt_path", "dummy", "--kv_cache_qformat", "none" ) - aq = load_recipe("general/auto_quantize/nvfp4_fp8_at_4p8bits").auto_quantize + 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": 4.8, "cost_model": "weight"} + assert inputs["constraints"] == {"effective_bits": 5.4, "cost_model": "weight"} assert inputs["kv_cache_quant_cfg"] is None assert inputs["method"] == "gradient" - assert inputs["num_score_steps"] == 128 + 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"] @@ -81,3 +85,34 @@ def test_autoquant_recipe_cost_excluded_layers_map_into_cost(monkeypatch): # 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"] + + +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) + + +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) diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index cbeb51ac979..4b66bad254e 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -80,7 +80,7 @@ def test_ptq_whisper(command): PTQCommand(quant="nvfp4_awq"), # autoquant (recipe-driven) PTQCommand( - recipe="general/auto_quantize/nvfp4_fp8_at_4p8bits", + recipe="general/auto_quantize/nvfp4_fp8_at_5p4bits", calib_batch_size=4, kv_cache_quant="none", ), @@ -89,12 +89,12 @@ def test_ptq_whisper(command): PTQCommand(quant="fp8", kv_cache_quant="fp8_cast", min_sm=89), # autoquant_kv_cache (recipe-driven; KV via --kv_cache_quant fallback) PTQCommand( - recipe="general/auto_quantize/nvfp4_fp8_at_4p8bits", + recipe="general/auto_quantize/nvfp4_fp8_at_5p4bits", kv_cache_quant="fp8", calib_batch_size=4, ), PTQCommand( - recipe="general/auto_quantize/nvfp4_fp8_at_4p8bits", + recipe="general/auto_quantize/nvfp4_fp8_at_5p4bits", kv_cache_quant="nvfp4", calib_batch_size=4, ), diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index a98238b2e4f..5acb2b1c36e 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -1720,7 +1720,7 @@ def test_load_recipe_autoquantize_minimal(tmp_path): assert isinstance(recipe, ModelOptAutoQuantizeRecipe) aq = recipe.auto_quantize assert aq.auto_quantize_method == "gradient" - assert aq.num_score_steps == 128 + assert aq.score_size == 128 assert aq.kv_cache is None assert aq.constraints.effective_bits == 4.8 assert aq.constraints.cost_model == "weight" @@ -1803,7 +1803,8 @@ def test_load_recipe_autoquantize_builtin_active_moe(): @pytest.mark.parametrize( "recipe_path", [ - "general/auto_quantize/nvfp4_fp8_at_4p8bits", + "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", From 2cf69fadf23d8fb98d0e9efef67213a5593e59e9 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Wed, 1 Jul 2026 22:09:39 +0000 Subject: [PATCH 14/17] examples/hf_ptq: harden export guard against effective_bits bypass (CodeRabbit) _match_candidate_to_preset matched candidates by exact model_dump equality, so a candidate built from a non-export-safe preset that also set a per-candidate effective_bits would fail the match, be classified 'custom', and slip past the export whitelist with only a warning. Exclude effective_bits (cost-only, export-irrelevant) from the match key so such a candidate is still identified as its base preset and rejected; preserve the override in the returned config. Shipped recipes are unaffected (they set no per-candidate effective_bits). (+test) Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/hf_ptq.py | 9 ++++++++- tests/examples/hf_ptq/test_hf_ptq_args.py | 15 +++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 2d01f615d48..a6cf02cd38c 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -288,10 +288,17 @@ def _match_candidate_to_preset(fmt) -> tuple[str | None, dict]: 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 == stripped: + 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() diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 2fed74dfc1d..fffba8554a3 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -116,3 +116,18 @@ def test_autoquant_warns_on_custom_candidate(monkeypatch): ) with pytest.warns(UserWarning, match="export compatibility cannot be verified"): hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) + + +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) From d108b1cef685f50d5faa47c2f67dfcfb1467f393 Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 2 Jul 2026 19:15:06 +0000 Subject: [PATCH 15/17] examples/hf_ptq: re-add AutoQuantize CLI as a deprecated on-the-fly-recipe shim Per review (Keval): keep the --auto_quantize_* flags working instead of hard-removing them. They convert into an AutoQuantizeConfig on the fly and run the same recipe path (DeprecationWarning); no new user flags. - _auto_quantize_config_from_cli(): builds the config from the flags; appends the shared base disabled + base cost-excluded layer sets (no model introspection). Base cost-excluded is appended unconditionally (harmless on non-VL, correct on VL). - Base layer-pattern sets loaded once as module constants in recipe/config.py, mirroring quantization/config.py's _default_disabled_quantizer_cfg (Shengliang). New shared unit configs/auto_quantize/units/base_cost_excluded_layers. - quantize_main resolves aq_config from a recipe OR the CLI flags. - Fix VL guards for the CLI path: skip the image-calib default AND the plain-PTQ extract_and_prepare_language_model_from_vl (else auto_quantize hits 'multiple modelopt states'); reject --low_memory_mode. - parser.sh / huggingface_example.sh: flag passthrough + auto-generated checkpoint path. - CHANGELOG: Backward-Breaking -> Deprecations (flags still work). README reframed. +test. Verified CLI == recipe (byte-identical) on the Qwen3.6 VL MoE. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- CHANGELOG.rst | 3 +- examples/hf_ptq/README.md | 17 +-- examples/hf_ptq/hf_ptq.py | 133 +++++++++++++++--- .../hf_ptq/scripts/huggingface_example.sh | 19 ++- examples/hf_ptq/scripts/parser.sh | 12 +- modelopt/recipe/config.py | 19 +++ .../units/base_cost_excluded_layers.yaml | 25 ++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 35 +++++ 8 files changed, 230 insertions(+), 33 deletions(-) create mode 100644 modelopt_recipes/configs/auto_quantize/units/base_cost_excluded_layers.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 460404fcbaa..030c3d41fd8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,10 +8,11 @@ Changelog - Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. - Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. -- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``) instead of CLI flags. The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags (and their ``scripts/huggingface_example.sh`` / ``parser.sh`` equivalents) are removed; ``--auto_quantize_checkpoint`` remains as a runtime save/restore path. The old AutoQuantize CLI remains available on the 0.45 release branch. See ``examples/hf_ptq/README.md`` and ``modelopt_recipes/general/auto_quantize/``. **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. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index cb784889a32..95888784f98 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -358,15 +358,16 @@ cost-excluded layers — see [`AutoQuantizeConfig`](../../modelopt/recipe/config recipes (carrying architecture-specific disabled layers — e.g. VL vision towers) live under `modelopt_recipes/huggingface//auto_quantize/`. -> *Migration: AutoQuantize is now recipe-only. The former `--auto_quantize_bits`, `--auto_quantize_method`, +> *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 removed and 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`. `--auto_quantize_checkpoint` is unchanged. Start from a shipped recipe under -> `modelopt_recipes/general/auto_quantize/` and adjust as needed. The removed AutoQuantize CLI -> remains available on the 0.45 release branch for anyone who needs the old flags.* +> 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) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index a6cf02cd38c..8dcc78afa27 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -351,6 +351,47 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) - } +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, @@ -506,13 +547,15 @@ def load_model(args: argparse.Namespace): is_nemotron_vl_model = is_nemotron_vl(full_model) - # Default to image-text calibration for VLM models. Skip for AutoQuantize recipes, 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. + # 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 @@ -565,10 +608,11 @@ 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. - if args.recipe is None: + # 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 ) @@ -1034,6 +1078,20 @@ def quantize_main( 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) @@ -1083,7 +1141,7 @@ def _is_layerwise(obj): else: sample_input_single_batch = None - run_auto_quant = isinstance(recipe, ModelOptAutoQuantizeRecipe) + run_auto_quant = aq_config is not None args.batch_size = get_max_batch_size( language_model, @@ -1104,8 +1162,7 @@ def _is_layerwise(obj): device, model_type, autoquant_gradient_recipe=( - isinstance(recipe, ModelOptAutoQuantizeRecipe) - and recipe.auto_quantize.auto_quantize_method == "gradient" + aq_config is not None and aq_config.auto_quantize_method == "gradient" ), ) @@ -1116,15 +1173,15 @@ def _is_layerwise(obj): args, full_model, model_type, tokenizer, calib_dataloader, is_nemotron_vl_model ) - if isinstance(recipe, ModelOptAutoQuantizeRecipe): - # Recipe-driven auto_quantize. For VL models the search walks the OUTER CausalLM - # (which carries lm_head and the LM-head forward path); architecture-specific - # exclusions come from the recipe's disabled_layers. + 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, - recipe.auto_quantize, + aq_config, full_model=full_model, ) @@ -1401,9 +1458,47 @@ def parse_args() -> argparse.Namespace: default=None, help=( "Path to checkpoint file for saving/restoring auto_quantize search state " - "(sensitivity scores, costs, etc.). Used with --recipe ." + "(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="[Deprecated: use an AutoQuantize --recipe] Sensitivity scoring method.", + ) + parser.add_argument( + "--auto_quantize_score_size", + type=int, + default=128, + 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="[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="[Deprecated: use an AutoQuantize --recipe] Routed-expert active ratio for the " + "'active_moe' cost model.", + ) parser.add_argument( "--moe_calib_experts_ratio", type=float, @@ -1448,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 diff --git a/examples/hf_ptq/scripts/huggingface_example.sh b/examples/hf_ptq/scripts/huggingface_example.sh index 2cc4ce5865e..84057e468c9 100755 --- a/examples/hf_ptq/scripts/huggingface_example.sh +++ b/examples/hf_ptq/scripts/huggingface_example.sh @@ -94,10 +94,9 @@ if [ "$LOW_MEMORY_MODE" = "true" ]; then PTQ_ARGS+=" --low_memory_mode " fi -# AutoQuantize is driven by an AutoQuantize --recipe (see modelopt_recipes/general/auto_quantize/). -# For an AutoQuantize recipe, auto-generate a checkpoint path (to save/restore the search state) -# when the user didn't supply one. Detected by the recipe path living under an auto_quantize/ dir. -if [ -z "$AUTO_QUANTIZE_CHECKPOINT" ] && [[ "$RECIPE" == *auto_quantize* ]]; then +# 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" @@ -106,6 +105,18 @@ 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_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 PTQ_ARGS+=" --dataset=$CALIB_DATASET " fi diff --git a/examples/hf_ptq/scripts/parser.sh b/examples/hf_ptq/scripts/parser.sh index 18a2c7d9746..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:,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:,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 @@ -73,6 +73,11 @@ parse_options() { --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_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;; @@ -172,6 +177,11 @@ parse_options() { 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_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 acc70b8d574..e0bece4c3b9 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -24,6 +24,7 @@ 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 @@ -128,6 +129,24 @@ class ModelOptPTQRecipe(ModelOptRecipeBase): 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).""" 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/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index fffba8554a3..7160dd38cb4 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -131,3 +131,38 @@ def test_autoquant_export_guard_not_bypassed_by_effective_bits(monkeypatch): ) with pytest.raises(ValueError, match="not supported for unified checkpoint export"): hf_ptq._mtq_inputs_from_auto_quantize_config(aq, args) + + +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*"] From 69807258094082f20aaf2afe62e12a9385d42bda Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 2 Jul 2026 20:22:59 +0000 Subject: [PATCH 16/17] auto_quantize base: disable *shared_expert_gate* (Qwen3.6 MoE export fusion) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Qwen3.6 MoE (e.g. Qwen/Qwen3.6-35B-A3B) fails HF export at linear fusion if the shared-expert gate is quantized (fusion partners get mismatched formats). On main this was a Qwen-specific introspection pattern (_QWEN36_AUTOQ_DISABLED_LAYERS); promote it to the shared base disabled set so the deprecated --auto_quantize_* CLI (which can't inject arch patterns) also disables it. Harmless elsewhere — matches nothing on non-MoE models. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- .../configs/auto_quantize/units/base_disabled_layers.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml index 65deb997293..fd27046d608 100644 --- a/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml +++ b/modelopt_recipes/configs/auto_quantize/units/base_disabled_layers.yaml @@ -32,3 +32,6 @@ - "*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*" From c53ae9f85a4e4108a3151e8715c42d94cd5b53fa Mon Sep 17 00:00:00 2001 From: Juhi Mittal Date: Thu, 2 Jul 2026 23:08:56 +0000 Subject: [PATCH 17/17] auto_quantize: allow a single candidate_format (one-format + bf16 search) Per review (Wei-Ming): support 'one format + bf16' for AutoQuantize. bf16/no-quant is always an implicit per-layer choice (mtq appends QuantRecipe(quant_cfg=None)), so a single explicit format already yields a real {format, bf16} search. Relax the candidate_formats validator from >=2 to >=1 (only an empty list is rejected). Works for both recipe (candidate_formats: [fp8]) and the CLI shim (--qformat fp8 --auto_quantize_bits ...). Updates the field description + README; retargets the loader test (empty rejected, single accepted). Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Juhi Mittal --- examples/hf_ptq/README.md | 3 +++ modelopt/recipe/config.py | 14 +++++++++----- tests/unit/recipe/test_loader.py | 20 ++++++++++++++++---- 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 95888784f98..3a7380a9f9e 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -386,6 +386,9 @@ keeps the more sensitive ones at higher precision (or unquantized), so the model 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. diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index e0bece4c3b9..2cf5a2f0cfb 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -201,7 +201,9 @@ class AutoQuantizeConfig(ModeloptBaseConfig): candidate_formats: list[QuantizeConfig] = ModeloptField( default=[], title="Candidate quantization formats", - description="Per-layer search space; each entry is a full QuantizeConfig. At least 2 required.", + 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( @@ -236,11 +238,13 @@ class AutoQuantizeConfig(ModeloptBaseConfig): @field_validator("candidate_formats") @classmethod - def _at_least_two_candidates(cls, v: list[QuantizeConfig]) -> list[QuantizeConfig]: - if len(v) < 2: + 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 2 candidate_formats. " - "For uniform quantization, use a PTQ recipe instead." + "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 diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index 5acb2b1c36e..3aaacaa3e0e 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -1766,18 +1766,30 @@ def test_load_recipe_autoquantize_missing_section_raises(tmp_path): load_recipe(bad) -def test_load_recipe_autoquantize_too_few_candidates_raises(tmp_path): - """candidate_formats with fewer than 2 entries is rejected.""" +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 - algorithm: max\n quant_cfg: []\n" + " candidate_formats: []\n" ) - with pytest.raises(ValueError, match="at least 2"): + 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"