diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 724876fcafa..78845904691 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add ``method="aumann_shapley"`` to ``mtq.auto_quantize``: label-free sensitivity scoring via Aumann-Shapley path-integral damage attributions with a measured-corner coverage calibration, so every allocation carries a ``predicted_damage`` quote in calibration-KL units. Method-specific settings (path nodes, damage link, a deterministic DP solver, and a ``max_predicted_damage`` bound mode) ride in a new optional ``auto_quantize(method_options=...)`` argument validated by the selected method. AutoQuantize scoring methods are now registered in ``modelopt.torch.quantization.algorithms.AUTO_QUANTIZE_SEARCHERS``. - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. *Megatron Framework (M-LM / M-Bridge)* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 72b2e2348ab..a2402be6bfa 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -382,7 +382,7 @@ 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`), `score_size`, +`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div` / `aumann_shapley`), `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 @@ -450,7 +450,7 @@ 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` or `aumann_shapley` methods (the latter is label-free but still backpropagates a KL loss). 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 0790f644308..1759b4fc1bd 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -495,7 +495,7 @@ 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": + elif inputs["method"] in ("kl_div", "aumann_shapley"): def forward_step(model, batch): inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch @@ -507,7 +507,8 @@ def forward_step(model, batch): else: raise ValueError( - f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'" + f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient', 'kl_div', " + "or 'aumann_shapley'" ) language_model, _ = mtq.auto_quantize( @@ -1574,7 +1575,7 @@ def parse_args() -> argparse.Namespace: "--auto_quantize_method", type=str, default="gradient", - choices=["gradient", "kl_div"], + choices=["gradient", "kl_div", "aumann_shapley"], help="[Deprecated: use an AutoQuantize --recipe] Sensitivity scoring method.", ) parser.add_argument( diff --git a/modelopt/recipe/config.py b/modelopt/recipe/config.py index a16cfe4401a..60711c5329c 100644 --- a/modelopt/recipe/config.py +++ b/modelopt/recipe/config.py @@ -254,10 +254,12 @@ class AutoQuantizeConfig(ModeloptBaseConfig): description="Optional per-module overrides for candidate formats and BF16/no-quant " "selectability. Matching is performed after runtime-fusion grouping.", ) - auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField( + auto_quantize_method: Literal["gradient", "kl_div", "aumann_shapley"] = ModeloptField( default="gradient", title="Sensitivity scoring method", - description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).", + description="'gradient' (Taylor + Fisher, needs labels), 'kl_div' (no labels), or " + "'aumann_shapley' (no labels; path-integral damage attributions with a predicted-damage " + "quote).", ) score_size: int = ModeloptField( default=128, diff --git a/modelopt/torch/quantization/_auto_quantize_shapley.py b/modelopt/torch/quantization/_auto_quantize_shapley.py new file mode 100644 index 00000000000..80e3e392d1c --- /dev/null +++ b/modelopt/torch/quantization/_auto_quantize_shapley.py @@ -0,0 +1,1056 @@ +# 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. + +"""Aumann-Shapley sensitivity scoring for AutoQuantize. + +The ``"aumann_shapley"`` method scores each (runtime group, candidate format) pair by how much +damage it causes, measured in nats of KL divergence against the model's own outputs. Because +the model supplies its own reference, scoring needs no labels. The reference keeps any fixed or +forced-single-format groups quantized, so scores are incremental KL relative to the baseline +recorded in ``damage_model["damage_reference"]`` (``{"type": "unquantized"}`` when nothing is +pinned). + +Scoring walks a path from the unquantized model to the quantized one. At each node +``t = (k + 1/2) / num_path_nodes`` every scored module emits ``y + t * (Q(y) - y)`` -- a blend +of its real and quantized outputs, produced by re-running the module with the candidate's +quantizers active -- and one backward pass accumulates ``
`` per (group, +format). Integrating along the path is what makes a KL objective usable: KL against the model's +own outputs sits at an exact zero minimum at the unquantized point, so its gradient vanishes +there and scoring at that single point would carry no signal. + +Cost per batch is one reference forward, one forward with every group at its most aggressive +format, and one forward+backward per (format, path node) -- independent of how many +configurations the solver later considers. + +Those raw attributions become solver scores by fitting them to a directly measured calibration +point: the damage of running every group at its most aggressive format. The fit reproduces that +measurement through ``damage = c * (1 - exp(-sum(b)))``, and the resulting per-group values +``b`` are written to ``candidate_stats["scores"]``, so the standard solve gives the best +allocation under this model and the chosen recipe carries a ``predicted_damage`` estimate in +the same measured units. Scores are additionally adjusted so a more aggressive format never +scores better than a less aggressive one, which keeps estimates conservative +(``b_unprojected`` retains the unadjusted values). ``predicted_damage`` is an estimate from +this model, not a bound on realized deployment KL; ``damage_model["valid"]``, +``approximation_flags`` and ``completeness`` record how far to trust it. + +An efficient implementation of the estimator in https://arxiv.org/abs/2607.12266, validated +empirically against it. + +Method-specific ``method_options``: + +- ``num_path_nodes`` (default 1): quadrature nodes for the path integral. +- ``damage_link`` (default ``"coverage"``): ``"coverage"`` or ``"additive"`` (raw scores). +- ``solver`` (default ``"lp"``): budget-mode solver -- ``"lp"`` (default, exact) or ``"dp"`` + (deterministic, grid-approximate). +- ``max_predicted_damage`` (default None): minimize weight cost subject to predicted damage <= + this bound (mutually exclusive with an ``effective_bits`` constraint). Always solved on + the grid-approximate DP path with conservative rounding, regardless of ``solver``. +""" + +import gc +import math +import types + +import numpy as np +import torch + +from modelopt.torch.opt.searcher import SearchConfig, SearchStateDict +from modelopt.torch.opt.utils import named_hparams +from modelopt.torch.utils import ( + create_param_grad_clear_hook, + print_rank_0, + report_memory, + warn_rank_0, +) +from modelopt.torch.utils.distributed import DistributedProcessGroup + +from .algorithms import ( + AUTO_QUANTIZE_SEARCHERS, + AutoQuantizeGradientSearcher, + QuantRecipe, + QuantRecipeHparam, + _AutoQuantizeBaseSearcher, + _get_kl_div_loss, + _get_lm_head, + _get_log_prob, +) + +__all__ = ["AutoQuantizeAumannShapleySearcher"] + +_DP_GRID = 4096 + + +def _as_seed_coverage(attributions, c, *, p=0.5, iters=200, tol=1e-10, damping=0.5): + """Invert ``AS_i = c a_i prod_{j != i} (1 - p a_j)`` for break rates given ceiling ``c``. + + Returns ``(a, b = -log(1 - a), converged)``. Zero attributions map to exactly zero break + rates; arbitrarily small positive attributions map to proportionally small ones. The + system is infeasible when the attribution mass is too large for the ceiling; the + iteration then diverges to the clip and the caller should inflate ``c`` and retry (see + :func:`_anchor_ceiling`). + """ + attributions = np.maximum(np.asarray(attributions, dtype=float), 0.0) + if not (attributions > 0).any(): + zeros = np.zeros_like(attributions) + return zeros, zeros.copy(), True + a = np.clip(attributions / max(c, 1e-12), 0.0, 0.999) + converged = False + for _ in range(iters): + log_prod = np.log(np.clip(1.0 - p * a, 1e-9, None)).sum() + discount = np.exp(log_prod - np.log(np.clip(1.0 - p * a, 1e-9, None))) + new = np.clip(attributions / (max(c, 1e-12) * np.clip(discount, 1e-6, None)), 0.0, 0.999) + nxt = damping * a + (1.0 - damping) * new + delta = float(np.abs(nxt - a).max()) + a = nxt + if delta < tol: + converged = bool(a.max() < 0.995) + break + b = -np.log(np.clip(1.0 - a, 1e-9, None)) + return a, b, converged + + +def _anchor_ceiling(as_by_key, f_corner, corner_mask_by_key, max_inflation=10.0): + """Invert per-format attributions into break rates with one corner-anchored ceiling. + + The ceiling starts at the measured corner damage and is inflated minimally until the + inversion converges for every format; ``b`` is then rescaled by ``kappa`` so the link stays + exact at the corner. Returns ``(c, b_by_key, kappa, inflation, converged)``. Raises + ``ValueError`` on non-finite inputs (callers must screen measurements first). + """ + if not math.isfinite(f_corner) or not all( + bool(np.isfinite(v).all()) for v in as_by_key.values() + ): + raise ValueError("corner damage and attributions must be finite") + f_corner = max(float(f_corner), 1e-12) + c = f_corner + # The exit conditions bound this well below the cap; the cap is a backstop only. + for _ in range(1 + math.ceil(math.log(max_inflation) / math.log(1.3))): + inversions = {k: _as_seed_coverage(v, c=c) for k, v in as_by_key.items()} + if all(conv for _a, _b, conv in inversions.values()) or c > max_inflation * f_corner: + break + c *= 1.3 + converged = all(conv for _a, _b, conv in inversions.values()) + b_by_key = {k: inv[1] for k, inv in inversions.items()} + + def corner_b_sum(): + return sum(float(b_by_key[k][mask].sum()) for k, mask in corner_mask_by_key.items()) + + kappa = 1.0 + tolerance = 0.01 * f_corner + if abs(_predict_damage(c, corner_b_sum()) - f_corner) > tolerance: + # Exact anchoring needs strict headroom above the corner (kappa solves + # c * (1 - exp(-kappa * sum(b_corner))) == f_corner, impossible at c == f_corner). + if c < 1.01 * f_corner: + c = 1.01 * f_corner + inversions = {k: _as_seed_coverage(v, c=c) for k, v in as_by_key.items()} + converged = all(conv for _a, _b, conv in inversions.values()) + b_by_key = {k: inv[1] for k, inv in inversions.items()} + total = corner_b_sum() + if total > 0: + kappa = -np.log(1.0 - f_corner / c) / total + b_by_key = {k: b * kappa for k, b in b_by_key.items()} + anchored = abs(_predict_damage(c, corner_b_sum()) - f_corner) <= tolerance + return float(c), b_by_key, float(kappa), float(c / f_corner), converged and anchored + + +def _predict_damage(c, b_sum): + return float(c * (1.0 - np.exp(-max(float(b_sum), 0.0)))) + + +def _mckp_max_value(values, costs, budget): + """Exact multiple-choice knapsack: max ``sum(values)`` s.t. ``sum(costs) <= budget``. + + Choice 0 of each row must have zero cost so every row has a feasible pick. + """ + values = np.asarray(values, dtype=float) + costs = np.asarray(costs, dtype=np.int64) + n, num_choices = values.shape + if (costs[:, 0] != 0).any(): + raise ValueError("choice 0 must have zero cost") + grid = int(budget) + dp = np.zeros(grid + 1) + choice = np.zeros((n, grid + 1), dtype=np.int32) + candidates = np.empty((num_choices, grid + 1)) + for i in range(n): + for k in range(num_choices): + cost = int(costs[i, k]) + if cost > grid: + candidates[k] = -np.inf + continue + candidates[k, :cost] = -np.inf + candidates[k, cost:] = dp[: grid + 1 - cost] + values[i, k] + best = candidates.argmax(axis=0) + choice[i] = best + dp = candidates[best, np.arange(grid + 1)] + selection = np.zeros(n, dtype=int) + g = grid + for i in range(n - 1, -1, -1): + k = int(choice[i, g]) + selection[i] = k + g -= int(costs[i, k]) + return selection, float(dp[grid]) + + +class AutoQuantizeAumannShapleySearcher(AutoQuantizeGradientSearcher): + """AutoQuantize searcher scoring with Aumann-Shapley damage attributions (see module doc).""" + + method_name = "aumann_shapley" + method_options_keys = frozenset( + {"num_path_nodes", "damage_link", "solver", "max_predicted_damage"} + ) + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + config = super().default_search_config + config.update( + { + "num_path_nodes": 1, + "damage_link": "coverage", + "solver": "lp", + "max_predicted_damage": None, + } + ) + return config + + @property + def default_state_dict(self) -> SearchStateDict: + """Get the default state dict for AutoQuantize.""" + state = super().default_state_dict + state["damage_model"] = None + state["scoring_signature"] = None + return state + + def _damage_reference(self, no_quant) -> dict: + """The baseline every score, corner, and quote is measured against. + + Groups pinned to one quantized format -- via ``fixed_quantization_config`` or a + single-candidate ``module_search_spaces`` entry with ``allow_no_quant=False`` -- stay + active during the reference passes, so all damage values are INCREMENTAL KL relative + to this resolved baseline, not total degradation from the unquantized model. + """ + forced_groups = { + name: str(stat["formats"][0]) + for name, stat in self.candidate_stats.items() + if len(stat["formats"]) == 1 and stat["formats"][0] != no_quant + } + if getattr(self, "fixed_quantization_config", None) is None and not forced_groups: + return {"type": "unquantized"} + return { + "type": "quantized_baseline", + "fixed_quantization_config_signature": getattr( + self, "fixed_quantization_config_signature", None + ), + "forced_groups": forced_groups, + } + + def _current_scoring_signature(self) -> dict: + # Settings that change what the stored scores MEAN; solver and max_predicted_damage only + # change how they are solved and may differ across a checkpoint resume. + return { + "version": 1, + "path_variant": "module_output_replay_v1", + "num_path_nodes": int(self.config["num_path_nodes"]), + "damage_link": self.config["damage_link"], + } + + def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: + """Sanitize the search config dict.""" + config = config or {} + for ignored_key in ["score_func", "loss_func", "forward_backward_step"]: + if config.get(ignored_key) is not None: + warn_rank_0( + f"`{ignored_key}` is ignored for Aumann-Shapley `auto_quantize`: the loss " + "is fixed to KL divergence against the model's own reference outputs." + ) + config.pop(ignored_key) + config = _AutoQuantizeBaseSearcher.sanitize_search_config(self, config) + assert config["forward_step"] is not None, ( + "`forward_step` must be provided for Aumann-Shapley `auto_quantize`. " + "`forward_step(model, data)` should return model logits." + ) + nodes = config["num_path_nodes"] + if not isinstance(nodes, int) or isinstance(nodes, bool) or nodes < 1: + raise ValueError(f"num_path_nodes must be an integer >= 1, got {nodes!r}") + if config["damage_link"] not in ("coverage", "additive"): + raise ValueError( + f"damage_link must be 'coverage' or 'additive', got {config['damage_link']!r}" + ) + if config["solver"] not in ("lp", "dp"): + raise ValueError(f"solver must be 'lp' or 'dp', got {config['solver']!r}") + bound = config["max_predicted_damage"] + if bound is not None and ( + not isinstance(bound, (int, float)) + or isinstance(bound, bool) + or not math.isfinite(bound) + or bound <= 0 + ): + raise ValueError( + f"max_predicted_damage must be a finite positive number, got {bound!r}" + ) + return config + + def validate_search_input(self, constraints, config) -> None: + """Reject ambiguous target combinations (runs before any model mutation).""" + if ( + config.get("max_predicted_damage") is not None + and (constraints or {}).get("effective_bits") is not None + ): + raise ValueError( + "Provide either constraints['effective_bits'] or " + "method_options['max_predicted_damage'], not both: the damage-bound mode " + "solves for the minimum effective bits itself." + ) + + def before_search(self) -> None: + """Prepare the model for search; damage-bound mode supplies the bit budget itself.""" + # Reject unsupported parallelism before ``super().before_search()`` calibrates every + # search recipe: scoring would raise the same error afterwards, but only after the + # user has paid the full multi-format calibration pass. + self._raise_if_vocab_sharded() + if self.config["max_predicted_damage"] is not None: + self.validate_search_input(self.constraints, self.config) + self.constraints = {"effective_bits": 16.0, **self.constraints} + # Stored scores are only reusable when their meaning is unchanged (see + # _current_scoring_signature); solver/SLA re-solves are allowed on resume. + current_signature = self._current_scoring_signature() + restored_signature = getattr(self, "scoring_signature", None) + if self.candidate_stats and restored_signature not in (None, current_signature): + raise ValueError( + f"Checkpoint scoring signature {restored_signature} does not match the " + f"current search config {current_signature}. Use a different checkpoint path." + ) + self.scoring_signature = current_signature + super().before_search() + + def _configurable_hparams(self) -> list[QuantRecipeHparam]: + return [ + hparam + for _name, hparam in named_hparams(self.model, unique=True) + if isinstance(hparam, QuantRecipeHparam) and hparam.is_configurable + ] + + @torch.enable_grad() + def _estimate_auto_quantize_scores(self, is_param_grad_enabled): + model = self.model + no_quant = QuantRecipe(quant_cfg=None) + lm_head = _get_lm_head(model) + self._raise_if_vocab_sharded() + num_nodes = int(self.config["num_path_nodes"]) + + hparams = self._configurable_hparams() + recipes = sorted({r for h in hparams for r in h.choices if r != no_quant}) + + self._as_recipe: QuantRecipe | None = None + self._as_t: float = 0.0 + self._corner_kl_sum: torch.Tensor | None = None + self._score_tokens: int = 0 + + def set_all_hparams(recipe_of) -> None: + for hparam in hparams: + hparam.active = recipe_of(hparam) + + def score_estimate_forward(module, input, *args, **kwargs): + recipe = self._as_recipe + if recipe is None: + # Reference/corner passes manage hparam state in the outer loop. + return module._forward_original(input, *args, **kwargs) + + # Only a pass that can repopulate the cache may clear it. Score modules nest + # (routed experts score at ``...mlp`` while shared experts inside that same + # mlp score at themselves), and the outer module's replay loop below re-enters + # the inner module's forward under ``no_grad``; clearing unconditionally here + # would discard the diffs the inner module captured during the base pass and + # leave it with a zero score. + grad_pass = torch.is_grad_enabled() + if grad_pass: + module._as_diffs = None + for hparam in module._hparams_for_scoring: + if hparam.is_configurable: + hparam.active = no_quant + output = module._forward_original(input, *args, **kwargs) + base = output[0] if isinstance(output, tuple) else output + + diffs: dict[QuantRecipeHparam, torch.Tensor] = {} + diff_total = None + with torch.no_grad(): + for hparam in module._hparams_for_scoring: + if not hparam.is_configurable or recipe not in hparam.choices: + continue + hparam.active = recipe + quant_output = module._forward_original(input, *args, **kwargs) + hparam.active = no_quant + quant_output = ( + quant_output[0] if isinstance(quant_output, tuple) else quant_output + ) + diff = (quant_output - base).detach() + diffs[hparam] = diff + diff_total = diff if diff_total is None else diff_total + diff + + if diff_total is None: + return output + if grad_pass and base.requires_grad: + module._as_diffs = diffs + # The shift must run in BOTH reentrant-checkpointing passes (identical streams); + # only the diff caching above gates on grad being enabled. + shifted = base + self._as_t * diff_total + if isinstance(output, tuple): + return (shifted, *output[1:]) + return shifted + + def backward_hook(module, grad_input, grad_output): + # Consume-then-clear keeps modules shared across checkpoint segments correct; a + # module invoked twice in ONE forward keeps only its last invocation's diffs. + diffs = getattr(module, "_as_diffs", None) + module._as_diffs = None + recipe = self._as_recipe + if not diffs or recipe is None or grad_output[0] is None: + return + with torch.no_grad(): + grad = grad_output[0].float() + for hparam, diff in diffs.items(): + contribution = (grad * diff.float()).sum() / num_nodes + if hparam._importance_dict[recipe][module] is None: + hparam._importance_dict[recipe][module] = contribution + else: + hparam._importance_dict[recipe][module] += contribution + + def setup_params_for_score_estimation(name, param, params_metadata, enable_grad=True): + params_metadata[name] = {"requires_grad": param.requires_grad} + param.requires_grad = enable_grad + if not enable_grad: + return + accum_grad, handle = create_param_grad_clear_hook(param) + params_metadata[name]["accum_grad"] = accum_grad + params_metadata[name]["handle"] = handle + + def setup_module_for_score_estimation(module): + module._forward_original = module.forward + module.forward = types.MethodType(score_estimate_forward, module) + module._backward_hook_handle = module.register_full_backward_hook(backward_hook) + + def cleanup_module_after_score_estimation(module): + module.forward = module._forward_original + del module._forward_original + module._backward_hook_handle.remove() + if hasattr(module, "_as_diffs"): + del module._as_diffs + + def cleanup_params_after_score_estimation(name, param, params_metadata): + param.requires_grad = params_metadata[name]["requires_grad"] + handle = params_metadata[name].get("handle") + if handle is not None: + handle.remove() + + score_modules = [] + seen: set[int] = set() + for _name, module in model.named_modules(): + if ( + hasattr(module, "_hparams_for_scoring") + and any(h.is_configurable for h in module._hparams_for_scoring) + and id(module) not in seen + ): + setup_module_for_score_estimation(module) + score_modules.append(module) + seen.add(id(module)) + + params_metadata: dict = {} + for name, param in model.named_parameters(): + setup_params_for_score_estimation( + name, param, params_metadata, is_param_grad_enabled(name, model) + ) + + gc.collect() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + report_memory("AutoQuantize(aumann_shapley): starting score estimation, ") + + def score_step(model_, data): + self._as_recipe = None + set_all_hparams(lambda _h: no_quant) + with torch.no_grad(): + ref_logits = self.config["forward_step"](model_, data) + ref_logprob = _get_log_prob(ref_logits, lm_head=lm_head).detach() + self._score_tokens += int(ref_logprob.numel() // ref_logprob.shape[-1]) + del ref_logits + + set_all_hparams(lambda h: h.choices[0]) + corner_logits = self.config["forward_step"](model_, data) + corner_loss = _get_kl_div_loss(ref_logprob, corner_logits, lm_head).detach() + self._corner_kl_sum = ( + corner_loss + if self._corner_kl_sum is None + else self._corner_kl_sum + corner_loss + ) + set_all_hparams(lambda _h: no_quant) + del corner_logits + + for recipe in recipes: + for node in range(num_nodes): + self._as_t = (node + 0.5) / num_nodes + self._as_recipe = recipe + logits = self.config["forward_step"](model_, data) + loss = _get_kl_div_loss(ref_logprob, logits, lm_head) + loss.backward() + del logits, loss + self._as_recipe = None + + try: + self._run_func( + score_step, + num_iters=self.config["num_score_steps"], + desc="Estimating aumann_shapley scores", + ) + finally: + for module in score_modules: + cleanup_module_after_score_estimation(module) + for name, param in model.named_parameters(): + cleanup_params_after_score_estimation(name, param, params_metadata) + del params_metadata + gc.collect() + + if torch.cuda.is_available(): + report_memory("AutoQuantize(aumann_shapley): after score estimation") + + def _loss_is_vocab_sharded(self) -> bool: + lm_head = _get_lm_head(self.model) + parallel_state = getattr(lm_head, "parallel_state", None) if lm_head is not None else None + return parallel_state is not None and parallel_state.tensor_parallel_group.is_initialized() + + def _raise_if_vocab_sharded(self) -> None: + # The score passes backprop through the KL loss; the vocab-sharded log-softmax + # uses in-place collectives that autograd cannot differentiate through. + if self._loss_is_vocab_sharded(): + raise NotImplementedError( + "aumann_shapley scoring does not support vocab-sharded (Megatron " + "tensor-parallel) losses yet. Use method='gradient' with a Megatron loss_func." + ) + + def _reduce_loss_scalar(self, value: float) -> float: + # A loss scalar is sharded over DP (disjoint batches) and, only when the loss itself + # is vocab-sharded, over TP; it is REPLICATED across EP ranks (unlike per-module + # importances, which get_score sums over all three groups). + module = self._any_score_parallel_module() + if module is None: + return value + parallel_state = module.parallel_state + sum_groups = [parallel_state.data_parallel_group] + if self._loss_is_vocab_sharded(): + sum_groups.append(parallel_state.tensor_parallel_group) + value = DistributedProcessGroup.get_dist_syncd_obj(value, sum_groups, sum) + return DistributedProcessGroup.get_dist_syncd_obj( + value, [parallel_state.expert_model_parallel_group], lambda a: a[0] + ) + + def _reduce_token_count(self, count: int) -> int: + module = self._any_score_parallel_module() + if module is None: + return count + return DistributedProcessGroup.get_dist_syncd_obj( + count, [module.parallel_state.data_parallel_group], sum + ) + + def _any_score_parallel_module(self): + for hparam in self._configurable_hparams(): + for module in hparam.score_modules: + if getattr(module, "parallel_state", None) is not None: + return module + return None + + def _exclude_non_finite_candidates( + self, no_quant + ) -> tuple[dict[str, list[str]], dict[str, str], bool]: + """Drop candidates whose measured attribution is non-finite. + + A non-finite score is a measurement of a format that destroys the reference output + (or of a numerically broken pass); zeroing or clamping it would make that candidate + look cheap to the solver, so it is removed from its group's ladder instead. When a + constraint is only reachable through removed candidates, the solve reports + ``is_satisfied=False``. Returns ``(excluded, forced, corner_removed)``: + + - ``excluded``: the removed format labels per group. + - ``forced``: groups with neither a finite candidate nor a no-quant fallback, mapped + to their retained format -- the least aggressive entry, kept as the forced choice + with a neutral solver score (the non-finite raw measurement is preserved). The + caller reports such searches unsatisfied and invalidates the damage model. + - ``corner_removed``: True when some group lost its most aggressive format. The + measured corner ran with that format active, so the anchor no longer corresponds + to the candidate corner and the caller invalidates the coverage fit. + """ + excluded: dict[str, list[str]] = {} + forced: dict[str, str] = {} + corner_removed = False + for name, stat in self.candidate_stats.items(): + if stat.get("is_fixed", False) or len(stat["formats"]) <= 1: + continue + keep = [ + index + for index, (recipe, raw) in enumerate( + zip(stat["formats"], stat["raw_scores"], strict=True) + ) + if recipe == no_quant or math.isfinite(raw) + ] + if len(keep) == len(stat["formats"]): + continue + if not keep: + keep = [len(stat["formats"]) - 1] + stat["scores"][keep[0]] = 0.0 + forced[name] = str(stat["formats"][keep[0]]) + if keep[0] != 0: + corner_removed = True + excluded[name] = [ + str(stat["formats"][index]) + for index in range(len(stat["formats"])) + if index not in keep + ] + for field in ("formats", "scores", "raw_scores", "costs"): + stat[field] = [stat[field][index] for index in keep] + # The base's running-min chain propagates an excluded candidate's non-finite + # score into every less aggressive entry, no_quant included (``min(0.0, -inf)``). + # A group left with no quantized candidate is dropped from the solver tables, so + # the coverage projection never rewrites these; without the reset a -inf reaches + # the LP objective. (+inf and nan already collapse to 0.0 through ``min``.) + stat["scores"] = [score if math.isfinite(score) else 0.0 for score in stat["scores"]] + if excluded: + warn_rank_0( + "aumann_shapley: excluding candidates with non-finite damage measurements " + f"from the search space: {excluded}" + ) + return excluded, forced, corner_removed + + def initialize_candidate_stats(self): + """Initialize candidate stats, then convert raw attributions through the damage link. + + The base implementation performs the distributed score reduction; the nonlinear + coverage inversion runs after it, on rank-identical values, and overwrites the + per-choice scores with log-headroom ``b`` so that minimizing their sum under the + standard solve is the coverage-optimal allocation. + """ + super().initialize_candidate_stats() + + no_quant = QuantRecipe(quant_cfg=None) + # Scoring-time semantics must be captured before pruning rewrites the ladders: the + # reference baseline, and which groups were configurable when scores were measured. + damage_reference = self._damage_reference(no_quant) + eligible = [ + name + for name, stat in self.candidate_stats.items() + if not stat.get("is_fixed", False) + and len(stat["formats"]) > 1 + and any(r != no_quant for r in stat["formats"]) + ] + excluded, forced_candidates, corner_removed = self._exclude_non_finite_candidates(no_quant) + tokens = self._reduce_token_count(int(getattr(self, "_score_tokens", 0))) + if tokens <= 0: + warn_rank_0("aumann_shapley: no scored tokens; leaving raw scores in place.") + return + corner_kl_sum = getattr(self, "_corner_kl_sum", None) + corner_kl_sum = 0.0 if corner_kl_sum is None else float(corner_kl_sum.item()) + f_corner = self._reduce_loss_scalar(corner_kl_sum) / tokens + + # A group that became a singleton through pruning stays in the fit (its remaining + # candidate is still a normalized decision the quote must account for); groups whose + # only remaining measurement is non-finite are forced and cannot be fitted. + names = [ + name + for name in eligible + if name not in forced_candidates + and any(r != no_quant for r in self.candidate_stats[name]["formats"]) + ] + # Equal formats can carry different auto-generated display names, so all tables are + # keyed by the canonical config signature; one representative provides the label. + recipe_by_key: dict[str, QuantRecipe] = {} + for name in names: + for recipe in self.candidate_stats[name]["formats"]: + if recipe != no_quant: + recipe_by_key.setdefault(recipe.checkpoint_signature, recipe) + keys = sorted(recipe_by_key) + labels = {key: str(recipe_by_key[key]) for key in keys} + + signed_by_key = {key: np.zeros(len(names)) for key in keys} + ladders = set() + for i, name in enumerate(names): + stat = self.candidate_stats[name] + ladder = [] + for recipe, raw_score in zip(stat["formats"], stat["raw_scores"], strict=True): + if recipe == no_quant: + continue + key = recipe.checkpoint_signature + # Raw (unclamped) values: the base's monotonicity clamp must not leak into + # the fit or the diagnostics. + signed_by_key[key][i] = raw_score / tokens + ladder.append(key) + ladders.add(tuple(ladder)) + heterogeneous = len(ladders) > 1 + if heterogeneous: + warn_rank_0( + "aumann_shapley: groups have differing candidate ladders; the joint coverage " + "interpretation is approximate for the formats not shared by all groups." + ) + as_by_key = {key: np.maximum(vector, 0.0) for key, vector in signed_by_key.items()} + # Candidate-level non-finite measurements were excluded from the ladders above; a + # non-finite corner cannot be localized to one candidate, so it invalidates the fit. + finite = math.isfinite(f_corner) + signed_total = sum(float(np.abs(v).sum()) for v in signed_by_key.values()) + negative_mass = sum( + float(np.abs(np.minimum(v, 0.0)).sum()) for v in signed_by_key.values() + ) / max(signed_total, 1e-12) + + corner_mask_by_key = {key: np.zeros(len(names), dtype=bool) for key in keys} + for i, name in enumerate(names): + corner = self.candidate_stats[name]["formats"][0] + # QuantRecipe.__lt__ pins no_quant last among equal-compression recipes, so the + # most aggressive entry is always a quantized format with a key in ``keys`` + # (built from quantized formats only) and matches what the corner forward ran. + assert not corner.is_no_quant, f"no_quant sorted first in the ladder for {name}" + corner_mask_by_key[corner.checkpoint_signature][i] = True + # Mathematical completeness concerns the SIGNED attribution sum; the clamp applied + # for the solver is reported separately as negative_attribution_mass. + corner_mass = sum( + float(signed_by_key[key][mask].sum()) for key, mask in corner_mask_by_key.items() + ) + + link = self.config["damage_link"] + valid = True + flags: list[str] = [] + if not finite: + flags.append("non_finite_measurements") + valid = False + if excluded: + flags.append("non_finite_scores_excluded") + if corner_removed: + # f_corner was measured with the original most-aggressive formats active; the + # anchor no longer describes the corner of the pruned candidate space. + flags.append("corner_format_excluded") + valid = False + if forced_candidates: + # Some group had neither a finite candidate nor a no-quant fallback: its damage + # is real but unquantifiable, so no quote derived from these scores is reliable. + flags.append("non_finite_candidate_forced") + valid = False + if heterogeneous: + flags.append("heterogeneous_ladders") + if negative_mass > 1e-3: + flags.append("negative_attribution_mass") + damage_model: dict = { + "link": link, + "f_corner": f_corner, + "n_score_tokens": tokens, + "damage_reference": damage_reference, + "negative_attribution_mass": negative_mass, + "as_scores": { + labels[key]: dict(zip(names, vector.tolist())) + for key, vector in signed_by_key.items() + }, + } + if excluded: + damage_model["excluded_candidates"] = excluded + if forced_candidates: + damage_model["forced_candidates"] = forced_candidates + + zero_tolerance = 1e-12 + positive_mass = sum(float(vector.sum()) for vector in as_by_key.values()) + coverage = link == "coverage" and bool(names) and bool(keys) + if coverage: + if not finite or corner_removed: + # Without a usable corner (non-finite, or measured with a since-pruned + # format) the inversion would anchor to an unrelated measurement; the + # normalized attributions themselves are the honest solver objective. + c, kappa, inflation, converged = 0.0, 1.0, 1.0, False + score_by_key = as_by_key + b_by_key = {key: np.zeros(len(names)) for key in keys} + elif f_corner <= zero_tolerance and positive_mass <= zero_tolerance: + # Quantization is measurably free here: an exact zero-damage model. + flags.append("zero_damage") + c, kappa, inflation, converged = 0.0, 1.0, 1.0, True + score_by_key = {key: np.zeros(len(names)) for key in keys} + b_by_key = score_by_key + elif f_corner <= zero_tolerance: + # Attributions claim damage the corner measurement does not show: the + # coverage fit is not meaningful. + flags.append("zero_corner_with_attribution_mass") + valid = False + c, kappa, inflation, converged = 0.0, 1.0, 1.0, False + score_by_key = as_by_key + b_by_key = {key: np.zeros(len(names)) for key in keys} + else: + c, b_by_key, kappa, inflation, converged = _anchor_ceiling( + as_by_key, f_corner, corner_mask_by_key + ) + valid = valid and bool(converged) and math.isfinite(c) + score_by_key = b_by_key + if not valid: + warn_rank_0( + "aumann_shapley: the coverage damage fit is not valid " + f"(flags={flags or ['inversion_not_converged']}); damage quotes are " + "unreliable and damage-bound searches will report is_satisfied=False." + ) + else: + score_by_key = as_by_key + + # Project solver scores onto the monotone ladder (more aggressive => at least as + # much damage) by raising the more aggressive entries: conservative, and never + # erases a less-aggressive format's real damage. + projected_by_key = {key: np.zeros(len(names)) for key in keys} + for i, name in enumerate(names): + stat = self.candidate_stats[name] + scores = [ + 0.0 if recipe == no_quant else float(score_by_key[recipe.checkpoint_signature][i]) + for recipe in stat["formats"] + ] + for k in range(len(scores) - 2, -1, -1): + scores[k] = max(scores[k], scores[k + 1]) + stat["scores"] = scores + for recipe, score in zip(stat["formats"], scores, strict=True): + if recipe != no_quant: + projected_by_key[recipe.checkpoint_signature][i] = score + + unprojected_total = sum(float(vector.sum()) for vector in score_by_key.values()) + projected_total = sum(float(vector.sum()) for vector in projected_by_key.values()) + adjustment = (projected_total - unprojected_total) / max(unprojected_total, 1e-12) + if adjustment > 1e-9: + flags.append("monotonicity_projection") + damage_model["monotonicity_adjustment"] = adjustment + + if coverage: + projected_corner_b = sum( + float(projected_by_key[key][mask].sum()) for key, mask in corner_mask_by_key.items() + ) + damage_model.update( + { + "c": c, + "kappa": kappa, + "ceiling_inflation": inflation, + "inversion_converged": bool(converged), + # The quote-operative link values (what the solver and predicted_damage + # use); the unprojected inversion output stays exactly corner-anchored. + "b": { + labels[key]: dict(zip(names, b.tolist())) + for key, b in projected_by_key.items() + }, + "b_unprojected": { + labels[key]: dict(zip(names, b.tolist())) for key, b in b_by_key.items() + }, + "projected_corner_damage": _predict_damage(c, projected_corner_b), + } + ) + + damage_model["valid"] = valid + damage_model["approximation_flags"] = flags + damage_model["completeness"] = corner_mass / max(f_corner, 1e-12) + self.damage_model = damage_model + + def run_search_with_stats(self, max_weight_size, verbose=False): + """Dispatch to the LP (default), the exact DP, or the SLA search.""" + max_predicted_damage = self.config.get("max_predicted_damage") + if max_predicted_damage is not None: + recipes, is_satisfied = self._run_damage_bound_search( + float(max_predicted_damage), verbose + ) + elif self.config.get("solver", "lp") == "dp": + recipes, is_satisfied = self._run_dp_budget_search(max_weight_size, verbose) + else: + recipes, is_satisfied = super().run_search_with_stats(max_weight_size, verbose) + flags = (getattr(self, "damage_model", None) or {}).get("approximation_flags", []) + if is_satisfied and "non_finite_candidate_forced" in flags: + warn_rank_0( + "AutoQuantize FAILED to find a valid solution! The selection includes a " + "forced candidate whose damage measurement was non-finite. " + ) + is_satisfied = False + return recipes, is_satisfied + + def run_search(self): + """Run the inherited search and attach the predicted-damage quote.""" + super().run_search() + self._attach_predicted_damage() + # The base flow only saves before solving; re-save so the checkpoint file carries the + # chosen recipe and can be re-solved offline. + self.save_search_checkpoint(verbose=self.config.get("verbose", False)) + + def _attach_predicted_damage(self) -> None: + damage_model = getattr(self, "damage_model", None) + if not damage_model or not self.best.get("recipe"): + return + no_quant = QuantRecipe(quant_cfg=None) + total_score = 0.0 + for name, recipe in self.best["recipe"].items(): + stat = self.candidate_stats[name] + if stat.get("is_fixed", False) or recipe == no_quant: + continue + total_score += stat["scores"][stat["formats"].index(recipe)] + if damage_model["link"] == "coverage" and "c" in damage_model: + predicted = _predict_damage(damage_model["c"], total_score) + else: + predicted = float(total_score) + self.best["predicted_damage"] = predicted + self.best["predicted_damage_valid"] = bool(damage_model.get("valid", True)) + if self.config.get("verbose"): + print_rank_0( + f"AutoQuantize(aumann_shapley) predicted damage: {predicted:.4e} " + "(mean per-token KL, calibration units)" + ) + + def _stat_tables(self): + stats = self.candidate_stats + names = list(stats) + n = len(names) + num_choices = max(len(stats[name]["formats"]) for name in names) + scores = np.full((n, num_choices), np.inf) + costs = np.full((n, num_choices), np.inf) + lengths = np.zeros(n, dtype=int) + uncompressed = np.zeros(n) + for i, name in enumerate(names): + stat = stats[name] + k = len(stat["formats"]) + lengths[i] = k + scores[i, :k] = stat["scores"] + costs[i, :k] = stat["costs"] + uncompressed[i] = stat.get("uncompressed_cost", max(stat["costs"])) + return names, scores, costs, lengths, uncompressed + + def _best_recipes_from_choice(self, names, choice_idx): + best_recipes = {} + for name, k in zip(names, choice_idx, strict=True): + stat = self.candidate_stats[name] + best_recipes[name] = { + "format": stat["formats"][int(k)], + "costs": stat["costs"][int(k)], + "scores": stat["scores"][int(k)], + } + return best_recipes + + @staticmethod + def _min_score_choices(scores, lengths): + # Ties resolve toward the least aggressive choice (scores are non-increasing along + # ascending compression, so scan from the end). + return np.array( + [ + int(lengths[i]) - 1 - int(np.argmin(scores[i, : lengths[i]][::-1])) + for i in range(len(lengths)) + ] + ) + + @staticmethod + def _minimize_within_budget(objective, constraint, budget): + """Per-row choices minimizing ``sum(objective)`` s.t. ``sum(constraint) <= budget``. + + Each row's true minimum constraint is subtracted before discretizing, so only the + increments above the mandatory baseline consume grid resolution (rows never overflow + the grid regardless of their count) and ceil rounding keeps the returned selection + feasible on the true budget. Returns None when even the per-row minima exceed it. + """ + n = len(objective) + row_min = np.where(np.isfinite(constraint), constraint, np.inf).min(axis=1) + remaining = float(budget - row_min.sum()) + if remaining < -1e-9 * max(abs(budget), 1.0): + return None + if remaining <= 0: + # No slack above the mandatory baseline: take each row's minimum-constraint + # column, breaking ties toward the lower objective. + at_minimum = constraint <= row_min[:, None] + return np.where(at_minimum, objective, np.inf).argmin(axis=1) + + quantum = remaining / _DP_GRID + increments = np.where( + np.isfinite(constraint), + np.ceil((constraint - row_min[:, None]) / quantum), + _DP_GRID + 1, + ).astype(np.int64) + # Reorder each row so its zero-increment column is first (DP precondition). + order = np.argsort(increments, axis=1, kind="stable") + rows = np.arange(n)[:, None] + values = np.where(np.isfinite(objective[rows, order]), -objective[rows, order], -np.inf) + selection, _ = _mckp_max_value(values, increments[rows, order], _DP_GRID) + return order[np.arange(n), selection] + + def _run_dp_budget_search(self, max_weight_size, verbose=False): + """Deterministic grid-approximate knapsack solve for budget mode (LP alternative).""" + names, scores, costs, lengths, _uncompressed = self._stat_tables() + n = len(names) + + choice_idx = self._minimize_within_budget(scores, costs, float(max_weight_size)) + if choice_idx is None: + warn_rank_0( + "AutoQuantize FAILED to find a solution! The searched model might not meet " + "all constraints. " + ) + # Best effort: each row's minimum true cost, ties toward the lower score. + at_minimum = costs <= np.where(np.isfinite(costs), costs, np.inf).min( + axis=1, keepdims=True + ) + choice_idx = np.where(at_minimum, scores, np.inf).argmin(axis=1) + realized = float(costs[np.arange(n), choice_idx].sum()) + is_satisfied = realized <= max_weight_size * (1 + 1e-12) + 1e-9 + if verbose: + print_rank_0( + f"AutoQuantize(dp): realized weight size {realized:.2f} " + f"(target {max_weight_size:.2f}), satisfied={is_satisfied}" + ) + return self._best_recipes_from_choice(names, choice_idx), is_satisfied + + def _run_damage_bound_search(self, max_predicted_damage, verbose=False): + """Minimize total weight cost subject to predicted damage <= ``max_predicted_damage``. + + Under the coverage link the bound maps to a score budget ``-log(1 - eps/c)``; score + costs are rounded UP on the constraint axis so the quote is certified. + """ + damage_model = getattr(self, "damage_model", None) or {} + link = damage_model.get("link", self.config.get("damage_link", "additive")) + names, scores, costs, lengths, _uncompressed = self._stat_tables() + n = len(names) + + if link == "coverage" and not damage_model.get("valid", False): + warn_rank_0( + "AutoQuantize FAILED to find a solution! The coverage damage fit is invalid " + f"(flags={damage_model.get('approximation_flags')}), so the damage bound " + "cannot be certified. Returning the minimum-damage configuration. " + ) + choice_idx = self._min_score_choices(scores, lengths) + return self._best_recipes_from_choice(names, choice_idx), False + + if link == "coverage" and "c" in damage_model: + c = float(damage_model["c"]) + budget = ( + np.inf if max_predicted_damage >= c else -math.log(1.0 - max_predicted_damage / c) + ) + else: + budget = float(max_predicted_damage) + + if not np.isfinite(budget): + choice_idx = np.zeros(n, dtype=int) + else: + choice_idx = self._minimize_within_budget(costs, scores, budget) if budget > 0 else None + if choice_idx is None: + warn_rank_0( + "AutoQuantize FAILED to find a solution! Even the least aggressive " + "choices exceed the damage budget. " + ) + choice_idx = self._min_score_choices(scores, lengths) + + total_score = float(scores[np.arange(n), choice_idx].sum()) + is_satisfied = (not np.isfinite(budget)) or bool(total_score <= budget + 1e-12) + if verbose: + realized = float(costs[np.arange(n), choice_idx].sum()) + print_rank_0( + f"AutoQuantize(sla): score total {total_score:.4e} (budget {budget:.4e}), " + f"weight size {realized:.2f}, satisfied={is_satisfied}" + ) + return self._best_recipes_from_choice(names, choice_idx), is_satisfied + + +AUTO_QUANTIZE_SEARCHERS[AutoQuantizeAumannShapleySearcher.method_name] = ( + AutoQuantizeAumannShapleySearcher +) diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 7beeef6ad7f..55aa95474df 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -17,6 +17,7 @@ import copy import fnmatch +import functools import gc import types import warnings @@ -268,6 +269,12 @@ def estimate_quant_compression_for_quantizer(quantizer_attr_cfg): return estimate_quant_compression_for_quantizer(cfgs) if cfgs else 1.0 +@functools.cache +def _no_quant_signature() -> str: + """Canonical signature of the no-quant recipe, used to pin it last in the format ladder.""" + return QuantRecipe(quant_cfg=None).checkpoint_signature + + class QuantRecipe(CustomHPType): """A subclass of QuantizeConfig enabling auto_quantize specific configurations. @@ -308,6 +315,11 @@ def checkpoint_signature(self) -> str: """Return the canonical identity used for ordering and checkpoint validation.""" return getattr(self, "_config_signature", self.config.model_dump_json()) + @property + def is_no_quant(self) -> bool: + """Whether this recipe leaves the module unquantized.""" + return self.checkpoint_signature == _no_quant_signature() + @staticmethod def get_auto_name_for_config(quant_cfg: str | dict[str, Any] | None) -> str | None: """Get a name for the quantization configuration.""" @@ -332,8 +344,15 @@ def __repr__(self) -> str: return self._str_repr def __lt__(self, other: "QuantRecipe"): - return (self.compression, self.checkpoint_signature) < ( + # no_quant's compression is 1.0, which a weight-preserving candidate (activation- or + # KV-cache-only) ties exactly; without an explicit tiebreak the config-JSON comparison + # decides, and no_quant can land anywhere in the ladder. Downstream code reads + # ``choices[0]`` / ``formats[0]`` as the most aggressive candidate and treats the last + # entry as the unquantized end, so pin no_quant last among equal-compression recipes. + # Ordering is unchanged whenever every candidate compresses weights. + return (self.compression, self.is_no_quant, self.checkpoint_signature) < ( other.compression, + other.is_no_quant, other.checkpoint_signature, ) @@ -595,6 +614,7 @@ 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)$") +_MLA_A_PROJ_RE = re.compile(r"^(.*?)\.(?:q_a_proj|kv_a_proj_with_mqa)$") def _linear_attn_qkvz_group_key(_model, name: str) -> str | None: @@ -607,6 +627,11 @@ def _linear_attn_ba_group_key(_model, name: str) -> str | None: return f"{m.group(1)}/ba" if m else None +def _mla_a_proj_group_key(_model, name: str) -> str | None: + m = _MLA_A_PROJ_RE.match(name) + return f"{m.group(1)}/qkv_a" if m else None + + def _module_search_space_signature(module_search_spaces) -> tuple: """Return a checkpoint-stable description of module-specific candidate spaces.""" return tuple( @@ -632,10 +657,12 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): # certain modules to share the same format. Sensitivity scores are computed from perturbations # at score modules. See AutoQuantizeGradientSearcher for detailed documentation. - candidate_stats: dict[str, dict[str, list[float]]] + candidate_stats: dict[str, dict[str, Any]] best: dict[str, Any] quantizer_states: dict method_name: str | None = None + # Config keys settable through ``auto_quantize(method_options=...)``. + method_options_keys: frozenset[str] = frozenset() quant_grouping_rules = [ r"^(.*?)\.(q_proj|k_proj|v_proj)$", # q_proj, k_proj, v_proj for llama like models @@ -647,6 +674,14 @@ 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 + # MLA low-rank input projections (DeepSeek/GLM lineage): TRT-LLM fuses them into + # fused_qkv_a_proj_with_mqa, so the shards must share one quantization format. + # A callable (not a regex) because a regex rule keys on match.group(1) -- the + # attention path -- which is the key the q_proj/k_proj/v_proj rule above already + # returns. On MLA built without a q LoRA rank (q_lora_rank=None, e.g. + # DeepSeek-V2-Lite) ``q_proj`` and ``kv_a_proj_with_mqa`` are siblings, so a + # shared key would merge the unfused q_proj into this group. + _mla_a_proj_group_key, # 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 @@ -709,6 +744,9 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: ) return config + def validate_search_input(self, constraints, config) -> None: + """Hook for cross-field input validation before the model is converted.""" + def load_search_checkpoint(self) -> bool: return super().load_search_checkpoint(strict=False) @@ -1056,7 +1094,7 @@ def initialize_candidate_stats(self): if not isinstance(hparam, QuantRecipeHparam): continue - formats, scores, costs = [], [], [] + formats, raw_scores, scores, costs = [], [], [], [] prev_score = float("inf") for recipe in hparam.solver_choices: formats.append(recipe) @@ -1064,12 +1102,16 @@ def initialize_candidate_stats(self): score = hparam.get_score(recipe) cost = hparam.get_cost(recipe) + raw_scores.append(score) score = min(score, prev_score) # TODO: Should we get rid of this? scores.append(score) costs.append(cost) prev_score = score self.candidate_stats[name]["formats"] = formats + # Unclamped values, aligned with ``formats``, for method-specific fitting and + # diagnostics (the reduction in get_score is a collective; run it only once). + self.candidate_stats[name]["raw_scores"] = raw_scores self.candidate_stats[name]["scores"] = scores self.candidate_stats[name]["costs"] = costs self.candidate_stats[name]["module_names"] = hparam.quant_module_names @@ -1946,6 +1988,13 @@ def run_search_with_stats(self, max_weight_size, verbose=False): # Backward compatibility alias (defaults to gradient-based searcher) AutoQuantizeSearcher = AutoQuantizeGradientSearcher +# Registry of auto_quantize sensitivity-scoring methods. Additional methods register +# themselves here on import (see e.g. _auto_quantize_shapley). +AUTO_QUANTIZE_SEARCHERS: dict[str, type[_AutoQuantizeBaseSearcher]] = { + AutoQuantizeGradientSearcher.method_name: AutoQuantizeGradientSearcher, + AutoQuantizeKLDivSearcher.method_name: AutoQuantizeKLDivSearcher, +} + def _as_list(value) -> list: if value is None: @@ -2078,14 +2127,12 @@ def _resolve_best_recipe(search_state, constraints, verbose=False): max_weight_size = total_weight_size * compression method = search_state["method"] - if method == "gradient": - searcher = AutoQuantizeGradientSearcher() - elif method == "kl_div": - searcher = AutoQuantizeKLDivSearcher() - else: + if method not in AUTO_QUANTIZE_SEARCHERS: raise ValueError( - f"Unknown autoquant search method: {method!r}. Expected 'gradient' or 'kl_div'." + f"Unknown autoquant search method: {method!r}. " + f"Expected one of {sorted(AUTO_QUANTIZE_SEARCHERS)}." ) + searcher = AUTO_QUANTIZE_SEARCHERS[method]() searcher.candidate_stats = candidate_stats searcher.cost_model = search_state.get("cost_model", COST_MODEL_WEIGHT) @@ -2103,6 +2150,11 @@ def _resolve_best_recipe(search_state, constraints, verbose=False): "cost": searcher.cost, "active_moe_expert_ratio": searcher.active_moe_expert_ratio, } + # Method-specific state (e.g. the aumann_shapley damage model) participates in the + # re-solve; restore whatever the searcher declares beyond the fields set above. + for key in searcher.default_state_dict: + if key in search_state and not hasattr(searcher, key): + setattr(searcher, key, search_state[key]) best_recipe_info, _ = searcher.run_search_with_stats(max_weight_size, verbose=verbose) best_recipe = {name: info["format"] for name, info in best_recipe_info.items()} diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 966a3643fe3..5c37e75498d 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -36,7 +36,9 @@ ) from modelopt.torch.utils import atomic_print -from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe +# The _auto_quantize_shapley import registers the "aumann_shapley" method on load. +from . import _auto_quantize_shapley # noqa: F401 +from .algorithms import AUTO_QUANTIZE_SEARCHERS, QuantRecipe 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 +284,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, + method_options: dict[str, Any] | None = None, ): r"""Perform optimal per-layer quantization by searching for the best quantization formats per-layer. @@ -439,9 +442,13 @@ def forward_backward_step(model, batch) -> None: verbose: If True, prints the search progress/intermediate results. method: Method to use for estimating sensitivity loss. Higher loss indicates greater sensitivity to quantization. Options are ``"gradient"`` (default; uses gradient-based loss estimation, - linear programming search, and requires ``loss_func`` or ``forward_backward_step``) and + linear programming search, and requires ``loss_func`` or ``forward_backward_step``), ``"kl_div"`` (uses KL divergence between unquantized and quantized outputs, relies on - threshold-based binary search, and only requires ``forward_step`` returning logits). + threshold-based binary search, and only requires ``forward_step`` returning logits), and + ``"aumann_shapley"`` (Aumann-Shapley path-integral damage attributions with a + measured-corner coverage calibration; label-free like ``"kl_div"`` but scores all + formats in a constant number of passes and attaches a ``predicted_damage`` quote -- + see :mod:`modelopt.torch.quantization._auto_quantize_shapley`). 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. @@ -459,6 +466,9 @@ def forward_backward_step(model, batch) -> None: active while searched modules are scored, is calibrated only with its own algorithm, and remains part of the effective-bits numerator and denominator. This is one integrated AutoQuantize operation, not staged PTQ followed by AutoQuantize. + method_options: Optional method-specific settings merged into the searcher config and + validated by the selected method (e.g. ``{"num_path_nodes": 2}`` or + ``{"max_predicted_damage": 1e-3}`` for ``method="aumann_shapley"``). Returns: A tuple (model, state_dict) where ``model`` is the searched and quantized model and ``state_dict`` contains the history and detailed stats of the search procedure. @@ -618,18 +628,12 @@ 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'.") + if method not in AUTO_QUANTIZE_SEARCHERS: + raise ValueError( + f"Invalid method: {method}. Valid options are {sorted(AUTO_QUANTIZE_SEARCHERS)}." + ) + searcher = AUTO_QUANTIZE_SEARCHERS[method]() - model = apply_mode( - model, - mode="auto_quantize", - registry=QuantizeModeRegistry, - ) search_config = { "quantization_formats": processed_quantization_formats, "fixed_quantization_config": processed_fixed_quantization_config, @@ -644,13 +648,37 @@ def _process_quantization_formats(formats, custom_name_prefix): "verbose": verbose, "checkpoint": checkpoint, } + if method_options is not None: + if not isinstance(method_options, dict): + raise TypeError(f"method_options must be a dict, got {type(method_options).__name__}") + # Only the selected method's declared options are accepted; core inputs (loaders, + # steps, checkpoint, ...) cannot be overridden here. + invalid = set(method_options) - searcher.method_options_keys + if invalid: + raise ValueError( + f"Invalid method_options {sorted(invalid)} for method={method!r}. " + f"Supported options: {sorted(searcher.method_options_keys)}." + ) + search_config.update(method_options) + # Validate the full search config (including method-option values and cross-field + # consistency with the constraints) before the model is converted, so a rejected + # configuration leaves the model untouched. The searcher re-sanitizes the + # already-sanitized config inside search(), which is a no-op. + search_config = searcher.sanitize_search_config(search_config) + search_constraints = cast("ConstraintsDict", constraints or {}) + searcher.validate_search_input(search_constraints, search_config) + + model = apply_mode( + model, + mode="auto_quantize", + registry=QuantizeModeRegistry, + ) # Disable all quantizers; AutoQuantize will enable the needed ones set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}]) if processed_fixed_quantization_config is not None: fixed_cfg, fixed_name = processed_fixed_quantization_config fixed_recipe = QuantRecipe(fixed_cfg, name=fixed_name) set_quantizer_by_cfg(model, fixed_recipe.config.quant_cfg) - search_constraints = cast("ConstraintsDict", constraints or {}) searcher.search(model, search_constraints, config=search_config) return model, searcher.state_dict() diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index e83f7fa0a70..f320d55f6c6 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -736,7 +736,7 @@ def test_active_moe_search_prefers_budget_lower_bound(): ) @pytest.mark.parametrize( "method", - ["gradient", "kl_div"], + ["gradient", "kl_div", "aumann_shapley"], ) def test_auto_quantize(model_cls, search_formats, min_bits, search_bits, method): model = model_cls() @@ -1085,7 +1085,7 @@ def test_estimate_quant_compression_per_entry_effective_bits(): ) -@pytest.mark.parametrize("method", ["gradient", "kl_div"]) +@pytest.mark.parametrize("method", ["gradient", "kl_div", "aumann_shapley"]) def test_auto_quantize_checkpoint_resume(method, tmp_path, capsys): """Test that checkpoint can be used to resume an interrupted search.""" model = SimpleLinear() @@ -1536,3 +1536,85 @@ def test_get_auto_quantize_config_emits_fused_expert_quantizer_names(with_persis 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 + + +def test_mla_projections_share_one_group(): + """TRT-LLM fuses q_a_proj + kv_a_proj_with_mqa, so they must share one quant format.""" + + class _MLAAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_a_proj = torch.nn.Linear(32, 32) + self.kv_a_proj_with_mqa = torch.nn.Linear(32, 32) + self.o_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.o_proj(self.q_a_proj(x) + self.kv_a_proj_with_mqa(x)) + + class _MLABlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _MLAAttention() + + def forward(self, x): + return self.self_attn(x) + + def get_input(self): + return torch.randn(1, 4, 32) + + model = _MLABlock() + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=2, + num_score_steps=2, + method="gradient", + ) + hparam = model.self_attn.q_a_proj.get_hparam("quant_recipe") + assert model.self_attn.kv_a_proj_with_mqa.get_hparam("quant_recipe") == hparam + assert model.self_attn.o_proj.get_hparam("quant_recipe") != hparam + + +def test_mla_group_does_not_absorb_unfused_q_proj(): + """With q_lora_rank=None there is no q_a_proj; q_proj is NOT fused with kv_a_proj.""" + + class _MLAAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_proj = torch.nn.Linear(32, 32) + self.kv_a_proj_with_mqa = torch.nn.Linear(32, 32) + self.o_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.o_proj(self.q_proj(x) + self.kv_a_proj_with_mqa(x)) + + class _MLABlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _MLAAttention() + + def forward(self, x): + return self.self_attn(x) + + def get_input(self): + return torch.randn(1, 4, 32) + + model = _MLABlock() + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.sum(), + num_calib_steps=2, + num_score_steps=2, + method="gradient", + ) + q_hparam = model.self_attn.q_proj.get_hparam("quant_recipe") + assert model.self_attn.kv_a_proj_with_mqa.get_hparam("quant_recipe") != q_hparam + assert model.self_attn.o_proj.get_hparam("quant_recipe") != q_hparam diff --git a/tests/unit/torch/quantization/test_autoquant_shapley.py b/tests/unit/torch/quantization/test_autoquant_shapley.py new file mode 100644 index 00000000000..815bc3486b7 --- /dev/null +++ b/tests/unit/torch/quantization/test_autoquant_shapley.py @@ -0,0 +1,967 @@ +# 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. + +"""Tests specific to the aumann_shapley AutoQuantize method. + +Shared behavior (search across models/formats, checkpoint resume) is covered by the method +parametrizations in test_autoquant.py; this module pins the method's own guarantees: config +parity with the standard builder, the damage model, the path-integral completeness property, +SLA certification, and the DP solver. +""" + +import copy +import itertools +import math + +import numpy as np +import pytest +import torch + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization._auto_quantize_shapley import ( + AutoQuantizeAumannShapleySearcher, + _anchor_ceiling, + _as_seed_coverage, + _predict_damage, +) +from modelopt.torch.quantization.algorithms import QuantRecipe + +SEARCH_FORMATS = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG] + + +class _Attention(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_proj = torch.nn.Linear(32, 32) + self.k_proj = torch.nn.Linear(32, 32) + self.v_proj = torch.nn.Linear(32, 32) + self.o_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + for layer in [self.q_proj, self.k_proj, self.v_proj, self.o_proj]: + x = layer(x) + return x + + +class _Block(torch.nn.Module): + def __init__(self, seed=0): + super().__init__() + torch.manual_seed(seed) + self.attn = _Attention() + self.mlp = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.mlp(self.attn(x)) + + def get_input(self): + return torch.randn(1, 4, 32) + + +class _OneLinear(torch.nn.Module): + def __init__(self): + super().__init__() + torch.manual_seed(0) + self.fc = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.fc(x) + + def get_input(self): + return torch.randn(1, 4, 32) + + +def _search(model, method="aumann_shapley", effective_bits=6.0, method_options=None, **kwargs): + return mtq.auto_quantize( + model, + constraints={"effective_bits": effective_bits} if effective_bits is not None else None, + quantization_formats=list(SEARCH_FORMATS), + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + loss_func=(lambda output, data: output.sum()) if method == "gradient" else None, + num_calib_steps=2, + num_score_steps=2, + method=method, + method_options=method_options, + **kwargs, + ) + + +@pytest.fixture(scope="module") +def shapley_state(): + _model, state = _search(_Block(), method_options={"num_path_nodes": 2}) + return state + + +def test_damage_model_and_score_monotonicity(shapley_state): + assert shapley_state["method"] == "aumann_shapley" + assert shapley_state["best"]["is_satisfied"] + assert shapley_state["best"]["predicted_damage"] >= 0 + assert shapley_state["best"]["predicted_damage_valid"] is True + + damage_model = shapley_state["damage_model"] + assert damage_model["link"] == "coverage" + assert damage_model["c"] >= damage_model["f_corner"] > 0 + assert damage_model["completeness"] > 0 + assert damage_model["valid"] is True + assert isinstance(damage_model["approximation_flags"], list) + assert damage_model["damage_reference"] == {"type": "unquantized"} + assert shapley_state["scoring_signature"]["num_path_nodes"] == 2 + + for stat in shapley_state["candidate_stats"].values(): + assert len(stat["formats"]) == len(stat["scores"]) == len(stat["costs"]) + assert all( + stat["scores"][i] >= stat["scores"][i + 1] - 1e-12 + for i in range(len(stat["scores"]) - 1) + ) + + +def test_identical_scores_emit_identical_configs(shapley_state): + """With identical stats and scores, the emitted config must match the gradient method's + dict-for-dict: config emission and solving are shared, only scoring differs.""" + _model, gradient_state = _search(_Block(), method="gradient") + + shapley = copy.deepcopy(shapley_state) + gradient = copy.deepcopy(gradient_state) + assert list(shapley["candidate_stats"]) == list(gradient["candidate_stats"]) + for i, name in enumerate(gradient["candidate_stats"]): + num_choices = len(gradient["candidate_stats"][name]["scores"]) + synthetic = [float(num_choices - j) * (1.0 + 0.1 * i) for j in range(num_choices)] + gradient["candidate_stats"][name]["scores"] = list(synthetic) + shapley["candidate_stats"][name]["scores"] = list(synthetic) + + for bits in (14.0, 9.0, 6.0): + config_gradient = mtq.get_auto_quantize_config(gradient, {"effective_bits": bits}) + config_shapley = mtq.get_auto_quantize_config(shapley, {"effective_bits": bits}) + assert config_shapley == config_gradient, f"configs diverge at effective_bits={bits}" + + +def test_config_applies_and_resolve_tightens(shapley_state): + config = mtq.get_auto_quantize_config(shapley_state) + assert config["algorithm"] == "max" + assert config["quant_cfg"][0] == {"quantizer_name": "*", "enable": False} + + model = _Block(seed=1) + mtq.quantize(model, config, lambda m: m(m.get_input())) + with torch.no_grad(): + model(model.get_input()) + + def enabled_entries(bits): + config = mtq.get_auto_quantize_config(shapley_state, {"effective_bits": bits}) + return sum(1 for entry in config["quant_cfg"] if entry.get("enable")) + + assert enabled_entries(5.0) >= enabled_entries(14.0) + + +def test_completeness_one_group(): + """With a single group the path integral must recover the measured corner damage.""" + _model, state = _search( + _OneLinear(), effective_bits=16.0, method_options={"num_path_nodes": 32} + ) + assert state["damage_model"]["completeness"] == pytest.approx(1.0, rel=0.05) + + +def test_sla_mode_certifies_the_quote(): + _model, state = _search(_Block(), method_options={"num_path_nodes": 2}) + epsilon = 0.5 * state["damage_model"]["f_corner"] + + _model, sla_state = _search( + _Block(), effective_bits=None, method_options={"max_predicted_damage": epsilon} + ) + assert sla_state["best"]["is_satisfied"] + assert sla_state["best"]["predicted_damage"] <= epsilon + 1e-12 + + +def _synthetic_searcher(n_groups=5, seed=0): + rng = np.random.default_rng(seed) + aggressive = QuantRecipe("INT4_BLOCKWISE_WEIGHT_ONLY_CFG") + moderate = QuantRecipe("INT8_DEFAULT_CFG") + no_quant = QuantRecipe(quant_cfg=None) + + searcher = AutoQuantizeAumannShapleySearcher() + searcher.candidate_stats = {} + b_total = 0.0 + for i in range(n_groups): + numel = float(rng.integers(100, 1000)) + b8 = float(rng.uniform(0.001, 0.05)) + b4 = b8 + float(rng.uniform(0.001, 0.1)) + b_total += b4 + searcher.candidate_stats[f"g{i}.quant_recipe"] = { + "formats": [aggressive, moderate, no_quant], + "scores": [b4, b8, 0.0], + "costs": [numel * aggressive.compression, numel * moderate.compression, numel], + "module_names": [f"g{i}"], + "quantizer_attrs": {f"g{i}": ("input_quantizer", "weight_quantizer")}, + "cost_weight": 1.0, + "allow_no_quant": True, + "is_fixed": False, + "uncompressed_cost": numel, + } + searcher.damage_model = { + "link": "coverage", + "c": 0.5, + "f_corner": 0.5 * (1 - np.exp(-b_total)), + "valid": True, + } + searcher.cost_model = "weight" + searcher.config = {**searcher.default_search_config} + return searcher + + +def _selected(searcher, best): + score = cost = 0.0 + for info in best.values(): + score += info["scores"] + cost += info["costs"] + return score, cost + + +def test_sla_search_certified_and_near_optimal(): + for seed in range(3): + searcher = _synthetic_searcher(seed=seed) + c = searcher.damage_model["c"] + for eps_frac in (0.9, 0.1, 0.02): + epsilon = eps_frac * searcher.damage_model["f_corner"] + searcher.config["max_predicted_damage"] = epsilon + best, is_satisfied = searcher.run_search_with_stats(max_weight_size=np.inf) + score, cost = _selected(searcher, best) + assert _predict_damage(c, score) <= epsilon + 1e-12 + assert is_satisfied + + budget = -np.log(1.0 - epsilon / c) + stats = searcher.candidate_stats + names = list(stats) + optimal = min( + ( + sum(stats[n]["costs"][k] for n, k in zip(names, combo, strict=True)) + for combo in itertools.product( + *[range(len(stats[n]["formats"])) for n in names] + ) + if sum(stats[n]["scores"][k] for n, k in zip(names, combo, strict=True)) + <= budget + 1e-12 + ), + default=np.inf, + ) + assert cost <= optimal * 1.05 + 1e-9 + + +def _brute_force_min_score(stats, budget): + names = list(stats) + return min( + ( + sum(stats[n]["scores"][k] for n, k in zip(names, combo, strict=True)) + for combo in itertools.product(*[range(len(stats[n]["formats"])) for n in names]) + if sum(stats[n]["costs"][k] for n, k in zip(names, combo, strict=True)) <= budget + 1e-9 + ), + default=np.inf, + ) + + +def test_solvers_optimal_within_their_contracts(): + """The LP is exact; the DP is exact up to the budget-grid resolution (its selection can + never beat the true optimum, and never trails the optimum of a grid-tightened budget).""" + for seed in range(3): + searcher = _synthetic_searcher(seed=seed) + stats = searcher.candidate_stats + total = sum(s["uncompressed_cost"] for s in stats.values()) + n_groups = len(stats) + for fraction in (0.9, 0.5): + budget = total * fraction + searcher.config["solver"] = "lp" + best_lp, satisfied_lp = searcher.run_search_with_stats(budget) + searcher.config["solver"] = "dp" + best_dp, satisfied_dp = searcher.run_search_with_stats(budget) + assert satisfied_lp and satisfied_dp + score_lp, cost_lp = _selected(searcher, best_lp) + score_dp, cost_dp = _selected(searcher, best_dp) + assert max(cost_lp, cost_dp) <= budget + 1e-6 + + optimum = _brute_force_min_score(stats, budget) + assert score_lp == pytest.approx(optimum, rel=1e-9) + tightened = _brute_force_min_score(stats, budget * (1 - (n_groups + 1) / 4096)) + assert optimum - 1e-9 <= score_dp <= tightened + 1e-9 + + +def test_coverage_inversion_recovers_forward_model(): + rng = np.random.default_rng(2) + n_groups, c, p = 16, 0.4, 0.5 + a_true = rng.uniform(0.001, 0.05, size=n_groups) + discounts = np.array([np.prod(1 - p * np.delete(a_true, i)) for i in range(n_groups)]) + attributions = c * a_true * discounts + + a, b, converged = _as_seed_coverage(attributions, c=c, p=p) + assert converged + assert np.allclose(a, a_true, rtol=1e-5) + + # One measured corner under-identifies the ceiling; the anchor must reproduce the + # corner exactly and preserve the allocation-relevant break-rate ratios. + b_true = -np.log(1 - a_true) + f_corner = c * (1 - np.exp(-b_true.sum())) + c_out, b_by_key, _kappa, _inflation, converged = _anchor_ceiling( + {"fmt": attributions}, f_corner, {"fmt": np.ones(n_groups, dtype=bool)} + ) + assert converged + b = b_by_key["fmt"] + assert _predict_damage(c_out, float(b.sum())) == pytest.approx(f_corner, rel=0.02) + assert np.allclose(b / b.sum(), b_true / b_true.sum(), rtol=1e-2) + + +def test_dp_scales_beyond_grid_group_counts(): + """Row count above the DP grid resolution must not report a false infeasibility.""" + searcher = _synthetic_searcher(n_groups=5000, seed=0) + searcher.config["solver"] = "dp" + total = sum(s["uncompressed_cost"] for s in searcher.candidate_stats.values()) + for fraction in (0.9, 0.5): + best, satisfied = searcher.run_search_with_stats(total * fraction) + assert satisfied + _score, cost = _selected(searcher, best) + assert cost <= total * fraction + 1e-6 + + +def test_method_options_validation(): + with pytest.raises(ValueError, match="Invalid method_options"): + _search(_Block(), method="kl_div", method_options={"num_path_nodes": 2}) + with pytest.raises(ValueError, match="Invalid method_options"): + _search(_Block(), method_options={"num_score_steps": 999}) # core input + with pytest.raises(ValueError, match="Invalid method_options"): + _search(_Block(), method_options={"unknown_option": 1}) + for bad_options in ( + {"num_path_nodes": True}, + {"num_path_nodes": 1.5}, + {"num_path_nodes": 0}, + {"damage_link": "unsupported"}, + {"solver": "unsupported"}, + {"max_predicted_damage": float("inf")}, + {"max_predicted_damage": -1.0}, + ): + with pytest.raises(ValueError): + _search(_Block(), method_options=bad_options) + with pytest.raises(TypeError, match="method_options must be a dict"): + _search(_Block(), method_options=[("num_path_nodes", 2)]) + + +def test_both_targets_rejected_before_model_conversion(): + from modelopt.torch.quantization.nn import TensorQuantizer + + model = _Block() + with pytest.raises(ValueError, match="not both"): + _search(model, effective_bits=8.0, method_options={"max_predicted_damage": 1e-3}) + assert type(model.mlp) is torch.nn.Linear + assert not any(isinstance(m, TensorQuantizer) for m in model.modules()) + + +def _inject_scores_and_corner(monkeypatch, injected, corner): + def inject(self, is_param_grad_enabled): + no_quant = QuantRecipe(quant_cfg=None) + self._corner_kl_sum = torch.tensor(float(corner)) + self._score_tokens = 1 + for hparam in self._configurable_hparams(): + for recipe in hparam.choices: + if recipe == no_quant: + continue + value = injected[str(recipe).split("(")[0]] + for module in hparam.score_modules: + hparam._importance_dict[recipe][module] = torch.tensor(value) + + monkeypatch.setattr(AutoQuantizeAumannShapleySearcher, "_estimate_auto_quantize_scores", inject) + + +def test_zero_corner_with_mixed_sign_attributions_is_invalid(monkeypatch): + """A zero measured corner with remaining positive attribution mass must invalidate the + fit: a signed cancellation is not a zero-damage model.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": -0.01, "INT8_DEFAULT_CFG": 0.02}, + corner=0.0, + ) + _model, state = _search(_OneLinear(), effective_bits=16.0) + damage_model = state["damage_model"] + assert "zero_corner_with_attribution_mass" in damage_model["approximation_flags"] + assert damage_model["valid"] is False + assert state["best"]["predicted_damage_valid"] is False + + +def test_projected_damage_model_matches_solver_scores(monkeypatch): + """The persisted link values must be the quote-operative (projected) ones, with the + projection recorded and the unprojected corner anchor kept alongside.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": 0.01, "INT8_DEFAULT_CFG": 0.02}, + corner=0.01, + ) + _model, state = _search(_OneLinear(), effective_bits=16.0) + damage_model = state["damage_model"] + assert "monotonicity_projection" in damage_model["approximation_flags"] + assert damage_model["monotonicity_adjustment"] > 0 + + (name,) = next(iter(damage_model["b"].values())).keys() + stat = state["candidate_stats"][name] + format_labels = [str(recipe) for recipe in stat["formats"]] + for label, values in damage_model["b"].items(): + assert values[name] == pytest.approx(stat["scores"][format_labels.index(label)]) + + # The unprojected link stays exactly anchored; the projected corner may exceed it. + unprojected_corner = sum( + values[name] + for label, values in damage_model["b_unprojected"].items() + if label.startswith("INT4") + ) + assert _predict_damage(damage_model["c"], unprojected_corner) == pytest.approx( + damage_model["f_corner"], rel=0.02 + ) + assert damage_model["projected_corner_damage"] >= damage_model["f_corner"] * 0.99 + + +def test_custom_format_identity_across_search_spaces(): + """Identical custom formats under different auto-generated names are one format.""" + custom = { + "quant_cfg": [{"quantizer_name": "*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}], + "algorithm": "max", + } + model = _Block() + with pytest.warns(UserWarning, match="custom quantization formats"): + _model, state = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + module_search_spaces=[ + {"module_name_patterns": ["*attn*"], "quantization_formats": [dict(custom)]}, + {"module_name_patterns": ["*mlp*"], "quantization_formats": [dict(custom)]}, + ], + fixed_quantization_config="INT8_DEFAULT_CFG", + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + num_calib_steps=1, + num_score_steps=1, + method="aumann_shapley", + ) + assert len(state["damage_model"]["as_scores"]) == 1 + assert state["damage_model"]["damage_reference"]["type"] == "quantized_baseline" + + +def test_heterogeneous_ladders_flagged(): + model = _Block() + with pytest.warns(UserWarning, match="differing candidate ladders"): + _model, state = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=list(SEARCH_FORMATS), + module_search_spaces=[ + { + "module_name_patterns": ["*mlp*"], + "quantization_formats": [mtq.INT8_DEFAULT_CFG], + } + ], + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + num_calib_steps=1, + num_score_steps=1, + method="aumann_shapley", + ) + assert "heterogeneous_ladders" in state["damage_model"]["approximation_flags"] + + +def test_corner_is_anchored_even_when_attributions_are_incomplete(): + """The damage link must reproduce the measured corner regardless of attribution mass.""" + rng = np.random.default_rng(0) + f_corner = 0.4 + for scale in (1.0, 0.1, 1e-6): # complete, incomplete, and nearly-vanished attributions + attributions = rng.uniform(0.001, 0.01, size=32) * scale + mask = {"fmt": np.ones(32, dtype=bool)} + c, b_by_key, _kappa, _inflation, converged = _anchor_ceiling( + {"fmt": attributions}, f_corner, mask + ) + assert converged + corner_prediction = _predict_damage(c, float(b_by_key["fmt"].sum())) + assert corner_prediction == pytest.approx(f_corner, rel=0.02) + + +def test_tiny_attributions_do_not_inflate(): + """Arbitrarily small positive attributions must not be floored into phantom damage.""" + a, b, converged = _as_seed_coverage(np.full(5000, 1e-15), c=0.4) + assert converged + assert float(b.sum()) < 1e-9 + + +def test_raw_scores_survive_the_base_monotonicity_clamp(monkeypatch): + """A negative attribution for one format must never overwrite its neighbor's positive + one: fitting and diagnostics read the unclamped values; only solver scores are + monotonized. Uses injected scores so the negative/positive case is deterministic.""" + injected = {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": -2.793747e-06, "INT8_DEFAULT_CFG": 2.738088e-07} + + def inject_scores(self, is_param_grad_enabled): + no_quant = QuantRecipe(quant_cfg=None) + self._corner_kl_sum = torch.tensor(0.4) + self._score_tokens = 1 + for hparam in self._configurable_hparams(): + for recipe in hparam.choices: + if recipe == no_quant: + continue + value = injected[str(recipe).split("(")[0]] + for module in hparam.score_modules: + hparam._importance_dict[recipe][module] = torch.tensor(value) + + monkeypatch.setattr( + AutoQuantizeAumannShapleySearcher, "_estimate_auto_quantize_scores", inject_scores + ) + _model, state = _search(_OneLinear(), effective_bits=16.0) + + as_scores = state["damage_model"]["as_scores"] + (name,) = next(iter(as_scores.values())).keys() + for label, values in as_scores.items(): + expected = injected[label.split("(")[0]] + assert values[name] == pytest.approx(expected, rel=1e-9) + + # The positive INT8 damage must reach the solver: the monotone projection may raise + # the more aggressive neighbor but must never erase a real score. + stat = state["candidate_stats"][name] + int8_index = [str(r).split("(")[0] for r in stat["formats"]].index("INT8_DEFAULT_CFG") + assert stat["scores"][int8_index] > 0 + assert state["damage_model"]["negative_attribution_mass"] > 0.5 + + +def test_invalid_method_options_leave_model_untouched(): + from modelopt.torch.quantization.nn import TensorQuantizer + + for options, exception in ( + (123, TypeError), + ([], TypeError), + ({"unknown_option": 1}, ValueError), + ({"num_score_steps": 999}, ValueError), + ({"num_path_nodes": 0}, ValueError), + ({"solver": "unsupported"}, ValueError), + ): + model = _Block() + with pytest.raises(exception): + _search(model, method_options=options) + assert type(model.mlp) is torch.nn.Linear, f"model mutated for {options!r}" + assert not any(isinstance(m, TensorQuantizer) for m in model.modules()) + + +def test_forced_single_format_group_recorded_as_baseline(): + """A single-candidate allow_no_quant=False group stays quantized in every reference + pass, so the damage reference must name it (quotes are incremental to it).""" + model = _Block() + _model, state = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=list(SEARCH_FORMATS), + module_search_spaces=[ + { + "module_name_patterns": ["*mlp*"], + "quantization_formats": [mtq.INT8_DEFAULT_CFG], + "allow_no_quant": False, + } + ], + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + num_calib_steps=1, + num_score_steps=1, + method="aumann_shapley", + ) + reference = state["damage_model"]["damage_reference"] + assert reference["type"] == "quantized_baseline" + assert any("mlp" in name for name in reference["forced_groups"]) + assert state["best"]["predicted_damage"] >= 0 + + +def test_mckp_choice_indices_above_255(): + from modelopt.torch.quantization._auto_quantize_shapley import _mckp_max_value + + values = np.arange(300, dtype=float)[None, :] + costs = np.zeros((1, 300), dtype=np.int64) + selection, total = _mckp_max_value(values, costs, budget=10) + assert selection[0] == 299 + assert total == 299.0 + + +def test_zero_attributions_invert_to_exact_zero(): + attributions = np.array([0.0, 0.02, 0.0, 0.05]) + a, b, converged = _as_seed_coverage(attributions, c=0.4) + assert converged + assert a[0] == a[2] == b[0] == b[2] == 0.0 + assert (a[[1, 3]] > 0).all() + + a, b, converged = _as_seed_coverage(np.zeros(4), c=0.4) + assert converged and (a == 0).all() and (b == 0).all() + + +def test_scoring_signature_guards_resume(tmp_path): + checkpoint = str(tmp_path / "state.pth") + _search(_Block(), checkpoint=checkpoint, method_options={"num_path_nodes": 2}) + + # Changing what the stored scores mean must be rejected. + with pytest.raises(ValueError, match="scoring signature"): + _search(_Block(), checkpoint=checkpoint, method_options={"num_path_nodes": 3}) + + # Changing only how they are solved reuses the stored scores. + _model, state = _search( + _Block(), checkpoint=checkpoint, method_options={"num_path_nodes": 2, "solver": "dp"} + ) + assert state["best"]["is_satisfied"] + + +def _shapley_data_parallel(rank, size, baseline): + from modelopt.torch.utils.distributed import DistributedProcessGroup + + _model, state = _search(_Block(seed=0), method_options={"num_path_nodes": 2}) + state_rank0 = DistributedProcessGroup.get_dist_syncd_obj( + state if rank == 0 else None, DistributedProcessGroup(None), lambda a: a[0] + ) + local = {k: v for k, v in state.items() if k != "quantizer_states"} + rank0 = {k: v for k, v in state_rank0.items() if k != "quantizer_states"} + assert local == rank0 + assert state["best"]["is_satisfied"] + + # Every rank scores the same batches, so correct DP reductions multiply the token count + # by the world size while leaving all per-token quantities equal to the single-process + # baseline; a dropped reduction shows up as a factor of the world size. Tolerances + # absorb float32 backward jitter (amplified by the coverage inversion), nothing more. + damage_model = state["damage_model"] + assert damage_model["n_score_tokens"] == size * baseline["n_score_tokens"] + assert damage_model["f_corner"] == pytest.approx(baseline["f_corner"], rel=1e-6) + for name, scores in baseline["scores"].items(): + got = state["candidate_stats"][name]["scores"] + assert got == pytest.approx(scores, rel=1e-3, abs=1e-9), f"{name}: {got} vs {scores}" + + +def test_data_parallel_aumann_shapley(skip_on_windows): + from functools import partial + + from _test_utils.torch.distributed.utils import spawn_multiprocess_job + + _model, single = _search(_Block(seed=0), method_options={"num_path_nodes": 2}) + baseline = { + "n_score_tokens": single["damage_model"]["n_score_tokens"], + "f_corner": single["damage_model"]["f_corner"], + "scores": {name: stat["scores"] for name, stat in single["candidate_stats"].items()}, + } + spawn_multiprocess_job(2, partial(_shapley_data_parallel, baseline=baseline), backend="gloo") + + +def test_anchor_ceiling_rejects_non_finite_measurements(): + attributions = np.array([0.01, 0.02]) + mask = {"fmt": np.ones(2, dtype=bool)} + for f_corner in (float("nan"), float("inf")): + with pytest.raises(ValueError, match="finite"): + _anchor_ceiling({"fmt": attributions}, f_corner, mask) + with pytest.raises(ValueError, match="finite"): + _anchor_ceiling({"fmt": np.array([0.01, float("nan")])}, 0.4, mask) + + +@pytest.mark.parametrize("corner", [float("nan"), float("inf")]) +def test_non_finite_corner_invalidates_damage_model(monkeypatch, corner): + """A non-finite corner KL must invalidate the fit (not hang) and keep scores finite.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": 2e-6, "INT8_DEFAULT_CFG": 1e-7}, + corner, + ) + _model, state = _search(_OneLinear(), effective_bits=16.0) + + damage_model = state["damage_model"] + assert damage_model["valid"] is False + assert "non_finite_measurements" in damage_model["approximation_flags"] + for stat in state["candidate_stats"].values(): + assert all(math.isfinite(score) for score in stat["scores"]) + + +def test_non_finite_attribution_excluded_and_anchor_invalidated(monkeypatch): + """A broken candidate leaves the search space; because it was the group's most + aggressive format, the measured corner no longer describes the pruned candidate space + and the fit must not certify quotes against it.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": float("nan"), "INT8_DEFAULT_CFG": 1e-7}, + 0.4, + ) + _model, state = _search(_OneLinear(), effective_bits=16.0) + + (name,) = state["candidate_stats"] + stat = state["candidate_stats"][name] + assert all("INT4_BLOCKWISE" not in str(recipe) for recipe in stat["formats"]) + # The solver objective must be the normalized additive attributions, NOT an inversion + # anchored to the removed format's corner measurement. + assert stat["scores"] == [pytest.approx(1e-7), 0.0] + damage_model = state["damage_model"] + assert "non_finite_scores_excluded" in damage_model["approximation_flags"] + assert "corner_format_excluded" in damage_model["approximation_flags"] + (dropped,) = damage_model["excluded_candidates"][name] + assert "INT4_BLOCKWISE" in dropped + assert damage_model["valid"] is False + assert state["best"]["predicted_damage_valid"] is False + + +def _search_no_bf16(model, effective_bits): + """Search where the only candidates are INT4/INT8 (no no-quant fallback).""" + return mtq.auto_quantize( + model, + constraints={"effective_bits": effective_bits}, + module_search_spaces=[ + { + "module_name_patterns": ["*"], + "quantization_formats": list(SEARCH_FORMATS), + "allow_no_quant": False, + } + ], + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + num_calib_steps=2, + num_score_steps=2, + method="aumann_shapley", + ) + + +def test_pruned_singleton_group_stays_fitted_with_unquantized_reference(monkeypatch): + """Pruning down to one candidate must not demote the group to a fixed-baseline one: + it was unquantized during the reference pass and its survivor still needs a + token-normalized fitted score.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": float("nan"), "INT8_DEFAULT_CFG": 1e-7}, + 0.4, + ) + _model, state = _search_no_bf16(_OneLinear(), effective_bits=8.0) + + (name,) = state["candidate_stats"] + stat = state["candidate_stats"][name] + assert [str(recipe).split("(")[0] for recipe in stat["formats"]] == ["INT8_DEFAULT_CFG"] + damage_model = state["damage_model"] + assert damage_model["damage_reference"] == {"type": "unquantized"} + (label,) = damage_model["as_scores"] + assert damage_model["as_scores"][label] == {name: pytest.approx(1e-7)} + assert stat["scores"] == [pytest.approx(1e-7)] + assert "corner_format_excluded" in damage_model["approximation_flags"] + assert damage_model["valid"] is False + + +def test_offline_resolve_preserves_forced_invalid_state(monkeypatch): + """get_auto_quantize_config re-solves on a bare searcher; the forced-candidate state + must survive the round trip so the re-solve cannot silently report a clean solution.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": float("inf"), "INT8_DEFAULT_CFG": float("nan")}, + 0.4, + ) + _model, state = _search_no_bf16(_OneLinear(), effective_bits=8.0) + assert not state["best"]["is_satisfied"] + + with pytest.warns(UserWarning, match="non-finite"): + config = mtq.get_auto_quantize_config(state, {"effective_bits": 8.0}) + assert config["algorithm"] == "max" + + +def test_all_non_finite_without_no_quant_reports_unsatisfied(monkeypatch): + """With every candidate non-finite and no no-quant fallback, the retained forced + choice must not report success, and the failed measurement must stay visible.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": float("inf"), "INT8_DEFAULT_CFG": float("nan")}, + 0.4, + ) + _model, state = _search_no_bf16(_OneLinear(), effective_bits=8.0) + + (name,) = state["candidate_stats"] + stat = state["candidate_stats"][name] + assert [str(recipe).split("(")[0] for recipe in stat["formats"]] == ["INT8_DEFAULT_CFG"] + assert not math.isfinite(stat["raw_scores"][0]) + assert not state["best"]["is_satisfied"] + damage_model = state["damage_model"] + assert damage_model["valid"] is False + assert "non_finite_candidate_forced" in damage_model["approximation_flags"] + (forced_format,) = damage_model["forced_candidates"].values() + assert "INT8_DEFAULT" in forced_format + assert state["best"]["predicted_damage_valid"] is False + + +def test_infinite_candidate_never_wins_the_allocation(monkeypatch): + """A non-finite measurement must not be zeroed into a free candidate: with finite INT4 + damage, infinite INT8 damage, and an 8-bit target, the solver must pick INT4.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": 2e-6, "INT8_DEFAULT_CFG": float("inf")}, + 0.4, + ) + _model, state = _search(_OneLinear(), effective_bits=8.0) + + (name,) = state["candidate_stats"] + assert all( + "INT8_DEFAULT" not in str(recipe) for recipe in state["candidate_stats"][name]["formats"] + ) + assert "INT4_BLOCKWISE" in str(state["best"]["recipe"][name]) + assert state["best"]["is_satisfied"] + assert state["damage_model"]["valid"] is True + + +def test_all_candidates_non_finite_falls_back_to_no_quant(monkeypatch): + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": float("inf"), "INT8_DEFAULT_CFG": float("nan")}, + 0.4, + ) + _model, state = _search(_OneLinear(), effective_bits=6.0) + + (name,) = state["candidate_stats"] + stat = state["candidate_stats"][name] + assert [str(recipe).split("(")[0] for recipe in stat["formats"]] == ["NONE"] + assert str(state["best"]["recipe"][name]).split("(")[0] == "NONE" + assert not state["best"]["is_satisfied"] + + +def test_nested_score_modules_are_scored(): + """A score module nested inside another must not be zeroed by the outer replay. + + Routed experts score at ``...mlp`` while shared experts inside that same mlp score at + themselves. The outer module's replay loop re-enters the inner forward under + ``no_grad``; if that clears the inner's cached diffs, the shared experts silently + score zero and the solver treats them as free to quantize. + """ + + class _Expert(torch.nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = torch.nn.Linear(32, 32) + self.up_proj = torch.nn.Linear(32, 32) + self.down_proj = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.down_proj(self.gate_proj(x) * self.up_proj(x)) + + class _MLP(torch.nn.Module): + def __init__(self): + super().__init__() + self.experts = torch.nn.ModuleList([_Expert() for _ in range(2)]) + self.shared_experts = _Expert() + + def forward(self, x): + out = self.shared_experts(x) + for expert in self.experts: + out = out + expert(x) + return out + + class _Layer(torch.nn.Module): + def __init__(self): + super().__init__() + self.mlp = _MLP() + + def forward(self, x): + return self.mlp(x) + + class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.layer = _Layer() + + def forward(self, x): + return self.layer(x) + + def get_input(self): + return torch.randn(1, 4, 32) + + torch.manual_seed(0) + model = _Model() + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda model, batch: model(batch), + num_calib_steps=2, + num_score_steps=2, + method="aumann_shapley", + ) + + def _quant_score(module): + hparam = module.get_hparam("quant_recipe") + return max( + hparam.get_score(recipe) for recipe in hparam.choices if "NONE" not in str(recipe) + ) + + # The nested (shared-expert) group must carry real attribution, like the routed group. + assert _quant_score(model.layer.mlp.shared_experts.gate_proj) > 0.0 + assert _quant_score(model.layer.mlp.experts[0].gate_proj) > 0.0 + + +def test_negative_inf_candidates_do_not_leak_into_solver_scores(monkeypatch): + """A -inf attribution must not survive candidate exclusion. + + Attributions here are signed and unclamped, so a candidate can measure -inf. The base + searcher's running-min chain then propagates it into every less aggressive entry + including no-quant, and a group left with no quantized candidate is dropped from the + solver tables, so the coverage projection never rewrites it. Unlike +inf and nan, which + collapse to 0.0 through ``min``, -inf would otherwise reach the LP objective. + """ + _inject_scores_and_corner( + monkeypatch, + { + "INT4_BLOCKWISE_WEIGHT_ONLY_CFG": float("-inf"), + "INT8_DEFAULT_CFG": float("-inf"), + }, + 0.4, + ) + _model, state = _search(_OneLinear(), effective_bits=6.0) + + (name,) = state["candidate_stats"] + stat = state["candidate_stats"][name] + assert [str(recipe).split("(")[0] for recipe in stat["formats"]] == ["NONE"] + assert all(math.isfinite(score) for score in stat["scores"]) + assert str(state["best"]["recipe"][name]).split("(")[0] == "NONE" + assert not state["best"]["is_satisfied"] + + +def test_vocab_sharded_loss_is_rejected_before_calibration(monkeypatch): + """The unsupported-parallelism error must fire before the calibration passes run.""" + monkeypatch.setattr( + AutoQuantizeAumannShapleySearcher, "_loss_is_vocab_sharded", lambda self: True + ) + + calibrated = [] + import modelopt.torch.quantization.model_quant as _model_quant + + real_calibrate = _model_quant.calibrate + + def _spy(*args, **kwargs): + calibrated.append(True) + return real_calibrate(*args, **kwargs) + + monkeypatch.setattr(_model_quant, "calibrate", _spy) + + with pytest.raises(NotImplementedError, match="vocab-sharded"): + _search(_OneLinear(), effective_bits=6.0) + assert not calibrated, "calibration ran before the unsupported-method check" + + +def test_no_quant_sorts_last_against_a_compression_tie(): + """The searcher reads formats[0] as the most aggressive candidate and treats the last + entry as unquantized. no_quant's compression is 1.0, which a config that leaves weights + at 16 bits ties exactly, so the ordering must pin no_quant last rather than let the + config-JSON tiebreak decide. + """ + no_quant = QuantRecipe(quant_cfg=None) + # Enabling a quantizer without a cfg leaves estimate_quant_compression at 1.0. + tied = QuantRecipe( + {"quant_cfg": [{"quantizer_name": "*input_quantizer", "enable": True}]}, + name="TIED_16BIT", + ) + assert tied.compression == no_quant.compression + assert not tied.is_no_quant and no_quant.is_no_quant + + ladder = sorted([QuantRecipe("NVFP4_DEFAULT_CFG"), no_quant, tied]) + assert not ladder[0].is_no_quant, "most aggressive entry must be a quantized format" + assert ladder[-1].is_no_quant, "no_quant must terminate the ladder" + + # Formats that compress weights are unaffected by the tiebreak. + standard = [QuantRecipe(c) for c in ("INT8_DEFAULT_CFG", "NVFP4_DEFAULT_CFG")] + [no_quant] + assert sorted(standard) == sorted( + standard, key=lambda r: (r.compression, r.checkpoint_signature) + )