diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..fc5ecbf6329 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -8,6 +8,7 @@ Changelog *Quantization* +- Add ``method="aumann_shapley"`` to ``mtq.auto_quantize`` for label-free path-integral sensitivity scoring, predicted calibration damage, and optional damage-bound search through ``method_options``. - Add ``mtq.temporarily_fold_weights`` for repeated frozen-weight inference and ``mtq.preserve_quantizer_attributes_context`` for restoring temporary quantizer property and type changes. Temporary folding snapshots affected fake-quant weights on a configurable device and restores them with their quantizer state; retained pre-quant scales are inactive, while shared weights, shared quantizers, and ``SequentialQuantizer`` weights are unsupported. - 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. diff --git a/modelopt/torch/quantization/_auto_quantize_shapley.py b/modelopt/torch/quantization/_auto_quantize_shapley.py new file mode 100644 index 00000000000..65ecc790adc --- /dev/null +++ b/modelopt/torch/quantization/_auto_quantize_shapley.py @@ -0,0 +1,1091 @@ +# 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 measures each candidate at a few points partway between the unquantized and the +quantized model, rather than only at the unquantized one. At step +``t = (k + 1/2) / num_path_nodes`` every scored module outputs ``y + t * (Q(y) - y)`` -- a blend +of its real and quantized output, produced by re-running the module with the candidate's +quantizers active -- and one backward pass accumulates ``
`` per (group, +format). Spreading the measurement over several steps is what makes a KL objective usable: KL +against the model's own outputs is exactly zero at the unquantized point, and flat there, so +measuring only at that point would give 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. Each scored module additionally re-runs its own +forward once per candidate within those passes. Replay assumes module forwards are +side-effect-free: they are re-run with no state reset between calls, as in the gradient method. + +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`` -- the fraction of the measured damage the +summed contributions reproduce, 1.0 being exact -- 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 2): how many steps between the unquantized and quantized + model to measure at. +- ``damage_link`` (default ``"coverage"``): how per-group scores combine into a damage + estimate. ``"coverage"`` uses the coverage form ``damage = c * (1 - exp(-sum(b)))`` of + https://arxiv.org/abs/2607.12266; ``"additive"`` sums the raw scores. +- ``max_predicted_damage`` (default None): minimize weight cost subject to predicted damage <= + this bound (mutually exclusive with an ``effective_bits`` constraint). +""" + +import gc +import math +from dataclasses import dataclass + +import numpy as np +import torch + +from modelopt.torch.opt.searcher import LPS, SearchConfig, SearchStateDict +from modelopt.torch.opt.utils import named_hparams +from modelopt.torch.utils import print_rank_0, report_memory, warn_rank_0 +from modelopt.torch.utils.distributed import DistributedProcessGroup + +from .algorithms import ( + AUTO_QUANTIZE_SEARCHERS, + QuantRecipe, + QuantRecipeHparam, + _AutoQuantizeBackwardScoringSearcher, + _AutoQuantizeCandidateReplayScoringSession, + _get_kl_div_loss, + _get_lm_head, + _get_log_prob, +) + +__all__ = ["AutoQuantizeAumannShapleySearcher"] + + +# Gauss-Legendre with 32 nodes is exact for coverage paths containing up to 64 groups. +_COVERAGE_QUADRATURE_ORDER = 32 +_COVERAGE_QUADRATURE_NODES, _COVERAGE_QUADRATURE_WEIGHTS = np.polynomial.legendre.leggauss( + _COVERAGE_QUADRATURE_ORDER +) +_COVERAGE_QUADRATURE_NODES = 0.5 * (_COVERAGE_QUADRATURE_NODES + 1.0) +_COVERAGE_QUADRATURE_WEIGHTS *= 0.5 + + +@dataclass(frozen=True) +class _AttributionData: + """Normalized attribution tables used by the damage fit.""" + + names: tuple[str, ...] + keys: tuple[str, ...] + labels: dict[str, str] + signed_by_key: dict[str, np.ndarray] + nonnegative_by_key: dict[str, np.ndarray] + corner_mask_by_key: dict[str, np.ndarray] + heterogeneous_ladders: bool + negative_mass: float + corner_mass: float + + +@dataclass(frozen=True) +class _DamageFit: + """Result of mapping attributions into the configured damage model.""" + + scores_by_key: dict[str, np.ndarray] + unprojected_b_by_key: dict[str, np.ndarray] + is_coverage: bool + valid: bool + flags: tuple[str, ...] = () + ceiling: float = 0.0 + kappa: float = 1.0 + ceiling_inflation: float = 1.0 + converged: bool = False + + +class _AumannShapleyScoringSession(_AutoQuantizeCandidateReplayScoringSession): + """Collect path-integral attributions using the shared backward-scoring lifecycle.""" + + def __init__( + self, + model, + hparams, + score_modules, + recipes, + num_path_nodes, + forward_step, + lm_head, + is_param_grad_enabled, + verbose=False, + ): + super().__init__(model, score_modules, is_param_grad_enabled, verbose=verbose) + self.hparams = tuple(hparams) + self.recipes = tuple(recipes) + self.num_path_nodes = num_path_nodes + self.forward_step = forward_step + self.lm_head = lm_head + self.current_recipe: QuantRecipe | None = None + self.path_position = 0.0 + self.corner_kl_sum: torch.Tensor | None = None + self.score_tokens = 0 + + def _set_all_hparams(self, recipe_of) -> None: + """Set every configurable recipe hparam.""" + for hparam in self.hparams: + hparam.active = recipe_of(hparam) + + def forward(self, module, *args, **kwargs): + """Emit a path-shifted output and cache the current candidate's differences.""" + recipe = self.current_recipe + if recipe is None: + return self.original_forward(module)(*args, **kwargs) + + output, base = self._run_unquantized(module, *args, **kwargs) + output_diffs = self._replay_candidates( + module, + base, + lambda hparam: (recipe,) if recipe in hparam.choices else (), + *args, + **kwargs, + ) + diff_total = None + for recipe_diffs in output_diffs.values(): + for output_diff in recipe_diffs.values(): + diff_total = output_diff if diff_total is None else diff_total + output_diff + + if diff_total is None: + return output + shifted = base + self.path_position * diff_total + if torch.is_grad_enabled() and shifted.requires_grad: + self._register_candidate_score_hook(module, shifted, output_diffs) + if not isinstance(output, tuple): + return shifted + if hasattr(output, "_fields"): + # Keep namedtuple attributes that downstream modules may access. + return output._replace(**{output._fields[0]: shifted}) + return (shifted, *output[1:]) + + def _score_contribution(self, grad_output, output_diff): + return (grad_output.float() * output_diff.float()).sum() / self.num_path_nodes + + def score_step(self, model, data) -> None: + """Score every format and path node for one calibration batch.""" + self.current_recipe = None + self._set_all_hparams(lambda _hparam: self.no_quant) + with torch.no_grad(): + ref_logits = self.forward_step(model, data) + ref_logprob = _get_log_prob(ref_logits, lm_head=self.lm_head).detach() + self.score_tokens += int(ref_logprob.numel() // ref_logprob.shape[-1]) + del ref_logits + + self._set_all_hparams(lambda hparam: hparam.choices[0]) + try: + corner_logits = self.forward_step(model, data) + corner_loss = _get_kl_div_loss(ref_logprob, corner_logits, self.lm_head).detach() + self.corner_kl_sum = ( + corner_loss if self.corner_kl_sum is None else self.corner_kl_sum + corner_loss + ) + del corner_logits + finally: + self._set_all_hparams(lambda _hparam: self.no_quant) + + try: + for recipe in self.recipes: + self.current_recipe = recipe + for node in range(self.num_path_nodes): + self.path_position = (node + 0.5) / self.num_path_nodes + try: + logits = self.forward_step(model, data) + loss = _get_kl_div_loss(ref_logprob, logits, self.lm_head) + loss.backward() + del logits, loss + finally: + self._clear_output_grad_hooks() + finally: + self.current_recipe = None + + +def _coverage_path_discounts(a): + """Integrate each group's coverage discount over the Aumann-Shapley path.""" + survival = np.clip(1.0 - _COVERAGE_QUADRATURE_NODES[:, None] * a[None, :], 1e-9, None) + log_survival = np.log(survival) + products_without_group = np.exp(log_survival.sum(axis=1, keepdims=True) - log_survival) + return _COVERAGE_QUADRATURE_WEIGHTS @ products_without_group + + +def _as_seed_coverage(attributions, c, *, iters=200, tol=1e-10): + """Invert the coverage Aumann-Shapley integral for per-group damage fractions. + + Each ``a_i`` is the fraction of the ceiling ``c`` that group ``i`` accounts for. + + Returns ``(a, b = -log(1 - a), converged)``. Zero attributions map to exactly zero + fractions; 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): + discount = _coverage_path_discounts(a) + new = np.clip(attributions / (max(c, 1e-12) * np.clip(discount, 1e-6, None)), 0.0, 0.999) + delta = float(np.abs(new - a).max()) + a = new + 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 per-group damage fractions, sharing one 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 + max_ceiling = max_inflation * f_corner + while True: + 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_ceiling: + break + c = min(c * 1.3, max_ceiling) + 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(): + """Total log-headroom over the corner candidates.""" + 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): + """Damage implied by a total score under the coverage link.""" + return float(c * (1.0 - np.exp(-max(float(b_sum), 0.0)))) + + +class AutoQuantizeAumannShapleySearcher(_AutoQuantizeBackwardScoringSearcher): + """AutoQuantize searcher scoring with Aumann-Shapley damage attributions (see module doc).""" + + method_name = "aumann_shapley" + method_options_keys = frozenset({"num_path_nodes", "damage_link", "max_predicted_damage"}) + + @property + def default_search_config(self) -> SearchConfig: + """Get the default config for the searcher.""" + config = super().default_search_config + config.update( + { + "forward_step": None, + "num_path_nodes": 2, + "damage_link": "coverage", + "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 determine what the stored scores mean.""" + # Scoring settings are checkpointed; max_predicted_damage only changes the re-solve. + 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 ignored_key in config: + if config[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 = super().sanitize_search_config(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}" + ) + 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 before ``super().before_search()`` calibrates every search recipe. + 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 or {})} + # Stored scores are only reusable when their meaning is unchanged (see + # _current_scoring_signature); damage-bound 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]: + """Every configurable quant-recipe hparam in the model.""" + 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): + """Accumulate path-integral damage attributions for each candidate.""" + model = self.model + no_quant = QuantRecipe(quant_cfg=None) + self._raise_if_vocab_sharded() + hparams = self._configurable_hparams() + recipes = sorted({r for h in hparams for r in h.choices if r != no_quant}) + scoring_session = _AumannShapleyScoringSession( + model, + hparams, + self._configurable_score_modules(), + recipes, + int(self.config["num_path_nodes"]), + self.config["forward_step"], + _get_lm_head(model), + is_param_grad_enabled, + verbose=self.config.get("verbose", False), + ) + with scoring_session: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + report_memory("AutoQuantize(aumann_shapley): starting score estimation, ") + self._run_func( + scoring_session.score_step, + num_iters=self.config["num_score_steps"], + desc="Estimating aumann_shapley scores", + ) + + self._corner_kl_sum = scoring_session.corner_kl_sum + self._score_tokens = scoring_session.score_tokens + gc.collect() + if torch.cuda.is_available(): + report_memory("AutoQuantize(aumann_shapley): after score estimation") + + def _loss_is_vocab_sharded(self) -> bool: + """Whether the loss is computed over a vocab-sharded lm_head.""" + 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: + """Reject vocab-sharded losses, which the score pass cannot backprop.""" + # 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: + """Reduce a loss scalar over the data-parallel group.""" + # 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: + """Reduce a token count over the data-parallel group.""" + 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): + """A module carrying a parallel state, if any. + + Falls back to quant modules so the loss and token reductions stay consistent with + get_score, which uses the same fallback when a score module is a plain container. + """ + hparams = self._configurable_hparams() + for collection in ("score_modules", "quant_modules"): + for hparam in hparams: + for module in getattr(hparam, collection, ()): + 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 -inf into every + # less aggressive entry; groups pruned to no_quant are never rewritten downstream. + 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 _collect_attributions( + self, names: list[str], no_quant: QuantRecipe, tokens: int + ) -> _AttributionData: + """Build normalized, format-keyed attribution tables from candidate stats.""" + 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 = tuple(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 index, 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 + signed_by_key[key][index] = raw_score / tokens + ladder.append(key) + ladders.add(tuple(ladder)) + + corner_mask_by_key = {key: np.zeros(len(names), dtype=bool) for key in keys} + for index, name in enumerate(names): + corner = self.candidate_stats[name]["formats"][0] + if corner.is_no_quant: + raise ValueError( + f"no_quant sorted first in the candidate ladder for {name}; " + "QuantRecipe ordering must keep it last." + ) + corner_mask_by_key[corner.checkpoint_signature][index] = True + + signed_total = sum(float(np.abs(vector).sum()) for vector in signed_by_key.values()) + negative_total = sum( + float(np.abs(np.minimum(vector, 0.0)).sum()) for vector in signed_by_key.values() + ) + corner_mass = sum( + float(signed_by_key[key][mask].sum()) for key, mask in corner_mask_by_key.items() + ) + return _AttributionData( + names=tuple(names), + keys=keys, + labels=labels, + signed_by_key=signed_by_key, + nonnegative_by_key={ + key: np.maximum(vector, 0.0) for key, vector in signed_by_key.items() + }, + corner_mask_by_key=corner_mask_by_key, + heterogeneous_ladders=len(ladders) > 1, + negative_mass=negative_total / max(signed_total, 1e-12), + corner_mass=corner_mass, + ) + + def _fit_damage_link( + self, + attributions: _AttributionData, + f_corner: float, + *, + corner_removed: bool, + valid: bool, + ) -> _DamageFit: + """Map normalized attributions to additive or coverage-link solver scores.""" + is_coverage = ( + self.config["damage_link"] == "coverage" + and bool(attributions.names) + and bool(attributions.keys) + ) + if not is_coverage: + return _DamageFit( + scores_by_key=attributions.nonnegative_by_key, + unprojected_b_by_key={}, + is_coverage=False, + valid=valid, + ) + + flags: list[str] = [] + zeros_by_key = {key: np.zeros(len(attributions.names)) for key in attributions.keys} + positive_mass = sum( + float(vector.sum()) for vector in attributions.nonnegative_by_key.values() + ) + if not math.isfinite(f_corner) or corner_removed: + return _DamageFit( + scores_by_key=attributions.nonnegative_by_key, + unprojected_b_by_key=zeros_by_key, + is_coverage=True, + valid=False, + ) + + if f_corner <= 1e-12: + if positive_mass <= 1e-12: + flags.append("zero_damage") + return _DamageFit( + scores_by_key=zeros_by_key, + unprojected_b_by_key=zeros_by_key, + is_coverage=True, + valid=valid, + flags=tuple(flags), + converged=True, + ) + flags.append("zero_corner_with_attribution_mass") + return _DamageFit( + scores_by_key=attributions.nonnegative_by_key, + unprojected_b_by_key=zeros_by_key, + is_coverage=True, + valid=False, + flags=tuple(flags), + ) + + ceiling, b_by_key, kappa, inflation, converged = _anchor_ceiling( + attributions.nonnegative_by_key, + f_corner, + attributions.corner_mask_by_key, + ) + if not converged: + flags.append("inversion_not_converged") + return _DamageFit( + scores_by_key=b_by_key, + unprojected_b_by_key=b_by_key, + is_coverage=True, + valid=valid and converged and math.isfinite(ceiling), + flags=tuple(flags), + ceiling=ceiling, + kappa=kappa, + ceiling_inflation=inflation, + converged=converged, + ) + + def _project_solver_scores( + self, + attributions: _AttributionData, + no_quant: QuantRecipe, + scores_by_key: dict[str, np.ndarray], + ) -> tuple[dict[str, np.ndarray], float]: + """Write monotone per-group scores and return their format-keyed projection.""" + projected_by_key = {key: np.zeros(len(attributions.names)) for key in attributions.keys} + for index, name in enumerate(attributions.names): + stat = self.candidate_stats[name] + scores = [ + 0.0 + if recipe == no_quant + else float(scores_by_key[recipe.checkpoint_signature][index]) + for recipe in stat["formats"] + ] + for choice in range(len(scores) - 2, -1, -1): + scores[choice] = max(scores[choice], scores[choice + 1]) + stat["scores"] = scores + for recipe, score in zip(stat["formats"], scores, strict=True): + if recipe != no_quant: + projected_by_key[recipe.checkpoint_signature][index] = score + + unprojected_total = sum(float(vector.sum()) for vector in scores_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) + return projected_by_key, adjustment + + @staticmethod + def _measurement_diagnostics( + attributions: _AttributionData, + f_corner: float, + excluded: dict[str, list[str]], + forced_candidates: dict[str, str], + corner_removed: bool, + ) -> tuple[bool, list[str]]: + """Return ordered diagnostic flags and whether the fit can be certified.""" + diagnostics = ( + ("non_finite_measurements", not math.isfinite(f_corner), True), + ("non_finite_scores_excluded", bool(excluded), False), + ("corner_format_excluded", corner_removed, True), + ("non_finite_candidate_forced", bool(forced_candidates), True), + ("heterogeneous_ladders", attributions.heterogeneous_ladders, False), + ("negative_attribution_mass", attributions.negative_mass > 1e-3, False), + ) + flags = [name for name, present, _invalidates in diagnostics if present] + valid = not any(present and invalidates for _name, present, invalidates in diagnostics) + return valid, flags + + def _finalize_damage_model( + self, + attributions: _AttributionData, + fit: _DamageFit, + no_quant: QuantRecipe, + *, + f_corner: float, + tokens: int, + damage_reference: dict, + excluded: dict[str, list[str]], + forced_candidates: dict[str, str], + flags: list[str], + ) -> dict: + """Project solver scores and assemble the persisted damage-model record.""" + damage_model = { + "link": self.config["damage_link"], + "f_corner": f_corner, + "n_score_tokens": tokens, + "damage_reference": damage_reference, + "negative_attribution_mass": attributions.negative_mass, + "as_scores": { + attributions.labels[key]: dict(zip(attributions.names, vector.tolist())) + for key, vector in attributions.signed_by_key.items() + }, + } + if excluded: + damage_model["excluded_candidates"] = excluded + if forced_candidates: + damage_model["forced_candidates"] = forced_candidates + + projected_by_key, adjustment = self._project_solver_scores( + attributions, no_quant, fit.scores_by_key + ) + if adjustment > 1e-9: + flags.append("monotonicity_projection") + damage_model["monotonicity_adjustment"] = adjustment + + if fit.is_coverage: + projected_corner_b = sum( + float(projected_by_key[key][mask].sum()) + for key, mask in attributions.corner_mask_by_key.items() + ) + damage_model.update( + { + "c": fit.ceiling, + "kappa": fit.kappa, + "ceiling_inflation": fit.ceiling_inflation, + "inversion_converged": fit.converged, + "b": { + attributions.labels[key]: dict(zip(attributions.names, vector.tolist())) + for key, vector in projected_by_key.items() + }, + "b_unprojected": { + attributions.labels[key]: dict(zip(attributions.names, vector.tolist())) + for key, vector in fit.unprojected_b_by_key.items() + }, + "projected_corner_damage": _predict_damage(fit.ceiling, projected_corner_b), + } + ) + + damage_model["valid"] = fit.valid + damage_model["approximation_flags"] = flags + damage_model["completeness"] = attributions.corner_mass / max(f_corner, 1e-12) + return damage_model + + 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.") + # Record an invalidated model rather than leaving it unset, with the same key + # shape as the normal path: the quote is keyed off damage_model. + self.damage_model = { + "link": self.config["damage_link"], + "valid": False, + "approximation_flags": ["no_scored_tokens"], + "completeness": float("nan"), + "f_corner": float("nan"), + "n_score_tokens": 0, + "damage_reference": None, + "as_scores": {}, + } + 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 + + 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"]) + ] + attributions = self._collect_attributions(names, no_quant, tokens) + if attributions.heterogeneous_ladders: + warn_rank_0( + "aumann_shapley: groups have differing candidate ladders; the joint coverage " + "interpretation is approximate for the formats not shared by all groups." + ) + + valid, flags = self._measurement_diagnostics( + attributions, + f_corner, + excluded, + forced_candidates, + corner_removed, + ) + fit = self._fit_damage_link( + attributions, + f_corner, + corner_removed=corner_removed, + valid=valid, + ) + valid = fit.valid + flags.extend(fit.flags) + if fit.is_coverage and 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." + ) + + self.damage_model = self._finalize_damage_model( + attributions, + fit, + no_quant, + f_corner=f_corner, + tokens=tokens, + damage_reference=damage_reference, + excluded=excluded, + forced_candidates=forced_candidates, + flags=flags, + ) + + def run_search_with_stats(self, max_weight_size, verbose=False): + """Solve either the effective-bits or predicted-damage constraint with LPS.""" + 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 + ) + else: + recipes, is_satisfied = self._run_linear_program_search(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: + """Record the damage quote for the selected recipe.""" + 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) + valid = bool(damage_model.get("valid", True)) + if not valid: + # An invalidated fit leaves c = 0.0, which would otherwise quote exactly 0.0. + predicted = float("nan") + self.best["predicted_damage"] = predicted + self.best["predicted_damage_valid"] = valid + if self.config.get("verbose"): + print_rank_0( + f"AutoQuantize(aumann_shapley) predicted damage: {predicted:.4e} " + "(mean per-token KL, calibration units)" + ) + + def _recipes_from_selections(self, selections) -> dict: + """Map one selected candidate index per group back to recipe metadata.""" + best_recipes = {} + for (name, stat), selected_idx in zip( + self.candidate_stats.items(), selections, strict=True + ): + best_recipes[name] = { + "format": stat["formats"][selected_idx], + "costs": stat["costs"][selected_idx], + "scores": stat["scores"][selected_idx], + } + return best_recipes + + def _least_damage_selections(self) -> list[int]: + """Choose each group's lowest-damage candidate, preferring less quantization on ties.""" + return [ + min(range(len(stat["formats"])), key=lambda index: (stat["scores"][index], -index)) + for stat in self.candidate_stats.values() + ] + + def _run_damage_bound_search(self, max_predicted_damage, verbose=False): + """Minimize weight cost while keeping predicted damage within the requested bound.""" + damage_model = getattr(self, "damage_model", None) or {} + if not damage_model.get("valid", False): + self.status = "Invalid damage model" + warn_rank_0( + "AutoQuantize FAILED to find a solution! The damage fit is invalid " + f"(flags={damage_model.get('approximation_flags')}). Returning the " + "minimum-damage configuration." + ) + return self._recipes_from_selections(self._least_damage_selections()), False + + if damage_model.get("link") == "coverage" and "c" in damage_model: + ceiling = float(damage_model["c"]) + score_budget = ( + math.inf + if max_predicted_damage >= ceiling + else -math.log1p(-max_predicted_damage / ceiling) + ) + else: + score_budget = float(max_predicted_damage) + + damage_constraint_costs = [ + [0.0] * len(stat["scores"]) if stat.get("is_fixed", False) else stat["scores"] + for stat in self.candidate_stats.values() + ] + if math.isinf(score_budget): + selections = [ + min( + range(len(stat["formats"])), + key=lambda index: (stat["costs"][index], stat["scores"][index], index), + ) + for stat in self.candidate_stats.values() + ] + self.status = "Optimal" + else: + lps = LPS( + name="AutoQuantizeDamageBound", + constraints={"damage_score": score_budget}, + constraints_to_candidate_costs={"damage_score": damage_constraint_costs}, + candidate_scores=[stat["costs"] for stat in self.candidate_stats.values()], + objective_type="minimize", + verbose=verbose, + ) + selections, self.status = lps() + + selected_score = sum( + scores[selected_idx] + for scores, selected_idx in zip(damage_constraint_costs, selections, strict=True) + ) + is_satisfied = self.status == "Optimal" and selected_score <= score_budget + 1e-12 + if not is_satisfied: + warn_rank_0( + "AutoQuantize FAILED to find a solution within the predicted-damage bound. " + "Returning the minimum-damage configuration." + ) + selections = self._least_damage_selections() + + if verbose: + selected_score = sum( + scores[selected_idx] + for scores, selected_idx in zip(damage_constraint_costs, selections, strict=True) + ) + total_cost = sum( + stat["costs"][selected_idx] + for stat, selected_idx in zip( + self.candidate_stats.values(), selections, strict=True + ) + ) + print_rank_0( + f"AutoQuantize(damage bound): score {selected_score:.4e} " + f"(budget {score_budget:.4e}), weight size {total_cost:.2f}, " + f"satisfied={is_satisfied}" + ) + return self._recipes_from_selections(selections), 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..ab817dc5ce0 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -17,13 +17,14 @@ import copy import fnmatch +import functools import gc import types import warnings from abc import ABC, abstractmethod from collections import defaultdict from collections.abc import Callable, Sequence -from contextlib import nullcontext +from contextlib import ExitStack, nullcontext from typing import Any import regex as re @@ -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: + """Return the canonical signature used to identify the no-quant recipe.""" + 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,10 @@ def __repr__(self) -> str: return self._str_repr def __lt__(self, other: "QuantRecipe"): - return (self.compression, self.checkpoint_signature) < ( + # Callers treat the last choice as the unquantized end of the format ladder. + return (self.compression, self.is_no_quant, self.checkpoint_signature) < ( other.compression, + other.is_no_quant, other.checkpoint_signature, ) @@ -413,8 +427,9 @@ def __init__( self.allow_no_quant = allow_no_quant self.is_fixed = fixed_recipe is not None - self.quant_modules = list(set(quant_modules or [])) - self.score_modules = list(set(score_modules or self.quant_modules)) + # Module hashes depend on object identity, so sets can produce different orders per rank. + self.quant_modules = list(dict.fromkeys(quant_modules or [])) + self.score_modules = list(dict.fromkeys(score_modules or self.quant_modules)) fixed_quantizers = ( { @@ -467,11 +482,11 @@ def __init__( quant_recipe: dict.fromkeys(self.score_modules) for quant_recipe in self.choices } - # Attach this hparam to each score_module's set of hparams it scores + # Registration order follows the rank-stable runtime-group construction order. for score_module in self.score_modules: if not hasattr(score_module, "_hparams_for_scoring"): - score_module._hparams_for_scoring = set() - score_module._hparams_for_scoring.add(self) + score_module._hparams_for_scoring = [] + score_module._hparams_for_scoring.append(self) @property def active(self) -> HPType: @@ -532,6 +547,15 @@ def get_score(self, recipe: QuantRecipe) -> float: continue parallel_state = getattr(score_module, "parallel_state", None) + if parallel_state is None: + parallel_state = next( + ( + state + for module in self.quant_modules + if (state := getattr(module, "parallel_state", None)) is not None + ), + None, + ) if parallel_state is None: total_score += importance.cpu().item() @@ -632,10 +656,11 @@ 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 + 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 @@ -709,6 +734,9 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: ) return config + def validate_search_input(self, constraints, config) -> None: + """Validate method-specific inputs before quantizing the model.""" + def load_search_checkpoint(self) -> bool: return super().load_search_checkpoint(strict=False) @@ -1056,7 +1084,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,6 +1092,7 @@ 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) @@ -1071,6 +1100,7 @@ def initialize_candidate_stats(self): self.candidate_stats[name]["formats"] = formats self.candidate_stats[name]["scores"] = scores + self.candidate_stats[name]["raw_scores"] = raw_scores self.candidate_stats[name]["costs"] = costs self.candidate_stats[name]["module_names"] = hparam.quant_module_names self.candidate_stats[name]["quantizer_attrs"] = hparam.quant_module_replay_attrs @@ -1358,6 +1388,46 @@ def _get_search_lower_bounds(self): return [0.99, 0.90, None] return [None, 0.99, 0.90] + def _run_linear_program_search(self, max_weight_size, verbose=False): + """Select recipes with the standard AutoQuantize linear program.""" + for lower_bound in self._get_search_lower_bounds(): + constraints, constraint_name = self._get_constraints_for_search( + max_weight_size, lower_bound + ) + lps = LPS( + name="AutoQuantize", + constraints=constraints, + constraints_to_candidate_costs={ + constraint_name: [ + candidate_stat["costs"] for candidate_stat in self.candidate_stats.values() + ] + }, + candidate_scores=[ + candidate_stat["scores"] for candidate_stat in self.candidate_stats.values() + ], + objective_type="minimize", + verbose=verbose, + ) + selections, self.status = lps() + if self.status == "Optimal": + break + + is_satisfied = self.status == "Optimal" + if not is_satisfied: + warnings.warn( + "AutoQuantize FAILED to find a solution! The searched model might not meet all constraints. " + ) + + best_recipes = {} + for name, selected_idx in zip(self.candidate_stats, selections, strict=True): + best_recipes[name] = { + "format": self.candidate_stats[name]["formats"][selected_idx], + "costs": self.candidate_stats[name]["costs"][selected_idx], + "scores": self.candidate_stats[name]["scores"][selected_idx], + } + + return best_recipes, is_satisfied + @abstractmethod def run_search_with_stats(self, max_weight_size, verbose=False): """Run the search with stats to get the best recipe and whether the constraints are satisfied.""" @@ -1434,11 +1504,256 @@ def _get_auto_quantize_score(grad_output, output_diff): return x.clamp(-1e10, 1e10).square().sum() -def _add_auto_quantize_score(grad_output, output_diff, score_tensor): - score_tensor += _get_auto_quantize_score(grad_output, output_diff) +class _AutoQuantizeBackwardScoringSession(ABC): + """Manage temporary model state used by activation-backward scoring.""" + + def __init__( + self, + model: nn.Module, + score_modules: Sequence[nn.Module], + is_param_grad_enabled: Callable, + verbose: bool = False, + ) -> None: + self.model = model + self.score_modules = tuple(score_modules) + self.is_param_grad_enabled = is_param_grad_enabled + self.verbose = verbose + self._stack = ExitStack() + self._original_forwards: dict[nn.Module, Callable] = {} + self._output_grad_hook_handles: set[Any] = set() + self._grad_accumulators: list[Any] = [] + + def __enter__(self): + """Install scoring hooks and parameter settings.""" + try: + hparams = list( + dict.fromkeys( + hparam + for module in self.score_modules + for hparam in module._hparams_for_scoring + ) + ) + for hparam in hparams: + self._stack.callback(setattr, hparam, "active", hparam.active) + + def patched_forward(module, *args, **kwargs): + return self.forward(module, *args, **kwargs) + + for module in self.score_modules: + original_forward = module.forward + self._original_forwards[module] = original_forward + had_instance_forward = "forward" in module.__dict__ + instance_forward = module.__dict__.get("forward") + module.forward = types.MethodType(patched_forward, module) + if had_instance_forward: + self._stack.callback(setattr, module, "forward", instance_forward) + else: + self._stack.callback(module.__dict__.pop, "forward", None) + + for name, param in self.model.named_parameters(): + requires_grad = param.requires_grad + enable_grad = self.is_param_grad_enabled(name, self.model) + param.requires_grad = enable_grad + self._stack.callback(setattr, param, "requires_grad", requires_grad) + if not enable_grad: + continue + if self.verbose: + print_rank_0(f"AutoQuantize: Enabling gradient for param {name}.") + accumulator, hook = create_param_grad_clear_hook(param) + self._grad_accumulators.append(accumulator) + self._stack.callback(hook.remove) + except Exception: + self._stack.close() + raise + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + """Restore all model state changed for scoring.""" + self._clear_output_grad_hooks() + self._stack.close() + self._original_forwards.clear() + self._grad_accumulators.clear() + + def original_forward(self, module: nn.Module) -> Callable: + """Return the forward method saved before scoring.""" + return self._original_forwards[module] + + def _clear_output_grad_hooks(self) -> None: + """Remove output hooks whose backward pass has not run.""" + for handle in self._output_grad_hook_handles: + handle.remove() + self._output_grad_hook_handles.clear() + + def _register_output_grad_hook(self, output: torch.Tensor, hook: Callable) -> None: + """Attach an invocation-specific output-gradient hook for this session.""" + + def run_once(grad): + try: + return hook(grad) + finally: + handle.remove() + self._output_grad_hook_handles.discard(handle) + + handle = output.register_hook(run_once) + self._output_grad_hook_handles.add(handle) + + @abstractmethod + def forward(self, module: nn.Module, *args, **kwargs): + """Run a score module forward pass and collect method-specific state.""" + + +class _AutoQuantizeCandidateReplayScoringSession(_AutoQuantizeBackwardScoringSession): + """Share baseline execution, candidate replay, and score accumulation.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.no_quant = QuantRecipe(quant_cfg=None) + + def _run_unquantized(self, module: nn.Module, *args, **kwargs): + """Run a score module with each configurable group unquantized.""" + for hparam in module._hparams_for_scoring: + if hparam.is_configurable: + hparam.active = self.no_quant + output = self.original_forward(module)(*args, **kwargs) + base = output[0] if isinstance(output, tuple) else output + return output, base + + def _replay_candidates(self, module: nn.Module, base, candidate_recipes, *args, **kwargs): + """Measure output differences for the requested candidates.""" + output_diffs = {} + with torch.no_grad(): + for hparam in module._hparams_for_scoring: + if not hparam.is_configurable: + continue + recipe_diffs = {} + for recipe in candidate_recipes(hparam): + if recipe == self.no_quant: + continue + hparam.active = recipe + try: + replay = self.original_forward(module)(*args, **kwargs) + finally: + hparam.active = self.no_quant + replay = replay[0] if isinstance(replay, tuple) else replay + recipe_diffs[recipe] = (replay - base).detach() + if recipe_diffs: + output_diffs[hparam] = recipe_diffs + return output_diffs + + def _accumulate_candidate_scores(self, module, output_diffs, grad_output) -> None: + """Apply the method's score functional to replayed output differences.""" + if grad_output is None: + return + with torch.no_grad(): + for hparam, recipe_diffs in output_diffs.items(): + for recipe, output_diff in recipe_diffs.items(): + contribution = self._score_contribution(grad_output, output_diff) + current = hparam._importance_dict[recipe][module] + hparam._importance_dict[recipe][module] = ( + contribution if current is None else current + contribution + ) + + def _register_candidate_score_hook(self, module, output, output_diffs) -> None: + """Bind replayed candidate differences to one output invocation.""" + self._register_output_grad_hook( + output, + lambda grad_output: self._accumulate_candidate_scores( + module, output_diffs, grad_output + ), + ) + + @abstractmethod + def _score_contribution(self, grad_output, output_diff): + """Return this method's score contribution for one replayed candidate.""" + +class _AutoQuantizeGradientScoringSession(_AutoQuantizeCandidateReplayScoringSession): + """Collect gradient-based scores while candidate recipes are replayed.""" + + def forward(self, module: nn.Module, *args, **kwargs): + """Run the reference forward and cache each recipe's output perturbation.""" + output, base = self._run_unquantized(module, *args, **kwargs) + + # Checkpointed modules recompute with gradients enabled during backward. + if not torch.is_grad_enabled() or not base.requires_grad: + return output -class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): + output_diffs = self._replay_candidates( + module, base, lambda hparam: hparam.choices, *args, **kwargs + ) + self._register_candidate_score_hook(module, base, output_diffs) + return output + + def _score_contribution(self, grad_output, output_diff): + return _get_auto_quantize_score(grad_output, output_diff) + + +class _AutoQuantizeBackwardScoringSearcher(_AutoQuantizeBaseSearcher): + """Share orchestration used by activation-backward scoring methods.""" + + score_module_rules = [ + # Score MoE projections together at their enclosing MLP or mixer output. + r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", + r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", + r"^(.*?)\.(\d+\.(w1|w2|w3))$", + r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", + ] + + _custom_support: list[tuple[Callable, Callable, Callable]] = [] + + @classmethod + def register_custom_support( + cls, + is_supported_checker: Callable, + grad_ckpt_context: Callable, + is_param_grad_enabled: Callable, + ) -> None: + """Register optional hooks for memory-efficient backward scoring. + + `is_supported_checker` selects models that use these hooks. + `grad_ckpt_context` enables their gradient-checkpointing context, and + `is_param_grad_enabled` selects the minimum parameters needed to propagate + activation gradients. + """ + cls._custom_support.append((is_supported_checker, grad_ckpt_context, is_param_grad_enabled)) + + def _configurable_score_modules(self) -> list[nn.Module]: + return [ + module + for module in self.model.modules() + if hasattr(module, "_hparams_for_scoring") + and any(hparam.is_configurable for hparam in module._hparams_for_scoring) + ] + + @abstractmethod + def _estimate_auto_quantize_scores(self, is_param_grad_enabled: Callable) -> None: + """Estimate scores while activation gradients are enabled.""" + + def estimate_sensitivity_scores(self) -> None: + """Run backward scoring with the first matching model-specific support hook.""" + self.model.eval() + + def default_is_param_grad_enabled(_name, _model): + return True + + grad_checkpointing_context = None + is_param_grad_enabled = default_is_param_grad_enabled + for is_supported, context_candidate, grad_candidate in self._custom_support: + if is_supported(self.model): + grad_checkpointing_context = context_candidate + is_param_grad_enabled = grad_candidate + break + + context = ( + grad_checkpointing_context(self.model) + if grad_checkpointing_context is not None + else nullcontext() + ) + with context: + self._estimate_auto_quantize_scores(is_param_grad_enabled) + + +class AutoQuantizeGradientSearcher(_AutoQuantizeBackwardScoringSearcher): """A searcher for AutoQuantize algorithm that uses gradient based score estimation. In AutoQuantize, we search for the best per-layer quantization configuration that minimizes the sum of per-layer @@ -1472,17 +1787,6 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): method_name = "gradient" - score_module_rules = [ - # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts) - r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", - r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts - r"^(.*?)\.(\d+\.(w1|w2|w3))$", # mixtral experts - r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$", # dbrx experts - ] - - # See `register_custom_support` for details - _custom_support: list[tuple[Callable, Callable, Callable]] = [] - @property def default_search_config(self): """Get the default config for the searcher.""" @@ -1500,7 +1804,8 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: """Sanitize the search config dict.""" config = config or {} if "score_func" in config: - warnings.warn("`score_func` is ignored for gradient based `auto_quantize`.") + if config["score_func"] is not None: + warnings.warn("`score_func` is ignored for gradient based `auto_quantize`.") config.pop("score_func") config = super().sanitize_search_config(config) if config["forward_backward_step"] is None: @@ -1511,30 +1816,6 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig: return config - @classmethod - def register_custom_support( - cls, - is_supported_checker: Callable, - grad_ckpt_context: Callable, - is_param_grad_enabled: Callable, - ) -> None: - """(Optional) Register custom support for `AutoQuantize` score estimation. - - This custom support is used to enable memory/compute efficient backward gradient propagation. This involves: - - - `grad_ckpt_context`: backward pass with gradient checkpointing enabled - - `is_param_grad_enabled`: AutoQuantize only needs activation gradients to be computed (not weight - gradients). `is_param_grad_enabled` is used to select which parameters should have gradients enabled, - limiting gradient computation to only what's needed for activation gradients. For LLMs, to trigger all - activation gradient computation, just enabling the embedding layer weight gradient is sufficient. This will - enable gradient computation for all the activation gradients downstream. - - If the `is_supported_checker(model)` returns True, the `grad_ckpt_context(model)` will be - used to enable gradient checkpointing and `is_param_grad_enabled(pname, model)` - will be used to select which parameters have gradients enabled to minimize gradient computation. - """ - cls._custom_support.append((is_supported_checker, grad_ckpt_context, is_param_grad_enabled)) - def _get_default_forward_backward_step(self): def forward_backward_step(model, data): output = self.config["forward_step"](model, data) @@ -1552,195 +1833,42 @@ def forward_backward_step(model, data): @torch.enable_grad() def _estimate_auto_quantize_scores(self, is_param_grad_enabled): - # TODO: remove the no-quant recipe - def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): - for hparam in module._hparams_for_scoring: - if hparam.is_configurable: - hparam.active = QuantRecipe(quant_cfg=None) - - output = module._forward_original(input, *args, **kwargs) - - # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass. - # With gradient checkpointing, gradients are computed in the local forward pass during backward pass - - # Lets compute the output_diff and save it in memory only if gradient is enabled to be memory efficient - if not torch.is_grad_enabled(): - return output - - module.output_diff_dict = {hparam: {} for hparam in module._hparams_for_scoring} - with torch.no_grad(): - for hparam in module._hparams_for_scoring: - if not hparam.is_configurable: - continue - for recipe in hparam.choices: - if recipe == QuantRecipe(quant_cfg=None): - continue - hparam.active = recipe - output_diff = module._forward_original(input, *args, **kwargs) - - if isinstance(output_diff, tuple): - output_diff = output_diff[0] - output[0] - else: - output_diff -= output - module.output_diff_dict[hparam][recipe] = output_diff.detach() - - # Disable the configurable hparam now that we have computed the diff - hparam.active = QuantRecipe(quant_cfg=None) - - return output - - def backward_hook(module, grad_input, grad_output): - for hparam, output_diff_dict in module.output_diff_dict.items(): - for recipe, output_diff in output_diff_dict.items(): - if hparam._importance_dict[recipe][module] is None: - hparam._importance_dict[recipe][module] = _get_auto_quantize_score( - grad_output[0], output_diff - ) - else: - _add_auto_quantize_score( - grad_output[0], output_diff, hparam._importance_dict[recipe][module] - ) - - def setup_params_for_score_estimation(name, param, params_metadata, enable_grad=True): - # Let us delete the gradient as soon as they are computed to save memory - params_metadata[name] = {"requires_grad": param.requires_grad} - param.requires_grad = enable_grad - if not enable_grad: - return - if self.config.get("verbose", False): - print_rank_0(f"AutoQuantize: Enabling gradient for param {name}.") - accum_grad, handle = create_param_grad_clear_hook(param) - params_metadata[name]["accum_grad"] = accum_grad # We need to keep the accum_grad alive - params_metadata[name]["handle"] = handle - - def setup_module_for_score_estimation(module): - module._forward_original = module.forward - module.forward = types.MethodType(auto_quantize_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() - - def cleanup_params_after_score_estimation(name, param, params_metadata): - param.requires_grad = params_metadata[name]["requires_grad"] - handle = params_metadata[name].get("handle", None) - if handle is not None: - handle.remove() - - score_modules = set() - for name, module in self.model.named_modules(): - if ( - hasattr(module, "_hparams_for_scoring") - and any(hparam.is_configurable for hparam in module._hparams_for_scoring) - and module not in score_modules - ): - # Monkey patch the forward methods to cache (Q(Y) - Y) - setup_module_for_score_estimation(module) - score_modules.add(module) - - params_metadata = {} - for name, param in self.model.named_parameters(): - setup_params_for_score_estimation( - name, param, params_metadata, is_param_grad_enabled(name, self.model) + score_modules = self._configurable_score_modules() + with _AutoQuantizeGradientScoringSession( + self.model, + score_modules, + is_param_grad_enabled, + verbose=self.config.get("verbose", False), + ) as scoring_session: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + report_memory("AutoQuantize: starting score estimation, ") + + def score_step(model, data): + try: + return self.config["forward_backward_step"](model, data) + finally: + scoring_session._clear_output_grad_hooks() + + self._run_func( + score_step, + num_iters=self.config["num_score_steps"], + desc="Estimating auto_quantize scores", ) - gc.collect() - if torch.cuda.is_available(): - torch.cuda.reset_peak_memory_stats() - report_memory("AutoQuantize: starting score estimation, ") - - self._run_func( - self.config["forward_backward_step"], - num_iters=self.config["num_score_steps"], - desc="Estimating auto_quantize scores", - ) - - if torch.cuda.is_available(): - report_memory("AutoQuantize: After score estimation") + if torch.cuda.is_available(): + report_memory("AutoQuantize: After score estimation") - for module in score_modules: - cleanup_module_after_score_estimation(module) - - for name, param in self.model.named_parameters(): - cleanup_params_after_score_estimation(name, param, params_metadata) - - # Delete the params_metadata - del params_metadata gc.collect() - def estimate_sensitivity_scores(self) -> None: - """Estimate sensitivity scores using hessian approximation.""" - self.model.eval() - - def _default_is_param_grad_enabled(pname, model): - return True - - grad_checkpointing_ctxt = None - is_param_grad_enabled = _default_is_param_grad_enabled - for is_supported_checker, ctxt_candidate, grad_enabled_candidate in self._custom_support: - if is_supported_checker(self.model): - grad_checkpointing_ctxt = ctxt_candidate - is_param_grad_enabled = grad_enabled_candidate - break - - with grad_checkpointing_ctxt(self.model) if grad_checkpointing_ctxt else nullcontext(): - self._estimate_auto_quantize_scores(is_param_grad_enabled) - def run_search_with_stats(self, max_weight_size, verbose=False): """Linear Programming Solve for gradient based auto_quantize. AutoQuantize uses Linear Programming Solver to find the optimal quantization configuration which minimizes the sum of per-layer auto_quantize scores while meeting the specified constraint. """ - # TODO: Do this only for rank 0 in the respective pipeline group - - for lower_bound in self._get_search_lower_bounds(): - # The LP solver for auto_quantize sometimes fails to find a solution if a lower bound is not - # specified. I dont know why this happens. - # As a workaround, lets specify a lower bound for the weight compression if previous - # search without lower bound fails. - constraints, constraint_name = self._get_constraints_for_search( - max_weight_size, lower_bound - ) - - lps = LPS( - name="AutoQuantize", - constraints=constraints, - constraints_to_candidate_costs={ - constraint_name: [ - candidate_stat["costs"] for candidate_stat in self.candidate_stats.values() - ] - }, - candidate_scores=[ - candidate_stat["scores"] for candidate_stat in self.candidate_stats.values() - ], - objective_type="minimize", - verbose=verbose, - ) - selections, self.status = lps() - if self.status == "Optimal": - break - - if self.status != "Optimal": - warnings.warn( - "AutoQuantize FAILED to find a solution! The searched model might not meet all constraints. " - ) - is_satisfied = False - else: - is_satisfied = True - - best_recipes = {} - for name, selected_idx in zip(self.candidate_stats.keys(), selections): - best_recipes[name] = { - "format": self.candidate_stats[name]["formats"][selected_idx], - "costs": self.candidate_stats[name]["costs"][selected_idx], - "scores": self.candidate_stats[name]["scores"][selected_idx], - } - - return best_recipes, is_satisfied + return self._run_linear_program_search(max_weight_size, verbose) @torch.compile(dynamic=True) @@ -1946,6 +2074,12 @@ def run_search_with_stats(self, max_weight_size, verbose=False): # Backward compatibility alias (defaults to gradient-based searcher) AutoQuantizeSearcher = AutoQuantizeGradientSearcher +# Registry of AutoQuantize scoring methods. Optional methods register on import. +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,15 +2212,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) searcher.cost = search_state.get("cost", {}) @@ -2103,8 +2234,10 @@ def _resolve_best_recipe(search_state, constraints, verbose=False): "cost": searcher.cost, "active_moe_expert_ratio": searcher.active_moe_expert_ratio, } + 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()} if verbose: total_cost = sum(info["costs"] for info in best_recipe_info.values()) diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 3f6040fd4ef..9baf3b4f551 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -38,7 +38,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 @@ -285,12 +287,14 @@ 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. ``auto_quantize`` uses sensitivity scores to rank the per-layer quantization formats and search - for the best quantization formats per-layer. The sensitivity score can be computed using gradient-based - methods (default) or KL divergence loss, controlled by the ``method`` parameter. + for the best quantization formats per-layer. The sensitivity score can be computed with + gradient-based methods (default), KL divergence loss, or Aumann-Shapley path-integral + attributions, controlled by the ``method`` parameter. Internally this API runs two main phases: @@ -442,9 +446,15 @@ 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"`` (path-integral damage attributions, calibrated against a + directly measured reference point; label-free like ``"kl_div"``, and additionally + reports a ``predicted_damage`` estimate for the selected recipe. Scoring passes grow + with the number of candidate formats and path nodes, not with the number of + whole-model configurations the search considers -- 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. @@ -462,6 +472,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. @@ -621,18 +634,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, @@ -647,13 +654,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..d067a0da6c8 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -15,6 +15,7 @@ import copy import io +import warnings from types import SimpleNamespace import pytest @@ -36,6 +37,7 @@ QuantRecipe, QuantRecipeHparam, _AutoQuantizeBaseSearcher, + _AutoQuantizeGradientScoringSession, _module_search_space_signature, estimate_quant_compression, ) @@ -98,6 +100,47 @@ def get_input(self): return torch.randn(1, 4, 32) +class _ScoredMoeExpert(torch.nn.Module): + def __init__(self): + super().__init__() + self.gate_proj = torch.nn.Linear(8, 8) + self.up_proj = torch.nn.Linear(8, 8) + self.down_proj = torch.nn.Linear(8, 8) + + def forward(self, x): + return self.down_proj(self.gate_proj(x) + self.up_proj(x)) + + +class _ScoredMoeMlp(torch.nn.Module): + def __init__(self): + super().__init__() + self.experts = torch.nn.ModuleList([_ScoredMoeExpert(), _ScoredMoeExpert()]) + + def forward(self, x): + output = torch.zeros_like(x) + for expert in self.experts: + output = output + expert(x) + return output + + +class _ScoredMoeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.mlp = _ScoredMoeMlp() + + def forward(self, x): + return self.mlp(x) + + +class _ScoredMoeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList([_ScoredMoeLayer()]) + + def forward(self, x): + return self.layers[0](x) + + @pytest.mark.parametrize( ("quant_cfg", "other_quant_cfg", "is_less_than"), [ @@ -577,6 +620,29 @@ def test_auto_quantize_rejects_empty_module_formats(formats): ) +def test_gradient_search_config_none_score_func_does_not_warn(): + """An explicitly empty score_func must not emit the ignored-value warning.""" + searcher = AutoQuantizeGradientSearcher() + + def forward_backward_step(model, data): + pass + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + config = searcher.sanitize_search_config( + { + "score_func": None, + "data_loader": [object()], + "forward_step": lambda model, data: None, + "forward_backward_step": forward_backward_step, + } + ) + + assert not any("`score_func` is ignored" in str(warning.message) for warning in caught) + assert "score_func" not in config + assert config["forward_backward_step"] is forward_backward_step + + def test_auto_quantize_fixed_module_isolated_from_unrelated_calibration(monkeypatch): model = TransformerBlock() calibration_states = [] @@ -736,7 +802,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() @@ -905,6 +971,229 @@ def test_data_parallel_auto_quantize(skip_on_windows): spawn_multiprocess_job(2, _test_data_parallel_auto_quantize, backend="gloo") +def _test_data_parallel_moe_score_module(rank, size): + torch.manual_seed(1234) + model = _ScoredMoeModel() + data_loader = [torch.randn(2, 3, 8) for _ in range(2)] + model, search_history = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=data_loader, + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, data: output.square().mean(), + num_calib_steps=2, + num_score_steps=2, + ) + + hparam = model.layers[0].mlp.experts[0].gate_proj.get_hparam("quant_recipe") + assert hparam.score_modules == [model.layers[0].mlp] + assert isinstance(model.layers[0].mlp._hparams_for_scoring, list) + + recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + local_score = sum(hparam._importance_dict[recipe][m] for m in hparam.score_modules) + candidate = next( + candidate + for candidate in search_history["candidate_stats"].values() + if "layers.0.mlp.experts.0.gate_proj" in candidate["module_names"] + ) + recipe_idx = candidate["formats"].index(recipe) + torch.testing.assert_close( + local_score * size, + torch.tensor( + candidate["scores"][recipe_idx], + device=local_score.device, + dtype=local_score.dtype, + ), + ) + + scores = { + name: candidate["scores"] for name, candidate in search_history["candidate_stats"].items() + } + rank_zero_scores = DistributedProcessGroup.get_dist_syncd_obj( + scores if rank == 0 else None, + DistributedProcessGroup(None), + lambda values: values[0], + ) + assert scores == rank_zero_scores + + +def test_data_parallel_moe_score_module(skip_on_windows): + spawn_multiprocess_job(2, _test_data_parallel_moe_score_module, backend="gloo") + + +def test_score_hparam_registration_preserves_order(): + quant_modules = [ + mtq.quantize(torch.nn.Linear(4, 4), mtq.INT8_DEFAULT_CFG), + mtq.quantize(torch.nn.Linear(4, 4), mtq.INT8_DEFAULT_CFG), + ] + score_module = torch.nn.Identity() + recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + + first = QuantRecipeHparam( + [recipe], + quant_modules=[quant_modules[0], quant_modules[1], quant_modules[0]], + score_modules=[score_module, score_module], + ) + second = QuantRecipeHparam( + [recipe], + quant_modules=[quant_modules[1]], + score_modules=[score_module], + ) + + assert first.quant_modules == quant_modules + assert first.score_modules == [score_module] + assert score_module._hparams_for_scoring == [first, second] + + +def test_gradient_scoring_tracks_reused_module_invocations(): + """Each autograd use of a shared score module retains its own replay difference.""" + no_quant_recipe = QuantRecipe(quant_cfg=None) + quant_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + + class TestHparam: + is_configurable = True + choices = [no_quant_recipe, quant_recipe] + active = no_quant_recipe + + class ScoreModule(torch.nn.Module): + def forward(self, x): + scale = 1.0 if hparam.active == no_quant_recipe else 2.0 + return scale * x + + class ReusedScoreModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.shared = ScoreModule() + + def forward(self, x): + first = self.shared(x) + second = self.shared(3.0 * x) + # This use does not participate in autograd. + self.shared(5.0 * x.detach()) + # This output requires grad but is intentionally unused by the loss. + self.shared(7.0 * x) + return first.sum() + 2.0 * second.sum() + + model = ReusedScoreModule() + score_module = model.shared + hparam = TestHparam() + hparam._importance_dict = {recipe: {score_module: None} for recipe in hparam.choices} + score_module._hparams_for_scoring = [hparam] + inputs = torch.tensor([[1.0, 2.0]], requires_grad=True) + + hparam._importance_dict[quant_recipe][score_module] = None + delayed_session = _AutoQuantizeGradientScoringSession(model, [score_module], lambda *_: True) + with delayed_session: + delayed_loss = model(inputs) + assert delayed_session._output_grad_hook_handles + + assert not delayed_session._output_grad_hook_handles + delayed_loss.backward() + assert hparam._importance_dict[quant_recipe][score_module] is None + + with _AutoQuantizeGradientScoringSession(model, [score_module], lambda *_: True): + model(inputs).backward() + + # First use: sum(x**2) = 5. Second use: sum((2 * 3x)**2) = 180. + importance = hparam._importance_dict[quant_recipe][score_module] + torch.testing.assert_close( + importance, + torch.tensor(185.0, device=importance.device), + rtol=0, + atol=1e-6, + ) + + +def test_gradient_scoring_restores_model_after_failure(): + model = SimpleLinear() + patched_modules = [] + + def fail_during_scoring(model, data): + model(data) + patched_modules.extend( + module + for module in model.modules() + if getattr(module.forward, "__name__", None) == "patched_forward" + ) + raise RuntimeError("stop after scoring forward") + + with pytest.raises(RuntimeError, match="stop after scoring forward"): + mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[model.get_input()], + forward_step=lambda model, batch: model(batch), + forward_backward_step=fail_during_scoring, + num_calib_steps=1, + num_score_steps=1, + ) + + assert patched_modules + assert all( + getattr(module.forward, "__name__", None) != "patched_forward" for module in patched_modules + ) + assert all("forward" not in module.__dict__ for module in patched_modules) + assert all(param.requires_grad for param in model.parameters()) + for module in model.modules(): + for hparam in getattr(module, "_hparams_for_scoring", []): + assert hparam.active == hparam.original + + +def test_backward_scoring_session_restores_partial_setup(): + model = torch.nn.Sequential(torch.nn.Linear(4, 4)) + score_module = model[0] + score_module._hparams_for_scoring = [] + score_module.weight.requires_grad = False + original_requires_grad = {name: param.requires_grad for name, param in model.named_parameters()} + + def fail_on_second_parameter(name, _model): + if name.endswith("bias"): + raise RuntimeError("stop during scoring setup") + return True + + session = _AutoQuantizeGradientScoringSession( + model, + [score_module], + fail_on_second_parameter, + ) + with pytest.raises(RuntimeError, match="stop during scoring setup"), session: + pytest.fail("scoring setup should not complete") + + assert "forward" not in score_module.__dict__ + assert { + name: param.requires_grad for name, param in model.named_parameters() + } == original_requires_grad + + +@pytest.mark.parametrize("instance_override", [False, True]) +def test_gradient_scoring_restores_forward_attribute_layout(instance_override): + module = torch.nn.Identity() + module._hparams_for_scoring = [] + + def original_forward(x): + return x + 1 + + original_override = original_forward + if instance_override: + module.forward = original_override + + session = _AutoQuantizeGradientScoringSession( + module, + [module], + lambda _name, _model: False, + ) + with pytest.raises(RuntimeError, match="stop during scoring"), session: + assert module.__dict__["forward"] is not original_override + raise RuntimeError("stop during scoring") + + if instance_override: + assert module.__dict__["forward"] is original_override + else: + assert "forward" not in module.__dict__ + + def test_auto_quantize_budget_uses_no_quant_candidate_cost(monkeypatch): class _BudgetCaptureSearcher(AutoQuantizeGradientSearcher): def run_search_with_stats(self, max_weight_size, verbose=False): @@ -1085,7 +1374,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() 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..042d91bf725 --- /dev/null +++ b/tests/unit/torch/quantization/test_autoquant_shapley.py @@ -0,0 +1,1180 @@ +# 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, +damage-bound search, and use of the shared linear-program solver. +""" + +import copy +import functools +import itertools +import math +from collections import namedtuple + +import numpy as np +import pytest +import torch +from _test_utils.torch.distributed.utils import spawn_multiprocess_job + +import modelopt.torch.quantization as mtq +import modelopt.torch.quantization.model_quant as _model_quant +from modelopt.torch.quantization._auto_quantize_shapley import ( + AutoQuantizeAumannShapleySearcher, + _anchor_ceiling, + _as_seed_coverage, + _predict_damage, +) +from modelopt.torch.quantization.algorithms import QuantRecipe, _AutoQuantizeBackwardScoringSearcher +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.utils.distributed import DistributedProcessGroup + +SEARCH_FORMATS = [mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG] + + +class _Attention(torch.nn.Module): + """Minimal attention block.""" + + 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): + """Minimal transformer block.""" + + def __init__(self, seed=0): + super().__init__() + with torch.random.fork_rng(devices=[]): + torch.manual_seed(seed) + self.attn = _Attention() + self.mlp = torch.nn.Linear(32, 32) + self._input_seed = seed + 1000 + + def forward(self, x): + return self.mlp(self.attn(x)) + + def get_input(self): + generator = torch.Generator().manual_seed(self._input_seed) + return torch.randn(1, 4, 32, generator=generator) + + +class _OneLinear(torch.nn.Module): + """Single-linear model.""" + + def __init__(self): + super().__init__() + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + self.fc = torch.nn.Linear(32, 32) + + def forward(self, x): + return self.fc(x) + + def get_input(self): + generator = torch.Generator().manual_seed(1000) + return torch.randn(1, 4, 32, generator=generator) + + +def _search(model, method="aumann_shapley", effective_bits=6.0, method_options=None, **kwargs): + """Run auto_quantize with the given method.""" + 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(): + """Shared aumann_shapley search state.""" + _model, state = _search(_Block(), method_options={"num_path_nodes": 2}) + return state + + +def test_damage_model_and_score_monotonicity(shapley_state): + """Damage model and score monotonicity.""" + 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 applies and resolve tightens.""" + 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_damage_bound_mode_respects_the_quote(shapley_state): + """Damage-bound mode keeps the selected recipe within the requested estimate.""" + epsilon = 0.5 * shapley_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): + """Searcher preloaded with synthetic candidate stats.""" + 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): + """Total score and cost summed over the selected entries.""" + score = cost = 0.0 + for info in best.values(): + score += info["scores"] + cost += info["costs"] + return score, cost + + +def test_damage_bound_search_is_exact(): + """Damage-bound LPS matches exhaustive allocation.""" + 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 == pytest.approx(optimal, rel=1e-9) + + +def _brute_force_min_score(stats, budget): + """Exhaustive minimum-score allocation.""" + 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_linear_program_search_is_exact(): + """The shared linear-program search matches exhaustive allocation.""" + for seed in range(3): + searcher = _synthetic_searcher(seed=seed) + stats = searcher.candidate_stats + total = sum(stat["uncompressed_cost"] for stat in stats.values()) + for fraction in (0.9, 0.5): + budget = total * fraction + best, is_satisfied = searcher.run_search_with_stats(budget) + assert is_satisfied + score, cost = _selected(searcher, best) + assert cost <= budget + 1e-6 + assert score == pytest.approx(_brute_force_min_score(stats, budget), rel=1e-9) + + +def test_coverage_inversion_recovers_forward_model(): + """Coverage inversion recovers coefficients integrated over the full path.""" + c = 0.9 + a_true = np.array([0.05, 0.1, 0.2, 0.3, 0.4, 0.5]) + attributions = [] + for index, a_i in enumerate(a_true): + coefficients = np.array([1.0]) + for a_j in np.delete(a_true, index): + coefficients = np.convolve(coefficients, [1.0, -a_j]) + integral = sum( + coefficient / (degree + 1) for degree, coefficient in enumerate(coefficients) + ) + attributions.append(c * a_i * integral) + attributions = np.asarray(attributions) + + a, b, converged = _as_seed_coverage(attributions, c=c) + assert converged, a + assert np.allclose(a, a_true, rtol=1e-5) + assert np.allclose(b, -np.log1p(-a_true), rtol=1e-5) + + +def test_anchor_ceiling_honors_max_inflation(): + """A fit that only converges beyond the ceiling backstop must remain invalid.""" + ceiling, _b, _kappa, inflation, converged = _anchor_ceiling( + {"fmt": np.array([10.0])}, + f_corner=1.0, + corner_mask_by_key={"fmt": np.array([True])}, + max_inflation=10.0, + ) + + assert ceiling == pytest.approx(10.0) + assert inflation == pytest.approx(10.0) + assert not converged + + +def test_both_targets_rejected_before_model_conversion(): + """Both targets rejected before model conversion.""" + + 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 test_damage_bound_searcher_accepts_none_constraints(monkeypatch): + """The searcher treats a missing constraint mapping like an empty mapping.""" + searcher = AutoQuantizeAumannShapleySearcher() + searcher.model = _Block() + searcher.constraints = None + searcher.config = { + **searcher.default_search_config, + "max_predicted_damage": 0.1, + } + searcher.candidate_stats = {} + monkeypatch.setattr(_AutoQuantizeBackwardScoringSearcher, "before_search", lambda self: None) + + searcher.before_search() + + assert searcher.constraints == {"effective_bits": 16.0} + + +def test_damage_bound_search_ignores_fixed_group_scores(): + """A fixed group's score cannot consume the configurable allocation's damage budget.""" + aggressive = QuantRecipe("INT4_BLOCKWISE_WEIGHT_ONLY_CFG") + moderate = QuantRecipe("INT8_DEFAULT_CFG") + no_quant = QuantRecipe(quant_cfg=None) + searcher = _synthetic_searcher(n_groups=1) + searcher.candidate_stats["g0.quant_recipe"].update( + formats=[aggressive, moderate, no_quant], + scores=[0.2, 0.1, 0.0], + costs=[4.0, 8.0, 16.0], + uncompressed_cost=16.0, + ) + searcher.candidate_stats["fixed.quant_recipe"] = { + "formats": [moderate], + "scores": [10.0], + "costs": [8.0], + "module_names": ["fixed"], + "quantizer_attrs": {"fixed": ("input_quantizer", "weight_quantizer")}, + "cost_weight": 1.0, + "allow_no_quant": False, + "is_fixed": True, + "uncompressed_cost": 16.0, + } + searcher.damage_model = {"link": "additive", "valid": True} + searcher.config["max_predicted_damage"] = 0.1 + + best, is_satisfied = searcher.run_search_with_stats(max_weight_size=np.inf) + + assert is_satisfied + assert best["g0.quant_recipe"]["format"] == moderate + assert best["fixed.quant_recipe"]["format"] == moderate + + +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 + # An invalidated fit leaves c = 0.0; the quote must not read as "zero damage". + assert math.isnan(state["best"]["predicted_damage"]) + + +def test_nonconvergent_coverage_inversion_is_flagged(monkeypatch): + """Persist the reason a finite coverage fit could not be inverted.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": 0.2, "INT8_DEFAULT_CFG": 0.2}, + corner=0.01, + ) + with pytest.warns(UserWarning, match="inversion_not_converged"): + _model, state = _search(_OneLinear(), effective_bits=16.0) + + damage_model = state["damage_model"] + assert damage_model["inversion_converged"] is False + assert damage_model["valid"] is False + assert "inversion_not_converged" in damage_model["approximation_flags"], damage_model + + +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(): + """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} + _inject_scores_and_corner(monkeypatch, injected, corner=0.4) + _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 + + +@pytest.mark.parametrize( + ("method", "options", "exception"), + [ + ("aumann_shapley", [("num_path_nodes", 2)], TypeError), + ("aumann_shapley", {"unknown_option": 1}, ValueError), + ("aumann_shapley", {"num_score_steps": 999}, ValueError), + ("aumann_shapley", {"num_path_nodes": True}, ValueError), + ("aumann_shapley", {"num_path_nodes": 1.5}, ValueError), + ("aumann_shapley", {"num_path_nodes": 0}, ValueError), + ("aumann_shapley", {"damage_link": "unsupported"}, ValueError), + ("aumann_shapley", {"solver": "lp"}, ValueError), + ("aumann_shapley", {"max_predicted_damage": float("inf")}, ValueError), + ("aumann_shapley", {"max_predicted_damage": -1.0}, ValueError), + ("kl_div", {"num_path_nodes": 2}, ValueError), + ], +) +def test_invalid_method_options_leave_model_untouched(method, options, exception): + """Reject invalid method options before converting the model.""" + model = _Block() + with pytest.raises(exception): + _search(model, method=method, method_options=options) + assert type(model.mlp) is torch.nn.Linear + assert not any(isinstance(module, TensorQuantizer) for module 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_zero_attributions_invert_to_exact_zero(): + """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): + """Scoring signature guards resume.""" + 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 the solve target reuses the stored scores. + _model, state = _search( + _Block(), + effective_bits=None, + checkpoint=checkpoint, + method_options={ + "num_path_nodes": 2, + "max_predicted_damage": 1.0, + }, + ) + assert state["best"]["is_satisfied"] + + +def _shapley_data_parallel(rank, size, baseline): + + _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): + """Data parallel aumann shapley.""" + _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, functools.partial(_shapley_data_parallel, baseline=baseline), backend="gloo" + ) + + +def test_anchor_ceiling_rejects_non_finite_measurements(): + """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): + """All candidates non finite falls back to no quant.""" + _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"] + + +@pytest.mark.parametrize("structured_output", [False, True], ids=["tensor", "namedtuple"]) +def test_nested_score_modules_are_scored(structured_output): + """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. + + The namedtuple case also ensures path shifting preserves structured output attributes. + """ + + mlp_output = namedtuple("MlpOutput", ["hidden_states"]) + + class _Expert(torch.nn.Module): + """Minimal expert block.""" + + 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): + """Minimal MoE MLP container.""" + + def __init__(self): + super().__init__() + self.experts = torch.nn.ModuleList([_Expert() for _ in range(2)]) + self.shared_experts = _Expert() + + def forward(self, hidden_states): + out = self.shared_experts(hidden_states) + for expert in self.experts: + out = out + expert(hidden_states) + return mlp_output(out) if structured_output else out + + class _Layer(torch.nn.Module): + """Minimal decoder layer.""" + + def __init__(self): + super().__init__() + self.mlp = _MLP() + + def forward(self, x): + output = self.mlp(hidden_states=x) + return output.hidden_states if structured_output else output + + class _Model(torch.nn.Module): + """Minimal model wrapper.""" + + 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) + ) + + # Assert the routing, not just that scores exist: this test kept passing when the + # score-module rules were dropped, because positive scores say nothing about where the + # attribution was measured. Routed experts share the MLP container; the shared expert's + # projections score at themselves, which is what creates the nesting. + routed = model.layer.mlp.experts[0].gate_proj.get_hparam("quant_recipe") + shared = model.layer.mlp.shared_experts.gate_proj.get_hparam("quant_recipe") + assert routed.score_modules == [model.layer.mlp] + assert set(shared.score_modules) == { + model.layer.mlp.shared_experts.gate_proj, + model.layer.mlp.shared_experts.up_proj, + } + + # 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_reused_score_module_accumulates_every_invocation(): + """Used calls of a reused score module must be attributed independently.""" + + class _ReusedLinear(_OneLinear): + def forward(self, x): + first = self.fc(x) + second = self.fc(x) + self.fc(torch.zeros_like(x)) + return first + second + + class _ScaledLinear(_OneLinear): + def forward(self, x): + return 2 * self.fc(x) + + repeated = _ReusedLinear() + scaled = _ScaledLinear() + scaled.load_state_dict(repeated.state_dict()) + data = [torch.randn(1, 4, 32) for _ in range(2)] + + def run(model): + _model, state = mtq.auto_quantize( + model, + constraints={"effective_bits": 12.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=data, + forward_step=lambda model, batch: model(batch), + num_calib_steps=2, + num_score_steps=2, + method="aumann_shapley", + ) + (stat,) = state["candidate_stats"].values() + return next( + score + for recipe, score in zip(stat["formats"], stat["raw_scores"], strict=True) + if not recipe.is_no_quant + ) + + assert run(repeated) == pytest.approx(run(scaled), rel=1e-5, abs=1e-8) + + +def test_model_specific_backward_support_is_used(monkeypatch): + """Aumann-Shapley should use the same model-specific backward setup as gradient scoring.""" + calls = {"context": 0, "parameters": 0} + + class _TrackingContext: + def __enter__(self): + calls["context"] += 1 + + def __exit__(self, *_exc_info): + calls["context"] += 1 + + def is_param_grad_enabled(_name, _model): + calls["parameters"] += 1 + return True + + monkeypatch.setattr( + AutoQuantizeAumannShapleySearcher, + "_custom_support", + [(lambda _model: True, lambda _model: _TrackingContext(), is_param_grad_enabled)], + ) + + _search(_OneLinear()) + + assert calls["context"] == 2 + assert calls["parameters"] > 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 = [] + + real_calibrate = _model_quant.calibrate + + def _spy(*args, **kwargs): + """Record that calibration ran.""" + 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) + ) + + +def test_no_scored_tokens_still_reports_an_invalid_quote(): + """Every failure path must signal through predicted_damage, never omit it. + + With no scored tokens the fit cannot be built at all. Returning without a damage model + would leave predicted_damage unset, so a consumer reading the documented key would get a + KeyError rather than the nan/valid=False signal the other failure paths produce. + """ + model = _OneLinear() + _m, state = mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=list(SEARCH_FORMATS), + data_loader=[model.get_input() for _ in range(2)], + forward_step=lambda m, batch: m(batch), + num_calib_steps=2, + num_score_steps=0, + method="aumann_shapley", + ) + + assert math.isnan(state["best"]["predicted_damage"]) + assert state["best"]["predicted_damage_valid"] is False + damage_model = state["damage_model"] + assert "no_scored_tokens" in damage_model["approximation_flags"] + # Same key shape as a normal run, so the documented contract does not KeyError. + assert math.isnan(damage_model["completeness"]) + for key in ("link", "f_corner", "n_score_tokens", "damage_reference", "as_scores"): + assert key in damage_model + + +def test_additive_link_bound_is_certified_when_the_fit_is_valid(monkeypatch): + """A valid additive-link fit still certifies its bound.""" + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": 0.01, "INT8_DEFAULT_CFG": 0.02}, + corner=0.4, + ) + _model, state = _search( + _OneLinear(), + effective_bits=None, + method_options={"damage_link": "additive", "max_predicted_damage": 1.0}, + ) + + assert state["damage_model"]["valid"] is True + assert state["best"]["is_satisfied"] is True + assert not math.isnan(state["best"]["predicted_damage"]) + + +def test_additive_link_does_not_certify_an_invalid_fit(monkeypatch): + """An invalidated fit cannot certify a bound under EITHER link. + + The validity gate used to sit inside the coverage branch, so an additive-link search + reported is_satisfied=True while the quote was NaN -- a self-contradictory result. + """ + _inject_scores_and_corner( + monkeypatch, + {"INT4_BLOCKWISE_WEIGHT_ONLY_CFG": 0.01, "INT8_DEFAULT_CFG": 0.02}, + corner=float("inf"), + ) + _model, state = _search( + _OneLinear(), + effective_bits=None, + method_options={"damage_link": "additive", "max_predicted_damage": 1e-3}, + ) + + assert state["damage_model"]["valid"] is False + assert state["best"]["is_satisfied"] is False + assert math.isnan(state["best"]["predicted_damage"])