diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index 57d9bebef43..9c692e5b7aa 100755 --- a/examples/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -42,6 +42,9 @@ ProcessorMixin, ) +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 except ImportError: @@ -51,6 +54,58 @@ SPECULATIVE_MODEL_LIST = ["Eagle", "Medusa"] +# TODO: Refactor into the config system. +_QWEN36_AUTOQ_DISABLED_LAYERS = ( + "*shared_expert_gate*", + "*linear_attn.in_proj_a*", + "*linear_attn.in_proj_b*", +) +_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, @@ -133,7 +188,6 @@ def is_nemotron_vl(model_or_config): # Try to get config from model, or use directly if it's a config if hasattr(model_or_config, "config"): config = model_or_config.config - from modelopt.torch.export.model_utils import is_multimodal_model if not is_multimodal_model(model_or_config): return False diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 3fe3f3ceb03..f3e7445e8b4 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -27,6 +27,8 @@ 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, build_quant_cfg, copy_custom_model_files, create_vlm_calibration_loop, @@ -72,7 +74,8 @@ save_expert_token_count_table, ) from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model -from modelopt.torch.quantization.config import _default_disabled_quantizer_cfg, need_calibration +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 from modelopt.torch.speculative.eagle.utils import ( @@ -132,6 +135,7 @@ def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool: "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", @@ -387,10 +391,14 @@ def forward_step(model, batch): "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_constraints["cost"] = { - "active_moe_expert_ratio": args.auto_quantize_active_moe_expert_ratio - } + 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, @@ -406,12 +414,7 @@ def forward_step(model, batch): len(calib_dataloader), max(auto_quantize_score_size // args.batch_size, 1) ), verbose=True, - # Disable all default disabled layers such as lm_head, mlp.gate, router etc. - disabled_layers=[ - entry["quantizer_name"] - for entry in _default_disabled_quantizer_cfg - if "parent_class" not in entry - ], + disabled_layers=_get_auto_quantize_disabled_layers(language_model), method=auto_quantize_method, checkpoint=auto_quantize_checkpoint, ) @@ -487,7 +490,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: + if is_nemotron_vl_model and not args.calib_with_images and args.auto_quantize_bits is None: print("Nemotron VL model detected. Enabling image-text calibration by default.") args.calib_with_images = True @@ -539,12 +542,10 @@ def load_model(args: argparse.Namespace): : len(args.dataset) ] - # We only quantize the language model for VLMs other than the type supported above. - # Recipe mode is the exception: in Qwen3.5/3.6-MoE VLMs, lm_head sits - # on the outer CausalLM, not the inner language backbone. A recipe that targets - # lm_head must therefore quantize against the full model and explicitly keep visual - # and MTP siblings disabled. - if args.recipe is None: + # 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: extracted_lm, extracted_model_type = extract_and_prepare_language_model_from_vl( full_model ) @@ -1070,9 +1071,16 @@ def _is_layerwise(obj): "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, - language_model, + full_model, calib_dataloader, auto_quantize_method=args.auto_quantize_method, auto_quantize_score_size=args.auto_quantize_score_size, @@ -1437,6 +1445,8 @@ 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 ): diff --git a/modelopt/torch/quantization/_auto_quantize_cost.py b/modelopt/torch/quantization/_auto_quantize_cost.py index f8f6cb36e7b..68f93770865 100644 --- a/modelopt/torch/quantization/_auto_quantize_cost.py +++ b/modelopt/torch/quantization/_auto_quantize_cost.py @@ -15,6 +15,7 @@ """Cost models for AutoQuantize effective-bits accounting.""" +import fnmatch from collections.abc import Callable, Iterable, Sequence from typing import Any, Final @@ -27,6 +28,7 @@ AUTO_QUANTIZE_CONSTRAINT_KEYS: Final = frozenset({"effective_bits", "cost_model", "cost"}) ACTIVE_MOE_EXPERT_RATIO_KEY: Final = "active_moe_expert_ratio" +EXCLUDED_MODULE_NAME_PATTERNS_KEY: Final = "excluded_module_name_patterns" COST_MODEL_WEIGHT: Final = "weight" COST_MODEL_ACTIVE_MOE: Final = "active_moe" @@ -90,11 +92,31 @@ def is_routed_moe_module_name(name: str) -> bool: return "shared_expert" not in name and _ROUTED_MOE_EXPERT_NAME_RE.search(name) is not None +def _get_module_weight_numel(module: nn.Module) -> int: + """Return the parameter count for a module's quantizable weights. + + Standard quantized linear modules have a single ``weight`` parameter. Fused + MoE expert containers expose projection tensors directly instead, so both + fused projections contribute to AutoQuantize cost accounting. + """ + weight = getattr(module, "weight", None) + if weight is not None: + return weight.numel() + + # Fused MoE expert containers expose projection tensors directly instead of + # a single ``weight`` parameter. + return sum( + param.numel() + for attr in ("gate_up_proj", "down_proj") + if (param := getattr(module, attr, None)) is not None + ) + + class AutoQuantizeCostModel: """Base class for AutoQuantize effective-bits cost accounting.""" name: str - supported_cost_keys: frozenset[str] = frozenset() + supported_cost_keys: frozenset[str] = frozenset({EXCLUDED_MODULE_NAME_PATTERNS_KEY}) def normalize_cost_constraints( self, model: nn.Module, cost_constraints: dict[str, Any] @@ -103,12 +125,35 @@ def normalize_cost_constraints( unknown_cost_keys = set(cost_constraints) - self.supported_cost_keys if unknown_cost_keys: raise ValueError(f"Unsupported auto_quantize cost constraints: {unknown_cost_keys}.") + excluded_patterns = cost_constraints.get(EXCLUDED_MODULE_NAME_PATTERNS_KEY) + if excluded_patterns is None: + return cost_constraints + if isinstance(excluded_patterns, str): + excluded_patterns = [excluded_patterns] + if not isinstance(excluded_patterns, Sequence) or not all( + isinstance(pattern, str) for pattern in excluded_patterns + ): + raise ValueError( + f"constraints['cost']['{EXCLUDED_MODULE_NAME_PATTERNS_KEY}'] must be a string " + "or a sequence of strings." + ) + cost_constraints[EXCLUDED_MODULE_NAME_PATTERNS_KEY] = list(excluded_patterns) return cost_constraints def module_cost_weight( self, module_names: Sequence[str], cost_constraints: dict[str, Any] ) -> float: """Return the cost multiplier for a group of modules.""" + excluded_patterns = cost_constraints.get(EXCLUDED_MODULE_NAME_PATTERNS_KEY, []) + if ( + module_names + and excluded_patterns + and all( + any(fnmatch.fnmatch(name, pattern) for pattern in excluded_patterns) + for name in module_names + ) + ): + return 0.0 return 1.0 def total_weight_size( @@ -119,7 +164,7 @@ def total_weight_size( ) -> float: """Return the cost denominator for the effective-bits constraint.""" return sum( - module.weight.numel() * self.module_cost_weight([name], cost_constraints) + _get_module_weight_numel(module) * self.module_cost_weight([name], cost_constraints) for name, module in named_modules if is_auto_quantize_module(module) ) @@ -135,7 +180,9 @@ class ActiveMoECostModel(AutoQuantizeCostModel): """Scale routed MoE expert weights by the active experts per-token ratio.""" name = COST_MODEL_ACTIVE_MOE - supported_cost_keys = frozenset({ACTIVE_MOE_EXPERT_RATIO_KEY}) + supported_cost_keys = frozenset( + {ACTIVE_MOE_EXPERT_RATIO_KEY, EXCLUDED_MODULE_NAME_PATTERNS_KEY} + ) def normalize_cost_constraints( self, model: nn.Module, cost_constraints: dict[str, Any] @@ -164,9 +211,12 @@ def normalize_cost_constraints( def module_cost_weight( self, module_names: Sequence[str], cost_constraints: dict[str, Any] ) -> float: + base_weight = super().module_cost_weight(module_names, cost_constraints) + if base_weight == 0.0: + return 0.0 if any(is_routed_moe_module_name(n) for n in module_names): return cost_constraints[ACTIVE_MOE_EXPERT_RATIO_KEY] - return 1.0 + return base_weight _COST_MODELS: Final = { diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 5a7c70b54ae..03bd801387c 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -45,6 +45,7 @@ AUTO_QUANTIZE_CONSTRAINT_KEYS, COST_MODEL_ACTIVE_MOE, COST_MODEL_WEIGHT, + _get_module_weight_numel, get_auto_quantize_cost_model, normalize_auto_quantize_constraints, ) @@ -54,6 +55,65 @@ from .utils import is_quantized_linear +def _is_hf_quant_fused_experts_module(module: nn.Module) -> bool: + """Return True for a converted HF fused-MoE-experts quantization wrapper.""" + # Late import avoids a circular import: the HF plugin registers AutoQuantize + # support from this module at import time. + try: + from .plugins.huggingface import _is_quant_fused_experts_module + except ImportError: + return False + return _is_quant_fused_experts_module(module) + + +# Quantizer attribute names that participate in AutoQuantize snapshot/restore. +_STD_QUANTIZER_ATTRS = ("input_quantizer", "weight_quantizer", "output_quantizer") +_FUSED_EXPERTS_QUANTIZER_ATTRS = ( + "gate_up_proj_input_quantizer", + "gate_up_proj_weight_quantizers", + "down_proj_input_quantizer", + "down_proj_weight_quantizers", +) +_FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS = ( + "gate_up_proj_input_quantizer", + "gate_up_proj_weight_quantizer", + "down_proj_input_quantizer", + "down_proj_weight_quantizer", +) + + +def _get_replay_quantizer_attr(attr_name: str) -> str: + """Return the quantizer name used by config matching/replay.""" + if attr_name.endswith("_quantizers"): + return attr_name.removesuffix("s") + return attr_name + + +def _get_quantizer_attrs(module: nn.Module) -> tuple[str, ...]: + """Return the quantizer attribute names that AutoQuantize must snapshot/restore. + + For fused MoE experts, this returns the four plural quantizer attrs (two + shared input quantizers + two ``ModuleList`` of per-expert weight quantizers). + For standard Linear-derived QuantModules, returns the canonical trio. + """ + if _is_hf_quant_fused_experts_module(module): + return _FUSED_EXPERTS_QUANTIZER_ATTRS + return _STD_QUANTIZER_ATTRS + + +def _make_fresh_quantizer_for_attr(module: nn.Module, attr_name: str) -> nn.Module: + """Return a fresh, default quantizer object suitable to overwrite ``module.``. + + For ModuleList attrs (per-expert quantizers on fused-experts modules), the + returned ModuleList preserves the original list length so per-expert + enumeration stays consistent across recipes. + """ + current = getattr(module, attr_name, None) + if isinstance(current, nn.ModuleList): + return nn.ModuleList(TensorQuantizer() for _ in range(len(current))) + return TensorQuantizer() + + def estimate_quant_compression(quant_cfg: QuantizeConfig) -> float: """Estimate the compression ratio of a quantization configuration. @@ -222,7 +282,11 @@ def __init__( self.name = name self.quant_module_names = quant_module_names or [] - assert cost_weight > 0.0, "cost_weight must be positive." + self.quant_module_replay_attrs = { + name: tuple(_get_replay_quantizer_attr(attr) for attr in _get_quantizer_attrs(module)) + for module, name in zip(quant_modules or [], self.quant_module_names) + } + assert cost_weight >= 0.0, "cost_weight must be non-negative." self.cost_weight = cost_weight self.quant_modules = list(set(quant_modules or [])) @@ -231,26 +295,26 @@ def __init__( # This is a hack; We dont want to make the input_quantizer, weight_quantizer, output_quantizer # a dynamic attribute for backward compatibility with the model_calib.py # TODO: Make input_quantizer, weight_quantizer, output_quantizer a dynamic attribute and get rid of this hack + # NOTE: For fused-experts modules, the relevant attrs are plural + # (``*_input_quantizer`` + ``*_weight_quantizers`` ModuleList) — see + # ``_get_quantizer_attrs``. Both layouts share the same snapshot dict + # shape so ``active.setter`` swaps the right child modules. self._all_quantizer_choices = {quant_recipe: {} for quant_recipe in self.choices} quant_recipe: QuantRecipe for quant_recipe in self.choices: for quant_module in self.quant_modules: - for quantizer_attr_name in [ - "input_quantizer", - "weight_quantizer", - "output_quantizer", - ]: - setattr(quant_module, quantizer_attr_name, TensorQuantizer()) + attr_names = _get_quantizer_attrs(quant_module) + for attr_name in attr_names: + setattr( + quant_module, + attr_name, + _make_fresh_quantizer_for_attr(quant_module, attr_name), + ) set_quantizer_by_cfg(quant_module, quant_recipe.config.quant_cfg) self._all_quantizer_choices[quant_recipe][quant_module] = { - quantizer_attr_name: getattr(quant_module, quantizer_attr_name) - for quantizer_attr_name in [ - "input_quantizer", - "weight_quantizer", - "output_quantizer", - ] + attr_name: getattr(quant_module, attr_name) for attr_name in attr_names } self.active = self.original @@ -360,6 +424,20 @@ def attrs(self) -> list[str]: return ["name", "cost_weight", *super().attrs] +_LINEAR_ATTN_QKVZ_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_qkv|in_proj_z)$") +_LINEAR_ATTN_BA_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_a|in_proj_b)$") + + +def _linear_attn_qkvz_group_key(_model, name: str) -> str | None: + m = _LINEAR_ATTN_QKVZ_RE.match(name) + return f"{m.group(1)}/qkvz" if m else None + + +def _linear_attn_ba_group_key(_model, name: str) -> str | None: + m = _LINEAR_ATTN_BA_RE.match(name) + return f"{m.group(1)}/ba" if m else None + + class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): """Base searcher for AutoQuantize algorithm.""" @@ -381,6 +459,13 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): r"^(.*?)\.(gate_proj|up_proj)$", # gate_proj, up_proj for llama like models r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts + # Qwen3.5/3.6 hybrid linear_attn: vLLM fuses (in_proj_qkv, in_proj_z) + # into ``in_proj_qkvz`` and (in_proj_a, in_proj_b) into ``in_proj_ba`` and + # requires fused shards to share quant_algo. Two callables (not one + # regex) so qkv+z and a+b produce DIFFERENT group keys; each pair + # stays with its own fusion partner. + _linear_attn_qkvz_group_key, + _linear_attn_ba_group_key, ] score_module_rules = [] @@ -411,6 +496,7 @@ def default_state_dict(self) -> SearchStateDict: "cost": {}, "active_moe_expert_ratio": None, "cost_denominator": None, + "disabled_layers": None, "candidate_stats": defaultdict(dict), "quantizer_states": {}, "best": {"recipe": {}, "constraints": {}, "score": float("inf"), "is_satisfied": False}, @@ -433,9 +519,15 @@ def load_search_checkpoint(self) -> bool: @staticmethod def _is_auto_quantize_module(module): - return ( - is_quantized_linear(module) or isinstance(module, QuantLinearConvBase) - ) and isinstance(module, QuantModule) + if (is_quantized_linear(module) or isinstance(module, QuantLinearConvBase)) and isinstance( + module, QuantModule + ): + return True + # Fused MoE experts: a single ``QuantModule`` that owns N per-expert + # weight quantizers in an ``nn.ModuleList`` plus shared input quantizers. + # All N experts in a layer share one search dimension (one recipe per + # fused module). + return _is_hf_quant_fused_experts_module(module) and isinstance(module, QuantModule) @staticmethod def _get_search_recipes(quantization_formats): @@ -629,6 +721,7 @@ def initialize_candidate_stats(self): self.candidate_stats[name]["scores"] = scores self.candidate_stats[name]["costs"] = costs self.candidate_stats[name]["module_names"] = hparam.quant_module_names + self.candidate_stats[name]["quantizer_attrs"] = hparam.quant_module_replay_attrs self.candidate_stats[name]["cost_weight"] = hparam.cost_weight def _run_func(self, func, num_iters=1, desc=""): @@ -677,6 +770,7 @@ def before_search(self): self.cost_model = self.config["cost_model"] self.cost = self.config["cost"] self.active_moe_expert_ratio = self.config["active_moe_expert_ratio"] + self.disabled_layers = self.config["disabled_layers"] self.cost_denominator = getattr(self, "cost_denominator", None) search_recipes = self._get_search_recipes(self.config["quantization_formats"]) @@ -765,11 +859,9 @@ def _print_recipe_summary(best_recipe, total_cost, total_weight_size, prefix="Au @staticmethod def _get_total_weight_size(modules): return sum( - ( - module.weight.numel() - if _AutoQuantizeBaseSearcher._is_auto_quantize_module(module) - else 0 - ) + _get_module_weight_numel(module) + if _AutoQuantizeBaseSearcher._is_auto_quantize_module(module) + else 0 for module in modules ) @@ -1372,6 +1464,32 @@ def run_search_with_stats(self, max_weight_size, verbose=False): AutoQuantizeSearcher = AutoQuantizeGradientSearcher +def _as_list(value) -> list: + if value is None: + return [] + if isinstance(value, list): + return value + if isinstance(value, tuple): + return list(value) + return [value] + + +def _get_replay_quantizer_attrs(candidate_stat: dict, module_name: str) -> tuple[str, ...]: + """Return quantizer attrs that a generated config should target for a searched module.""" + quantizer_attrs = candidate_stat.get("quantizer_attrs") + if isinstance(quantizer_attrs, dict): + attrs = quantizer_attrs.get(module_name) + if attrs: + return tuple(attrs) + + # Backward-compatible fallback for search checkpoints saved before + # ``quantizer_attrs`` was persisted. Structural HF fused experts are searched + # as modules named ``...mlp.experts`` and expose gate/up + down quantizers. + if module_name.endswith((".mlp.experts", ".mixer.experts")): + return _FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS + return _STD_QUANTIZER_ATTRS + + def get_auto_quantize_config(search_state, constraints=None, verbose=False): """Build a flat quant config dict from auto_quantize search_state. @@ -1401,16 +1519,22 @@ def _cfg_to_dict(v): return v quant_cfg: list[dict] = [{"quantizer_name": "*", "enable": False}] - _per_module_attrs = ("input_quantizer", "weight_quantizer", "output_quantizer") + quant_cfg.extend( + {"quantizer_name": pattern, "enable": False} + for pattern in _as_list(search_state.get("disabled_layers")) + ) + per_module_entries: list[dict] = [] + _per_module_attrs = (*_STD_QUANTIZER_ATTRS, *_FUSED_EXPERTS_REPLAY_QUANTIZER_ATTRS) # Track global (non per-module) recipe entries. Last recipe wins for each pattern. global_entries: dict[str, dict] = {} for hparam_name, recipe in best_recipe.items(): if recipe == QuantRecipe(quant_cfg=None): continue - module_names = search_state["candidate_stats"][hparam_name]["module_names"] + candidate_stat = search_state["candidate_stats"][hparam_name] + module_names = candidate_stat["module_names"] for module_name in module_names: - for quantizer_attr in _per_module_attrs: + for quantizer_attr in _get_replay_quantizer_attrs(candidate_stat, module_name): matched_cfg, matched_enable = _match_quantizer_cfg( recipe.config.quant_cfg, quantizer_attr ) @@ -1421,7 +1545,7 @@ def _cfg_to_dict(v): } if matched_cfg is not None: entry["cfg"] = _cfg_to_dict(matched_cfg) - quant_cfg.append(entry) + per_module_entries.append(entry) # Collect non-per-module entries (e.g. *[kv]_bmm_quantizer) from winning recipes. for recipe_entry in recipe.config.quant_cfg: @@ -1438,7 +1562,10 @@ def _cfg_to_dict(v): ge["cfg"] = _cfg_to_dict(cfg) global_entries[pattern] = ge + # Keep path-scoped recipe entries before explicit module entries so selected + # modules override default disables such as ``*lm_head*``. quant_cfg.extend(global_entries.values()) + quant_cfg.extend(per_module_entries) warnings.warn( "get_auto_quantize_config: returned config uses algorithm='max'. " "Per-recipe calibration algorithms (e.g. smoothquant, awq) are not preserved. " @@ -1502,6 +1629,9 @@ def _match_quantizer_cfg(quant_cfg, quantizer_attr): matched = None matched_enable = None for entry in quant_cfg: + parent_class = entry.get("parent_class") if hasattr(entry, "get") else entry.parent_class + if parent_class is not None: + continue pattern = entry["quantizer_name"] cfg = entry.get("cfg") enable = entry.get("enable", True) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 314d630e07f..7dbdd36d04e 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -318,7 +318,10 @@ def auto_quantize( constraints = { "effective_bits": 4.8, "cost_model": "active_moe", - "cost": {"active_moe_expert_ratio": 0.25}, + "cost": { + "active_moe_expert_ratio": 0.25, + "excluded_module_name_patterns": ["*visual*", "*vision_tower*", "*mtp*"], + }, } quantization_formats: A list of quantization format config dictionaries or string names to search for. diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 1873ecda528..631226dd090 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -946,6 +946,11 @@ def fold_weight(self, keep_attrs: bool = False): delattr(q, attr_name) +def _is_quant_fused_experts_module(module): + """Return True for a converted HF fused-MoE-experts quantization wrapper.""" + return isinstance(module, _QuantFusedExperts) + + class _QuantDbrxFFN(_QuantSparseSequentialMoe): @property def num_experts(self): diff --git a/tests/examples/llm_ptq/test_hf_ptq_args.py b/tests/examples/llm_ptq/test_hf_ptq_args.py new file mode 100644 index 00000000000..94f0b079df7 --- /dev/null +++ b/tests/examples/llm_ptq/test_hf_ptq_args.py @@ -0,0 +1,110 @@ +# 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. + +import importlib +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "llm_ptq" + + +def _import_hf_ptq(monkeypatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + 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]) + parsed_args = hf_ptq.parse_args() + parsed_args.dataset = ( + parsed_args.dataset.split(",") + if isinstance(parsed_args.dataset, str) + else parsed_args.dataset + ) + parsed_args.calib_size = [int(num_sample) for num_sample in parsed_args.calib_size.split(",")] + 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*", + "*linear_attn.in_proj_a*", + "*linear_attn.in_proj_b*", + } + + 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) diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 4378d7bbc9e..d241c5df5bb 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -24,7 +24,11 @@ import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq -from modelopt.torch.quantization._auto_quantize_cost import infer_active_moe_expert_ratio +from modelopt.torch.quantization._auto_quantize_cost import ( + EXCLUDED_MODULE_NAME_PATTERNS_KEY, + get_auto_quantize_cost_model, + infer_active_moe_expert_ratio, +) from modelopt.torch.quantization.algorithms import ( AutoQuantizeGradientSearcher, QuantRecipe, @@ -154,6 +158,71 @@ def test_quant_recipe_hparam_cost_weight(): assert int8_cost == pytest.approx(model_test.weight.numel() * 0.25 * 0.5) +def test_quant_recipe_hparam_zero_cost_weight(): + model_test = mtq.quantize(torch.nn.Linear(4, 16), mtq.INT8_DEFAULT_CFG) + hparam = QuantRecipeHparam( + [QuantRecipe(mtq.INT8_DEFAULT_CFG)], + quant_modules=[model_test], + quant_module_names=["visual.blocks.0.attn.qkv"], + cost_weight=0.0, + ) + + assert hparam.get_cost(QuantRecipe(quant_cfg=None)) == pytest.approx(0.0) + assert hparam.get_cost(QuantRecipe(mtq.INT8_DEFAULT_CFG)) == pytest.approx(0.0) + + +def test_auto_quantize_cost_model_excludes_module_name_patterns(): + visual = torch.nn.Linear(4, 16) + mtp = torch.nn.Linear(4, 16) + lm_head = torch.nn.Linear(4, 16) + cost_model = get_auto_quantize_cost_model("weight") + cost_constraints = {EXCLUDED_MODULE_NAME_PATTERNS_KEY: ["*visual*", "*vision_tower*", "*mtp*"]} + + total_weight_size = cost_model.total_weight_size( + [ + ("model.visual.blocks.0.attn.qkv", visual), + ("model.mtp.layers.0.mlp", mtp), + ("lm_head", lm_head), + ], + is_auto_quantize_module=lambda module: True, + cost_constraints=cost_constraints, + ) + + assert total_weight_size == pytest.approx(lm_head.weight.numel()) + assert cost_model.module_cost_weight(["model.visual.blocks.0.attn.qkv"], cost_constraints) == 0 + assert cost_model.module_cost_weight(["model.mtp.layers.0.mlp"], cost_constraints) == 0 + assert ( + cost_model.module_cost_weight( + ["model.visual.blocks.0.attn.qkv", "lm_head"], cost_constraints + ) + == 1.0 + ) + + +def test_active_moe_cost_model_counts_fused_experts_without_weight(): + fused_experts = torch.nn.Module() + fused_experts.gate_up_proj = torch.nn.Parameter(torch.empty(2, 3, 5)) + fused_experts.down_proj = torch.nn.Parameter(torch.empty(2, 5, 3)) + visual = torch.nn.Linear(4, 16) + cost_model = get_auto_quantize_cost_model("active_moe") + + total_weight_size = cost_model.total_weight_size( + [ + ("layers.0.mlp.experts", fused_experts), + ("model.visual.blocks.0.attn.qkv", visual), + ], + is_auto_quantize_module=lambda module: True, + cost_constraints={ + "active_moe_expert_ratio": 0.25, + EXCLUDED_MODULE_NAME_PATTERNS_KEY: ["*visual*"], + }, + ) + + assert total_weight_size == pytest.approx( + (fused_experts.gate_up_proj.numel() + fused_experts.down_proj.numel()) * 0.25 + ) + + @pytest.mark.parametrize("num_experts_attr", ["num_experts", "num_local_experts"]) def test_auto_quantize_active_moe_cost_model(num_experts_attr): model = _AutoQuantMoeModel(num_experts_attr) @@ -620,3 +689,62 @@ def test_get_auto_quantize_config(method): fresh_model = mtq.quantize(fresh_model, config, forward_loop=lambda m: m(model.get_input())) output = fresh_model(model.get_input()) assert output is not None + + +def test_get_auto_quantize_config_keeps_selected_lm_head_enabled(): + recipe_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + recipe_config["quant_cfg"].append({"quantizer_name": "*lm_head*", "enable": False}) + recipe = QuantRecipe(recipe_config, name="explicit_lm_head_disable") + search_state = { + "best": {"recipe": {"lm_head.quant_recipe": recipe}}, + "candidate_stats": {"lm_head.quant_recipe": {"module_names": ["lm_head"]}}, + "disabled_layers": ["*visual*", "*mtp*"], + } + + config = mtq.get_auto_quantize_config(search_state) + quant_cfg = config["quant_cfg"] + quantizer_names = [entry["quantizer_name"] for entry in quant_cfg] + + default_disable_idx = next( + idx for idx, entry in enumerate(quant_cfg) if entry["quantizer_name"] == "*lm_head*" + ) + weight_idx = next( + idx + for idx, entry in enumerate(quant_cfg) + if entry["quantizer_name"] == "lm_head.weight_quantizer" + ) + weight_entry = quant_cfg[weight_idx] + + assert "*visual*" in quantizer_names + assert "*mtp*" in quantizer_names + assert default_disable_idx < weight_idx + assert weight_entry["enable"] is True + assert weight_entry["cfg"]["num_bits"] == (4, 3) + + +@pytest.mark.parametrize("with_persisted_attrs", [True, False]) +def test_get_auto_quantize_config_emits_fused_expert_quantizer_names(with_persisted_attrs): + recipe = QuantRecipe(copy.deepcopy(mtq.FP8_DEFAULT_CFG), name="fp8") + module_name = "layers.0.mlp.experts" + candidate_stat = {"module_names": [module_name]} + if with_persisted_attrs: + candidate_stat["quantizer_attrs"] = { + module_name: [ + "gate_up_proj_input_quantizer", + "gate_up_proj_weight_quantizer", + "down_proj_input_quantizer", + "down_proj_weight_quantizer", + ] + } + search_state = { + "best": {"recipe": {f"{module_name}.quant_recipe": recipe}}, + "candidate_stats": {f"{module_name}.quant_recipe": candidate_stat}, + "disabled_layers": [], + } + + config = mtq.get_auto_quantize_config(search_state) + quantizer_names = {entry["quantizer_name"] for entry in config["quant_cfg"]} + + assert f"{module_name}.gate_up_proj_weight_quantizer" in quantizer_names + assert f"{module_name}.down_proj_weight_quantizer" in quantizer_names + assert f"{module_name}.weight_quantizer" not in quantizer_names diff --git a/tests/unit/torch/quantization/test_config_validation.py b/tests/unit/torch/quantization/test_config_validation.py index 17ee64a6d8a..93a60924792 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/tests/unit/torch/quantization/test_config_validation.py @@ -481,6 +481,22 @@ def test_last_match_wins(self): matched, _ = _match_quantizer_cfg(quant_cfg, "weight_quantizer") assert matched.model_dump(exclude_unset=True) == {"num_bits": 4} + def test_parent_class_scoped_entries_are_ignored_for_bare_autoquant_lookup(self): + """parent_class-scoped globals should not override AutoQuantize bare-name lookup.""" + quant_cfg = normalize_quant_cfg_list( + [ + {"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8}}, + { + "parent_class": "nn.BatchNorm1d", + "quantizer_name": "*", + "enable": False, + }, + ] + ) + matched, enable = _match_quantizer_cfg(quant_cfg, "weight_quantizer") + assert matched.model_dump(exclude_unset=True) == {"num_bits": 8} + assert enable is True + def test_no_match_returns_none(self): """No matching entry returns (None, None).""" quant_cfg = normalize_quant_cfg_list(