diff --git a/CHANGELOG.rst b/CHANGELOG.rst index aadd0f0563f..3920525673b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -28,6 +28,7 @@ Experimental - Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. +- Add group-boundary AutoQuantize scoring and make it the default for gradient scoring. Quantizable descendants of common Hugging Face and Megatron attention parents are scored at their parent attention output without grouping recipe decisions; attention containers resolve to their invoked child, while fused, shared, and unfused MoE experts retain the established parent-MLP scoring boundary. Explicit mutable cache or recurrent-state inputs are rejected during parent replay, and configuration-level cache changes are transactional. Set ``score_boundary: local`` in an AutoQuantize recipe to restore leaf-output attention scoring. KL-divergence retains local scoring. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 950f3640aaf..7c2f30298a2 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -352,7 +352,8 @@ 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, optional fixed PTQ baseline, `effective_bits` target, cost model, scoring method, +candidate formats, optional fixed PTQ baseline, `effective_bits` target, cost model, scoring method +and boundary, 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 @@ -360,11 +361,13 @@ recipes (carrying architecture-specific disabled layers — e.g. VL vision tower `modelopt_recipes/huggingface//auto_quantize/`. > *Migration: prefer an AutoQuantize `--recipe`. The `--auto_quantize_bits`, `--auto_quantize_method`, -> `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and `--auto_quantize_active_moe_expert_ratio` +> `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and +> `--auto_quantize_active_moe_expert_ratio` > CLI flags are **deprecated but still work** — they are converted into an `AutoQuantizeConfig` on the fly > (with a `DeprecationWarning`) and will be removed in a future release. They map to recipe fields: > `--auto_quantize_bits` → `constraints.effective_bits`, `--auto_quantize_method` → `auto_quantize_method`, -> `--auto_quantize_score_size` → `score_size`, `--auto_quantize_cost_model` → `constraints.cost_model`, +> `--auto_quantize_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` @@ -385,7 +388,15 @@ keeps the more sensitive ones at higher precision (or unquantized), so the model `constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div`), `score_size`, `module_search_spaces` (optional per-module candidate overrides), `disabled_layers` (excluded from the search), and `cost_excluded_layers` (kept out of the bit-budget accounting — e.g. VL vision -towers). Recipes can splice a shared base `disabled_layers` set via `$import` (see +towers), and `score_boundary` (`local` / `group`). Group scoring is the default for gradient; it +changes the measurement boundary but does not force attention projections to share one recipe +decision. Set `score_boundary: local` in the recipe to restore leaf-output attention scoring; +expert projections retain their established parent MLP/mixer scoring boundary. Attention parent +names ending in `attn` or `attention` are recognized across fused and split projection layouts. +For attention stored in a `ModuleList` or `ModuleDict`, scoring resolves the invoked child instead +of the uncalled container. Group scoring rejects explicit past/cache/recurrent-state inputs because +replaying mutable state would compare candidates at different sequence positions. Recipes can splice +a shared base `disabled_layers` set via `$import` (see `modelopt_recipes/configs/auto_quantize/units/base_disabled_layers`). AutoQuantize recipes support two mutually exclusive search-space styles: @@ -450,7 +461,9 @@ The example scripts above also have an additional flag `--tasks`, where the actu > *If GPU out-of-memory error is reported running the scripts, please try editing the scripts and reducing the max batch size to save GPU memory.* -> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` method. The `kl_div` method does not require backpropagation.* +> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support +> (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` method. The `kl_div` +> method does not require backpropagation.* ## Real Quant diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 57a3dd6e264..5aaa54162b8 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -371,6 +371,7 @@ def _mtq_inputs_from_auto_quantize_config( "disabled_layers": aq_config.disabled_layers, "kv_cache_quant_cfg": kv_cache_quant_cfg, "method": aq_config.auto_quantize_method, + "score_boundary": aq_config.score_boundary, "score_size": aq_config.score_size, } @@ -501,6 +502,7 @@ def forward_step(model, batch): verbose=True, disabled_layers=inputs["disabled_layers"], method=inputs["method"], + score_boundary=inputs["score_boundary"], checkpoint=args.auto_quantize_checkpoint, ) diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index a16cfe4401a..b951720719d 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -259,6 +259,13 @@ class AutoQuantizeConfig(ModeloptBaseConfig): title="Sensitivity scoring method", description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).", ) + score_boundary: Literal["local", "group"] | None = ModeloptField( + default=None, + title="Sensitivity score boundary", + description="'local' scores attention projections at their leaf outputs; 'group' scores " + "them at the parent attention output. Expert projections retain their established parent " + "MLP boundary in either mode. Defaults to 'group' for gradient and 'local' for kl_div.", + ) score_size: int = ModeloptField( default=128, title="Scoring sample count", @@ -293,6 +300,15 @@ def _has_search_space(self): ) return self + @model_validator(mode="after") + def _validate_scoring_configuration(self): + boundary = self.score_boundary or ( + "local" if self.auto_quantize_method == "kl_div" else "group" + ) + if self.auto_quantize_method == "kl_div" and boundary != "local": + raise ValueError("auto_quantize_method='kl_div' requires score_boundary='local'.") + return self + class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase): """Our config class for AutoQuantize recipes.""" diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 7beeef6ad7f..a36746ace67 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -18,6 +18,7 @@ import copy import fnmatch import gc +import inspect import types import warnings from abc import ABC, abstractmethod @@ -595,6 +596,56 @@ def attrs(self) -> list[str]: _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)$") +_FUSED_EXPERTS_PARENT_RE = re.compile(r"^((?:.*\.)?(?:mlp|mixer))\.experts$") +_LIKELY_ATTENTION_PROJECTION_NAMES = frozenset( + { + "c_attn", + "k_proj", + "linear_proj", + "linear_qkv", + "o_proj", + "q_proj", + "qkv_proj", + "query_key_value", + "v_proj", + "wqkv", + } +) +_MUTABLE_REPLAY_INPUT_NAMES = frozenset( + { + "cache", + "cache_params", + "inference_context", + "inference_params", + "kv_cache", + "layer_past", + "past_key_value", + "past_key_values", + "state", + } +) + +AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL = "local" +AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP = "group" +AUTO_QUANTIZE_SCORE_BOUNDARIES = frozenset( + {AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL, AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP} +) + + +def _resolve_auto_quantize_score_boundary(method: str, score_boundary: str | None) -> str: + if method not in {"gradient", "kl_div"}: + raise ValueError(f"Invalid method: {method}. Valid options are 'gradient' or 'kl_div'.") + if score_boundary is None: + score_boundary = ( + AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL + if method == "kl_div" + else AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP + ) + if score_boundary not in AUTO_QUANTIZE_SCORE_BOUNDARIES: + raise ValueError(f"score_boundary must be one of {sorted(AUTO_QUANTIZE_SCORE_BOUNDARIES)}.") + if method != "gradient" and score_boundary != AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL: + raise ValueError("score_boundary='group' is supported only with method='gradient'.") + return score_boundary def _linear_attn_qkvz_group_key(_model, name: str) -> str | None: @@ -607,6 +658,105 @@ def _linear_attn_ba_group_key(_model, name: str) -> str | None: return f"{m.group(1)}/ba" if m else None +def _is_attention_module_name(name: str) -> bool: + return name.lower().endswith(("attention", "attn")) + + +def _module_implements_forward(module: nn.Module) -> bool: + """Return whether a module implements a callable forward boundary.""" + return type(module).forward is not nn.Module.forward + + +def _attention_parent_score_module(model, name: str) -> str | None: + """Return the nearest callable attention ancestor for a quantizable descendant.""" + components = name.split(".") + for idx in range(len(components) - 2, -1, -1): + if not _is_attention_module_name(components[idx]): + continue + + parent_name = ".".join(components[: idx + 1]) + try: + parent = model.get_submodule(parent_name) + except AttributeError: + return None + if _module_implements_forward(parent): + return parent_name + + # Some models use a named attention ModuleList/ModuleDict as a container + # and invoke one of its children directly. Resolve the first callable + # descendant on the path to the quantized projection instead of hooking + # the uncalled container. + for child_idx in range(idx + 1, len(components) - 1): + child_name = ".".join(components[: child_idx + 1]) + try: + child = model.get_submodule(child_name) + except AttributeError: + return None + if _module_implements_forward(child): + return child_name + return None + return None + + +def _looks_like_unmapped_attention_projection(name: str) -> bool: + return name.rsplit(".", 1)[-1].lower() in _LIKELY_ATTENTION_PROJECTION_NAMES + + +def _bound_forward_inputs(forward, args, kwargs) -> dict[str, Any]: + """Best-effort map of positional and keyword inputs for a bound forward method.""" + inputs = dict(kwargs) + try: + signature = inspect.signature(forward) + bound = signature.bind_partial(*args, **kwargs) + except (TypeError, ValueError): + return inputs + + for name, value in bound.arguments.items(): + parameter = signature.parameters[name] + if parameter.kind == inspect.Parameter.VAR_KEYWORD and isinstance(value, dict): + inputs.update(value) + else: + inputs[name] = value + return inputs + + +def _is_mutable_replay_input_name(name: str) -> bool: + name = name.lower() + if name == "use_cache": + return False + return name in _MUTABLE_REPLAY_INPUT_NAMES or name.endswith(("_cache", "_state")) + + +def _validate_group_score_replay_inputs(module: nn.Module, args, kwargs) -> None: + """Reject explicit mutable cache or recurrent state before replaying a score module.""" + inputs = _bound_forward_inputs(module._forward_original, args, kwargs) + mutable_inputs = sorted( + name + for name, value in inputs.items() + if value is not None and _is_mutable_replay_input_name(name) + ) + if inputs.get("use_cache") is True: + mutable_inputs.append("use_cache=True") + if mutable_inputs: + raise ValueError( + "Group-boundary AutoQuantize scoring cannot replay a parent module with explicit " + f"mutable state/cache inputs ({', '.join(mutable_inputs)}). Disable use_cache and " + "omit past/cache/recurrent state in forward_step so every candidate starts from the " + "same sequence state." + ) + + +def _fused_experts_parent_score_module(model, name: str) -> str | None: + match = _FUSED_EXPERTS_PARENT_RE.match(name) + if match is None: + return None + try: + module = model.get_submodule(name) + except AttributeError: + return None + return match.group(1) if _is_hf_quant_fused_experts_module(module) else None + + def _module_search_space_signature(module_search_spaces) -> tuple: """Return a checkpoint-stable description of module-specific candidate spaces.""" return tuple( @@ -675,6 +825,7 @@ def default_search_config(self): "cost_model": COST_MODEL_WEIGHT, "cost": {}, "active_moe_expert_ratio": None, + "score_boundary": None, } @property @@ -685,6 +836,7 @@ def default_state_dict(self) -> SearchStateDict: "cost_model": "weight", "cost": {}, "active_moe_expert_ratio": None, + "score_boundary": AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL, "cost_denominator": None, "quantization_formats_signature": None, "fixed_quantization_config_signature": None, @@ -707,8 +859,19 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: assert config["forward_step"] is not None, ( "`forward_step` must be provided for `auto_quantize`." ) + assert self.method_name is not None + config["score_boundary"] = _resolve_auto_quantize_score_boundary( + self.method_name, config["score_boundary"] + ) return config + def _get_score_module_rules(self): + rules = [] + if self.config["score_boundary"] == AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP: + rules.append(_attention_parent_score_module) + rules.extend(self.score_module_rules) + return rules + def load_search_checkpoint(self) -> bool: return super().load_search_checkpoint(strict=False) @@ -895,6 +1058,7 @@ def insert_hparams_after_merge_rules( # Map from group key to list of (quant_module, name, disabled, score_module) search_map: dict[str, list[tuple[nn.Module, str, bool, nn.Module]]] = {} + unmapped_attention_projections = [] for name, module in model.named_modules(): if not self._is_auto_quantize_module(module): continue @@ -917,12 +1081,20 @@ def insert_hparams_after_merge_rules( # Apply score_module_rules to determine the score module name, then get the actual module score_module_name = name # Default: score from same module - for rule in self.score_module_rules: + score_rule_matched = False + for rule in self._get_score_module_rules(): result = self._apply_score_group_rule(name, rule) if result is not None: score_module_name = result + score_rule_matched = True # We support only one rule for matching per module break + if ( + self.config["score_boundary"] == AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP + and not score_rule_matched + and _looks_like_unmapped_attention_projection(name) + ): + unmapped_attention_projections.append(name) # Get the actual score module object immediately score_module = self._get_score_module_from_name(model, score_module_name, module) @@ -932,6 +1104,16 @@ def insert_hparams_after_merge_rules( else: search_map[group_key].append((module, name, disabled, score_module)) + if unmapped_attention_projections: + warnings.warn( + "Group-boundary AutoQuantize scoring could not map these likely attention " + "projections to a parent attention module and will score them locally: " + f"{unmapped_attention_projections}. Use a conventional attention parent name " + "(for example, attn, self_attn, self_attention, or linear_attn), register a " + "model-specific score rule, or set score_boundary='local'.", + stacklevel=2, + ) + for group_key, module_info_list in search_map.items(): quant_modules = [module for module, _, _, _ in module_info_list] disabled = any(disabled for _, _, disabled, _ in module_info_list) @@ -1099,6 +1281,11 @@ def before_search(self): super().before_search() self.constraints = normalize_auto_quantize_constraints(self.model, self.constraints) + if ( + self.method_name != "gradient" + and self.config["score_boundary"] != AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL + ): + raise ValueError("score_boundary='group' is supported only with method='gradient'.") self.config["cost_model"] = self.constraints["cost_model"] self.config["cost"] = self.constraints.get("cost", {}) self.config["active_moe_expert_ratio"] = self.config["cost"].get( @@ -1113,6 +1300,9 @@ def before_search(self): ) restored_cost_model = getattr(self, "cost_model", "weight") restored_active_moe_expert_ratio = getattr(self, "active_moe_expert_ratio", None) + restored_score_boundary = getattr( + self, "score_boundary", AUTO_QUANTIZE_SCORE_BOUNDARY_LOCAL + ) if self.candidate_stats and ( restored_cost_model != self.config["cost_model"] or restored_active_moe_expert_ratio != self.config["active_moe_expert_ratio"] @@ -1123,10 +1313,17 @@ def before_search(self): f"current=({self.config['cost_model']}, {self.config['active_moe_expert_ratio']}). " "Use a different checkpoint path." ) + if self.candidate_stats and restored_score_boundary != self.config["score_boundary"]: + raise ValueError( + "Checkpoint AutoQuantize score boundary does not match current search config: " + f"checkpoint={restored_score_boundary}, current={self.config['score_boundary']}. " + "Use a different checkpoint path." + ) self.method = self.method_name self.cost_model = self.config["cost_model"] self.cost = self.config["cost"] self.active_moe_expert_ratio = self.config["active_moe_expert_ratio"] + self.score_boundary = self.config["score_boundary"] self.disabled_layers = self.config["disabled_layers"] self.cost_denominator = getattr(self, "cost_denominator", None) @@ -1473,11 +1670,16 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): method_name = "gradient" score_module_rules = [ - # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts) + # Use the parent MLP output for unfused routed experts. r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", + # Apply the same parent boundary to shared-expert projections. + r"^((?:.*\.)?mlp)\.shared_expert\.(gate_proj|up_proj|down_proj)$", r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts + # Fused expert containers are one quant module; preserve the established + # expert behavior by measuring their perturbation at the parent MLP output. + _fused_experts_parent_score_module, ] # See `register_custom_support` for details @@ -1553,12 +1755,15 @@ def forward_backward_step(model, data): @torch.enable_grad() def _estimate_auto_quantize_scores(self, is_param_grad_enabled): # TODO: remove the no-quant recipe - def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): + def auto_quantize_score_estimate_forward(module, *args, **kwargs): + if self.config["score_boundary"] == AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP: + _validate_group_score_replay_inputs(module, args, kwargs) + for hparam in module._hparams_for_scoring: if hparam.is_configurable: hparam.active = QuantRecipe(quant_cfg=None) - output = module._forward_original(input, *args, **kwargs) + output = module._forward_original(*args, **kwargs) # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass. # With gradient checkpointing, gradients are computed in the local forward pass during backward pass @@ -1576,7 +1781,7 @@ def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): if recipe == QuantRecipe(quant_cfg=None): continue hparam.active = recipe - output_diff = module._forward_original(input, *args, **kwargs) + output_diff = module._forward_original(*args, **kwargs) if isinstance(output_diff, tuple): output_diff = output_diff[0] - output[0] @@ -1623,6 +1828,7 @@ def cleanup_module_after_score_estimation(module): del module._forward_original module._backward_hook_handle.remove() + del module._backward_hook_handle def cleanup_params_after_score_estimation(name, param, params_metadata): param.requires_grad = params_metadata[name]["requires_grad"] @@ -1652,20 +1858,28 @@ def cleanup_params_after_score_estimation(name, param, params_metadata): torch.cuda.reset_peak_memory_stats() report_memory("AutoQuantize: starting score estimation, ") - self._run_func( - self.config["forward_backward_step"], - num_iters=self.config["num_score_steps"], - desc="Estimating auto_quantize scores", - ) - - if torch.cuda.is_available(): - report_memory("AutoQuantize: After score estimation") - - for module in score_modules: - cleanup_module_after_score_estimation(module) + cache_originals = [] + try: + if self.config["score_boundary"] == AUTO_QUANTIZE_SCORE_BOUNDARY_GROUP: + cache_originals = _set_model_use_cache(self.model, False) + self._run_func( + self.config["forward_backward_step"], + num_iters=self.config["num_score_steps"], + desc="Estimating auto_quantize scores", + ) - for name, param in self.model.named_parameters(): - cleanup_params_after_score_estimation(name, param, params_metadata) + if torch.cuda.is_available(): + report_memory("AutoQuantize: After score estimation") + finally: + try: + _restore_model_use_cache(cache_originals) + finally: + try: + for module in score_modules: + cleanup_module_after_score_estimation(module) + finally: + for name, param in self.model.named_parameters(): + cleanup_params_after_score_estimation(name, param, params_metadata) # Delete the params_metadata del params_metadata @@ -1743,6 +1957,44 @@ def run_search_with_stats(self, max_weight_size, verbose=False): return best_recipes, is_satisfied +def _get_model_config_objects(model: nn.Module) -> list[Any]: + """Return unique model configs that may carry generation cache defaults.""" + configs = [] + seen = set() + for module in model.modules(): + for attr_name in ("config", "generation_config"): + config = getattr(module, attr_name, None) + if config is None or id(config) in seen: + continue + seen.add(id(config)) + configs.append(config) + for nested_name in ("text_config", "language_config"): + nested = getattr(config, nested_name, None) + if nested is not None and id(nested) not in seen: + seen.add(id(nested)) + configs.append(nested) + return configs + + +def _set_model_use_cache(model: nn.Module, use_cache: bool) -> list[tuple[Any, Any]]: + """Transactionally set use_cache on model configs and return values to restore.""" + originals = [] + try: + for config in _get_model_config_objects(model): + if hasattr(config, "use_cache"): + originals.append((config, config.use_cache)) + config.use_cache = use_cache + except BaseException: + _restore_model_use_cache(reversed(originals)) + raise + return originals + + +def _restore_model_use_cache(originals) -> None: + for config, use_cache in originals: + config.use_cache = use_cache + + @torch.compile(dynamic=True) def _get_log_softmax_dist(logits: torch.Tensor, tp_group) -> torch.Tensor: dtype = logits.dtype diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 966a3643fe3..3acbb9deadc 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -36,7 +36,12 @@ ) from modelopt.torch.utils import atomic_print -from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe +from .algorithms import ( + AutoQuantizeGradientSearcher, + AutoQuantizeKLDivSearcher, + QuantRecipe, + _resolve_auto_quantize_score_boundary, +) from .algorithms import get_auto_quantize_config as _get_auto_quantize_config from .config import QuantizeAlgoCfgType from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg @@ -282,6 +287,7 @@ def auto_quantize( checkpoint: str | None = None, module_search_spaces: list[dict[str, Any]] | None = None, fixed_quantization_config: dict[str, Any] | str | None = None, + score_boundary: str | None = None, ): r"""Perform optimal per-layer quantization by searching for the best quantization formats per-layer. @@ -442,6 +448,16 @@ def forward_backward_step(model, batch) -> None: linear programming search, and requires ``loss_func`` or ``forward_backward_step``) and ``"kl_div"`` (uses KL divergence between unquantized and quantized outputs, relies on threshold-based binary search, and only requires ``forward_step`` returning logits). + score_boundary: Boundary used to measure perturbations. ``"local"`` + scores attention projections at their leaf outputs. ``"group"`` scores them at the + nearest parent attention output. Attention parent names ending in ``attn`` or + ``attention`` are recognized, covering + common Hugging Face and Megatron layouts. Expert projections retain their established + parent MLP/mixer score boundary in either mode. This does not group recipe decisions + or force attention modules to use the same quantization format. Parent replay rejects + explicit mutable cache inputs; omit past/cache state from ``forward_step``. Defaults + to ``"group"`` for ``method="gradient"`` and ``"local"`` for + ``method="kl_div"``. checkpoint: (Optional) Path to checkpoint file for saving/restoring auto_quantize search state. If the checkpoint file exists, the search state will be restored from it, skipping the expensive score estimation step. @@ -506,6 +522,9 @@ def forward_backward_step(model, batch) -> None: might not be readily deployable to TensorRT-LLM yet. """ + # Validate and resolve this before any quantization mode is applied. A rejected call must leave + # the caller's model untouched so it can be retried with corrected configuration. + score_boundary = _resolve_auto_quantize_score_boundary(method, score_boundary) def _process_quantization_formats(formats, custom_name_prefix): processed = [] @@ -620,10 +639,9 @@ def _process_quantization_formats(formats, custom_name_prefix): # Select the appropriate searcher based on method if method == "gradient": searcher = AutoQuantizeGradientSearcher() - elif method == "kl_div": - searcher = AutoQuantizeKLDivSearcher() else: - raise ValueError(f"Invalid method: {method}. Valid options are 'gradient' or 'kl_div'.") + assert method == "kl_div" + searcher = AutoQuantizeKLDivSearcher() model = apply_mode( model, @@ -640,6 +658,7 @@ def _process_quantization_formats(formats, custom_name_prefix): "forward_backward_step": forward_backward_step, "num_calib_steps": num_calib_steps, "num_score_steps": num_score_steps, + "score_boundary": score_boundary, "disabled_layers": disabled_layers, "verbose": verbose, "checkpoint": checkpoint, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 30ab3226a64..4fe6dcca7d6 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -56,6 +56,7 @@ def test_autoquant_recipe_builds_mtq_inputs(monkeypatch): assert inputs["constraints"] == {"effective_bits": 5.4, "cost_model": "weight"} assert inputs["kv_cache_quant_cfg"] is None assert inputs["method"] == "gradient" + assert inputs["score_boundary"] is None assert inputs["score_size"] == 128 assert inputs["fixed_quantization_config"] is None assert inputs["module_search_spaces"] == [] @@ -188,6 +189,7 @@ def test_autoquant_config_from_deprecated_cli_flags(monkeypatch): 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_boundary is None 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] == [ diff --git a/tests/unit/recipe/test_loader.py b/tests/unit/recipe/test_loader.py index e15b897a224..a54a181b328 100644 --- a/tests/unit/recipe/test_loader.py +++ b/tests/unit/recipe/test_loader.py @@ -27,6 +27,8 @@ import modelopt.torch.quantization.config as qcfg from modelopt.recipe.config import ( + AutoQuantizeConfig, + AutoQuantizeConstraints, ModelOptAutoQuantizeRecipe, ModelOptDFlashRecipe, ModelOptEagleRecipe, @@ -35,7 +37,11 @@ ) from modelopt.recipe.loader import _apply_dotlist, load_config, load_recipe from modelopt.torch.opt.config_loader import _load_raw_config, _schema_type -from modelopt.torch.quantization.config import QuantizerAttributeConfig, normalize_quant_cfg_list +from modelopt.torch.quantization.config import ( + QuantizeConfig, + QuantizerAttributeConfig, + normalize_quant_cfg_list, +) # --------------------------------------------------------------------------- # Static YAML fixtures @@ -1721,6 +1727,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.score_boundary is None assert aq.score_size == 128 assert aq.kv_cache is None assert aq.constraints.effective_bits == 4.8 @@ -1800,6 +1807,16 @@ def test_load_recipe_autoquantize_effective_bits_out_of_range_raises(tmp_path): load_recipe(bad) +def test_load_recipe_autoquantize_rejects_group_boundary_for_kl_div(): + with pytest.raises(ValueError, match=r"requires score_boundary='local'"): + AutoQuantizeConfig( + constraints=AutoQuantizeConstraints(effective_bits=5.4), + candidate_formats=[QuantizeConfig(quant_cfg=[])], + auto_quantize_method="kl_div", + score_boundary="group", + ) + + 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") diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index e83f7fa0a70..0e868bb012c 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -73,6 +73,143 @@ def get_input(self): return torch.randn(1, 4, 32) +class _LinearAttentionLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.in_proj_qkv = torch.nn.Linear(32, 32) + self.in_proj_z = torch.nn.Linear(32, 32) + self.in_proj_a = torch.nn.Linear(32, 32) + self.in_proj_b = torch.nn.Linear(32, 32) + self.out_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + x = self.in_proj_qkv(x) + self.in_proj_z(x) + x = x + self.in_proj_a(x) + self.in_proj_b(x) + return self.out_proj(x) + + +class _GroupBoundaryModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + use_cache=True, + text_config=SimpleNamespace(use_cache=True), + ) + self.use_cache_seen = [] + self.self_attn = _AttentionLayer() + self.linear_attn = _LinearAttentionLayer() + + def forward(self, x): + self.use_cache_seen.append((self.config.use_cache, self.config.text_config.use_cache)) + return self.linear_attn(x=self.self_attn(x=x)) + + def get_input(self): + return torch.randn(1, 4, 32) + + +class _ExplicitCacheGroupBoundaryModel(torch.nn.Module): + class CacheAwareAttention(_AttentionLayer): + def forward(self, x, cache_params=None): + if cache_params is not None: + cache_params["positions"].append(len(cache_params["positions"]) + 1) + return super().forward(x) + + def __init__(self): + super().__init__() + self.config = SimpleNamespace(use_cache=False) + self.cache = {"positions": []} + self.self_attn = self.CacheAwareAttention() + + def forward(self, x): + # Pass the mutable cache positionally to verify signature binding, not just kwargs. + return self.self_attn(x, self.cache) + + def get_input(self): + return torch.randn(1, 4, 32) + + +class _ExplicitMutableStateGroupBoundaryModel(torch.nn.Module): + class StateAwareAttention(_AttentionLayer): + def forward(self, x, state=None, cache=None): + mutable = state if state is not None else cache + if mutable is not None: + mutable["positions"].append(len(mutable["positions"]) + 1) + return super().forward(x) + + def __init__(self, input_name, as_keyword): + super().__init__() + self.config = SimpleNamespace(use_cache=False) + self.mutable = {"positions": []} + self.input_name = input_name + self.as_keyword = as_keyword + self.self_attn = self.StateAwareAttention() + + def forward(self, x): + if self.as_keyword: + return self.self_attn(x, **{self.input_name: self.mutable}) + if self.input_name == "state": + return self.self_attn(x, self.mutable) + return self.self_attn(x, None, self.mutable) + + def get_input(self): + return torch.randn(1, 4, 32) + + +class _ExplicitUseCacheGroupBoundaryModel(torch.nn.Module): + class CacheFlagAttention(_AttentionLayer): + def forward(self, x, use_cache=False): + return super().forward(x) + + def __init__(self, use_cache): + super().__init__() + self.config = SimpleNamespace(use_cache=False) + self.use_cache = use_cache + self.self_attn = self.CacheFlagAttention() + + def forward(self, x): + return self.self_attn(x, use_cache=self.use_cache) + + def get_input(self): + return torch.randn(1, 4, 32) + + +class _NestedAttentionContainerModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(use_cache=False) + self.self_attn = torch.nn.ModuleList([_AttentionLayer(), _AttentionLayer()]) + + def forward(self, x): + for attention in self.self_attn: + x = attention(x) + return x + + def get_input(self): + return torch.randn(1, 4, 32) + + +class _FailingUseCacheConfig: + def __init__(self): + self._use_cache = True + + @property + def use_cache(self): + return self._use_cache + + @use_cache.setter + def use_cache(self, value): + self._use_cache = value + if value is False: + raise RuntimeError("cannot disable use_cache") + + +class _PartiallyFailingCacheConfigModel(_GroupBoundaryModel): + def __init__(self): + super().__init__() + self.failing_config_module = torch.nn.Module() + self.failing_config_module.config = _FailingUseCacheConfig() + + class _AutoQuantMoeModel(torch.nn.Module): def __init__(self, num_experts_attr="num_experts"): super().__init__() @@ -708,6 +845,375 @@ def test_active_moe_search_prefers_budget_lower_bound(): assert best_recipes["layers.0.mlp.quant_recipe"]["format"] == "near_budget" +def test_auto_quantize_local_boundary_scores_shared_expert_at_parent_mlp(): + searcher = AutoQuantizeGradientSearcher() + searcher.config = {"score_boundary": "local"} + + score_modules = [] + for projection in ("gate_proj", "up_proj", "down_proj"): + name = f"model.layers.0.mlp.shared_expert.{projection}" + score_module = next( + ( + result + for rule in searcher._get_score_module_rules() + if (result := searcher._apply_score_group_rule(name, rule)) is not None + ), + name, + ) + score_modules.append(score_module) + + assert score_modules == ["model.layers.0.mlp"] * 3 + + +@pytest.mark.parametrize("parent_name", ["mlp", "mixer"]) +def test_auto_quantize_local_boundary_scores_fused_experts_at_parent(monkeypatch, parent_name): + model = torch.nn.Module() + parent = torch.nn.Module() + parent.experts = torch.nn.Module() + setattr(model, parent_name, parent) + monkeypatch.setattr( + "modelopt.torch.quantization.algorithms._is_hf_quant_fused_experts_module", + lambda module: module is parent.experts, + ) + searcher = AutoQuantizeGradientSearcher() + searcher.model = model + searcher.config = {"score_boundary": "local"} + + score_module = next( + result + for rule in searcher._get_score_module_rules() + if ( + result := searcher._apply_score_group_rule( + f"{parent_name}.experts", + rule, + ) + ) + is not None + ) + + assert score_module == parent_name + + +@pytest.mark.parametrize( + ("module_name", "expected_parent"), + [ + ("model.layers.0.self_attn.qkv_proj", "model.layers.0.self_attn"), + ("model.layers.0.self_attn.o_proj", "model.layers.0.self_attn"), + ("decoder.layers.0.self_attention.linear_qkv", "decoder.layers.0.self_attention"), + ("decoder.layers.0.self_attention.linear_proj", "decoder.layers.0.self_attention"), + ("transformer.h.0.attn.c_attn", "transformer.h.0.attn"), + ("transformer.h.0.crossattention.c_attn", "transformer.h.0.crossattention"), + ("encoder.layer.0.attention.self.query", "encoder.layer.0.attention"), + ("model.layers.0.linear_attn.in_proj_qkv", "model.layers.0.linear_attn"), + ], +) +def test_auto_quantize_group_boundary_maps_supported_attention_layouts( + module_name, expected_parent +): + class LookupModel: + def get_submodule(self, name): + assert name == expected_parent + return _AttentionLayer() + + searcher = AutoQuantizeGradientSearcher() + searcher.model = LookupModel() + searcher.config = {"score_boundary": "group"} + + score_module = next( + result + for rule in searcher._get_score_module_rules() + if (result := searcher._apply_score_group_rule(module_name, rule)) is not None + ) + + assert score_module == expected_parent + + +def test_auto_quantize_group_boundary_resolves_invoked_child_under_attention_container(): + model = _NestedAttentionContainerModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + ) + + first_q_proj = model.self_attn[0].q_proj.get_hparam("quant_recipe") + second_o_proj = model.self_attn[1].o_proj.get_hparam("quant_recipe") + assert first_q_proj.score_modules == [model.self_attn[0]] + assert second_o_proj.score_modules == [model.self_attn[1]] + assert not hasattr(model.self_attn, "_hparams_for_scoring") + assert any( + score > 0 for stats in search_state["candidate_stats"].values() for score in stats["scores"] + ) + + +def test_auto_quantize_group_boundary_warns_on_unmapped_attention_projection(): + class RootAttentionProjection(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.q_proj(x) + + model = RootAttentionProjection() + batch = torch.randn(1, 4, 32) + with pytest.warns(UserWarning, match=r"could not map.*q_proj"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + data_loader=[batch], + forward_step=lambda model, data: model(data), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + ) + + +def test_auto_quantize_group_score_boundary_does_not_group_recipe_decisions(): + model = _GroupBoundaryModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + ) + + self_qkv = model.self_attn.q_proj.get_hparam("quant_recipe") + self_o = model.self_attn.o_proj.get_hparam("quant_recipe") + linear_qkvz = model.linear_attn.in_proj_qkv.get_hparam("quant_recipe") + linear_ba = model.linear_attn.in_proj_a.get_hparam("quant_recipe") + linear_out = model.linear_attn.out_proj.get_hparam("quant_recipe") + + assert self_qkv is not self_o + assert linear_qkvz is not linear_ba + assert linear_qkvz is not linear_out + assert linear_ba is not linear_out + assert self_qkv.score_modules == [model.self_attn] + assert self_o.score_modules == [model.self_attn] + assert linear_qkvz.score_modules == [model.linear_attn] + assert linear_ba.score_modules == [model.linear_attn] + assert linear_out.score_modules == [model.linear_attn] + assert search_state["score_boundary"] == "group" + assert search_state["method"] == "gradient" + assert any( + score > 0 for stats in search_state["candidate_stats"].values() for score in stats["scores"] + ) + assert model.config.use_cache is True + assert model.config.text_config.use_cache is True + assert (False, False) in model.use_cache_seen + + +def test_auto_quantize_group_boundary_rejects_explicit_mutable_cache_before_replay(): + model = _ExplicitCacheGroupBoundaryModel() + + with pytest.raises(ValueError, match=r"mutable state/cache inputs \(cache_params\)"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + ) + + # Each format runs one calibration forward. The parent replay must fail before a third, + # baseline/candidate forward can mutate and reuse the same cache. + assert model.cache["positions"] == [1, 2] + + +@pytest.mark.parametrize("input_name", ["state", "cache"]) +@pytest.mark.parametrize("as_keyword", [False, True]) +def test_auto_quantize_group_boundary_rejects_generic_mutable_state_before_replay( + input_name, as_keyword +): + model = _ExplicitMutableStateGroupBoundaryModel(input_name, as_keyword) + + with pytest.raises(ValueError, match=rf"mutable state/cache inputs \({input_name}\)"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + ) + + assert model.mutable["positions"] == [1, 2] + + +@pytest.mark.parametrize("use_cache", [False, True]) +def test_auto_quantize_group_boundary_handles_explicit_use_cache_flag(use_cache): + model = _ExplicitUseCacheGroupBoundaryModel(use_cache) + kwargs = { + "constraints": {"effective_bits": 8.0}, + "quantization_formats": [ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + "data_loader": [model.get_input()], + "forward_step": lambda model, batch: model(batch), + "loss_func": lambda output, data: output.sum(), + "num_calib_steps": 1, + "num_score_steps": 1, + "method": "gradient", + } + + if use_cache: + with pytest.raises(ValueError, match=r"mutable state/cache inputs \(use_cache=True\)"): + mtq.auto_quantize(model, **kwargs) + else: + _, search_state = mtq.auto_quantize(model, **kwargs) + assert any( + score > 0 + for stats in search_state["candidate_stats"].values() + for score in stats["scores"] + ) + + +def test_auto_quantize_group_boundary_restores_cleanup_when_cache_disabling_fails(): + model = _PartiallyFailingCacheConfigModel() + + with pytest.raises(RuntimeError, match="cannot disable use_cache"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + ) + + assert model.config.use_cache is True + assert model.config.text_config.use_cache is True + assert model.failing_config_module.config.use_cache is True + for score_module in (model.self_attn, model.linear_attn): + assert not hasattr(score_module, "_forward_original") + assert not hasattr(score_module, "_backward_hook_handle") + + +def test_auto_quantize_gradient_defaults_to_group_boundary_and_supports_keyword_calls(): + model = _GroupBoundaryModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + ) + + assert search_state["score_boundary"] == "group" + assert search_state["method"] == "gradient" + assert any( + score > 0 for stats in search_state["candidate_stats"].values() for score in stats["scores"] + ) + assert model.config.use_cache is True + assert model.config.text_config.use_cache is True + assert (False, False) in model.use_cache_seen + + +def test_auto_quantize_gradient_local_boundary_opt_out(): + model = _GroupBoundaryModel() + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + score_boundary="local", + ) + + assert search_state["score_boundary"] == "local" + assert (False, False) not in model.use_cache_seen + + +@pytest.mark.parametrize( + ("method", "score_boundary", "match"), + [ + ("gradient", "unsupported", "score_boundary must be one of"), + ("kl_div", "group", "supported only with method='gradient'"), + ], +) +def test_auto_quantize_validates_score_boundary_before_mutating_model( + method, score_boundary, match +): + model = TransformerBlock() + original_mlp_type = type(model.mlp) + + with pytest.raises(ValueError, match=match): + mtq.auto_quantize(model, method=method, score_boundary=score_boundary) + + assert type(model.mlp) is original_mlp_type + assert not hasattr(model.mlp, "weight_quantizer") + + +def test_auto_quantize_can_retry_after_rejected_score_boundary(): + model = TransformerBlock() + with pytest.raises(ValueError, match="supported only with method='gradient'"): + mtq.auto_quantize(model, method="kl_div", score_boundary="group") + + _, search_state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_WEIGHT_ONLY_CFG, + ], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=1, + num_score_steps=1, + method="gradient", + score_boundary="local", + ) + + assert search_state["score_boundary"] == "local" + + # use this config to test custom quantization config INT8_CUSTOM_QUANT_TEST_CFG = { "quant_cfg": [ @@ -1434,6 +1940,38 @@ def interrupt_after_calibration(self): ) +def test_auto_quantize_checkpoint_rejects_score_boundary_mismatch(tmp_path): + checkpoint_path = str(tmp_path / "autoquant_group_boundary.pth") + model = _GroupBoundaryModel() + common = { + "constraints": {"effective_bits": 8.0}, + "quantization_formats": [ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.INT8_DEFAULT_CFG, + ], + "forward_step": lambda model, batch: model(batch), + "loss_func": lambda output, data: output.sum(), + "num_calib_steps": 1, + "num_score_steps": 1, + "checkpoint": checkpoint_path, + } + mtq.auto_quantize( + model, + data_loader=[model.get_input()], + score_boundary="group", + **common, + ) + + resumed_model = _GroupBoundaryModel() + with pytest.raises(ValueError, match="score boundary does not match"): + mtq.auto_quantize( + resumed_model, + data_loader=[resumed_model.get_input()], + score_boundary="local", + **common, + ) + + @pytest.mark.parametrize("method", ["gradient", "kl_div"]) def test_get_auto_quantize_config(method): model = TransformerBlock() @@ -1452,6 +1990,7 @@ def test_get_auto_quantize_config(method): # Verify search_state has method and module_names assert search_state["method"] == method + assert search_state["score_boundary"] == ("local" if method == "kl_div" else "group") for stats in search_state["candidate_stats"].values(): assert "module_names" in stats assert len(stats["module_names"]) > 0