Add Aumann-Shapley AutoQuantize recipe integration

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.
 

examples/hf_ptq/README.md

@@ -366,15 +366,41 @@ export HF_PATH=<the downloaded checkpoint from the Hugging Face hub, or simply t
 scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits --calib_batch_size 4
 ```
 
-The recipe quantizes the less accuracy-sensitive layers with the more aggressive format (e.g. NVFP4) and
-keeps the more sensitive ones at higher precision (or unquantized), so the model meets the recipe's
+The recipe quantizes the less accuracy-sensitive layers with the more aggressive format (e.g. NVFP4)
+and keeps the more sensitive ones at higher precision (or unquantized), so the model meets the recipe's
 `effective_bits` target. To author your own, copy a shipped recipe and adjust `candidate_formats`,
-`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div`), `score_size`,
-`module_search_spaces` (optional per-module candidate overrides), `disabled_layers` (excluded from
-the search), and `cost_excluded_layers` (kept out of the bit-budget accounting — e.g. VL vision
-towers). Recipes can splice a shared base `disabled_layers` set via `$import` (see
+`constraints.effective_bits`, `auto_quantize_method` (`gradient`, `kl_div`, or `aumann_shapley`),
+`method_options`, `score_size`, `module_search_spaces` (optional per-module candidate overrides),
+`disabled_layers` (excluded from the search), and `cost_excluded_layers` (kept out of the bit-budget
+accounting — e.g. VL vision towers). Recipes can splice a shared base `disabled_layers` set via `$import` (see
 `modelopt_recipes/configs/auto_quantize/units/base_disabled_layers`).
 
+An Aumann-Shapley recipe can target effective bits:
+
+```yaml
+auto_quantize:
+  constraints:
+    effective_bits: 5.4
+  auto_quantize_method: aumann_shapley
+  method_options:
+    num_path_nodes: 2
+    damage_link: coverage
+```
+
+Or it can minimize weight cost while keeping predicted damage under a bound:
+
+```yaml
+auto_quantize:
+  constraints: {}
+  auto_quantize_method: aumann_shapley
+  method_options:
+    num_path_nodes: 2
+    max_predicted_damage: 0.01
+```
+
+These fragments omit the candidate formats and other shared recipe fields. Choose either
+`constraints.effective_bits` or `method_options.max_predicted_damage`; recipes that set both are rejected.
+
 AutoQuantize recipes support two mutually exclusive search-space styles:
 
 1. Set top-level `auto_quantize.candidate_formats` to search every unmatched quantizable module, with
@@ -437,7 +463,7 @@ The example scripts above also have an additional flag `--tasks`, where the actu
 
 > *If GPU out-of-memory error is reported running the scripts, please try editing the scripts and reducing the max batch size to save GPU memory.*
 
-> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` method. The `kl_div` method does not require backpropagation.*
+> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` or `aumann_shapley` methods. Aumann-Shapley does not need labels, but it still backpropagates its comparison loss. The `kl_div` method does not require backpropagation.*
 
 ## Real Quant
 

examples/hf_ptq/hf_ptq.py

@@ -349,6 +349,10 @@ def _mtq_inputs_from_auto_quantize_config(
     to ``--kv_cache_qformat`` when the recipe omits it.
     """
     constraints = aq_config.constraints.model_dump(exclude_none=True)
+    method_options = aq_config.method_options
+    if aq_config.uses_predicted_damage_target:
+        # The recipe validator rejects an explicit bit budget; remove only the schema default.
+        constraints.pop("effective_bits", None)
     # cost_excluded_layers (sibling of disabled_layers) maps to the mtq cost key: these layers are
     # kept out of the bit-budget denominator (cost_weight 0) — e.g. VL vision towers — distinct from
     # disabled_layers, which removes them from the search.
@@ -387,6 +391,7 @@ def _mtq_inputs_from_auto_quantize_config(
         "disabled_layers": aq_config.disabled_layers,
         "kv_cache_quant_cfg": kv_cache_quant_cfg,
         "method": aq_config.auto_quantize_method,
+        "method_options": method_options,
         "score_size": aq_config.score_size,
     }
 
@@ -454,7 +459,7 @@ def forward_step(model, batch):
             inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch
             return model(**inputs_)
 
-    elif inputs["method"] == "kl_div":
+    elif inputs["method"] in ("kl_div", "aumann_shapley"):
 
         def forward_step(model, batch):
             inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch
@@ -464,9 +469,13 @@ def forward_step(model, batch):
                 return full_model.lm_head(output.last_hidden_state)
             return output.logits
 
+        if inputs["method"] == "aumann_shapley":
+            loss_func = None
+
     else:
         raise ValueError(
-            f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'"
+            f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient', 'kl_div', "
+            "or 'aumann_shapley'"
         )
 
     language_model, _ = mtq.auto_quantize(
@@ -483,6 +492,7 @@ def forward_step(model, batch):
         verbose=True,
         disabled_layers=inputs["disabled_layers"],
         method=inputs["method"],
+        method_options=inputs["method_options"],
         checkpoint=args.auto_quantize_checkpoint,
     )
 

modelopt/recipe/config.py

@@ -19,7 +19,7 @@
 
 import warnings
 from enum import Enum
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import Field, field_validator, model_validator
 
@@ -235,10 +235,17 @@ class AutoQuantizeConfig(ModeloptBaseConfig):
         description="Optional per-module overrides for candidate formats and BF16/no-quant "
         "selectability. Matching is performed after runtime-fusion grouping.",
     )
-    auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField(
+    auto_quantize_method: Literal["gradient", "kl_div", "aumann_shapley"] = ModeloptField(
         default="gradient",
         title="Sensitivity scoring method",
-        description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).",
+        description="'gradient' uses task labels; 'kl_div' and 'aumann_shapley' compare the "
+        "model's own outputs without labels.",
+    )
+    method_options: dict[str, Any] | None = ModeloptField(
+        default=None,
+        title="Method-specific scoring options",
+        description="Options forwarded to the selected scoring method. Each method validates "
+        "its supported keys and values.",
     )
     score_size: int = ModeloptField(
         default=128,
@@ -265,8 +272,24 @@ class AutoQuantizeConfig(ModeloptBaseConfig):
         "the --kv_cache_qformat CLI flag when omitted.",
     )
 
+    @property
+    def uses_predicted_damage_target(self) -> bool:
+        """Whether predicted damage, rather than effective bits, is the search target."""
+        return (self.method_options or {}).get("max_predicted_damage") is not None
+
     @model_validator(mode="after")
-    def _has_search_space(self):
+    def _validate_search_targets_and_space(self):
+        if self.uses_predicted_damage_target and self.auto_quantize_method != "aumann_shapley":
+            raise ValueError(
+                "method_options.max_predicted_damage requires "
+                "auto_quantize_method='aumann_shapley'."
+            )
+        has_explicit_bit_budget = "effective_bits" in self.constraints.model_fields_set
+        if self.uses_predicted_damage_target and has_explicit_bit_budget:
+            raise ValueError(
+                "A damage-bound AutoQuantize recipe must omit constraints.effective_bits; "
+                "max_predicted_damage supplies the search target."
+            )
         if not self.candidate_formats and not self.module_search_spaces:
             raise ValueError(
                 "auto_quantize requires candidate_formats or at least one module_search_spaces "