Add Aumann-Shapley sensitivity scoring method to auto_quantize

CHANGELOG.rst

@@ -8,6 +8,7 @@ Changelog
 
 *Quantization*
 
+- Add ``method="aumann_shapley"`` to ``mtq.auto_quantize``: label-free sensitivity scoring via Aumann-Shapley path-integral damage attributions with a measured-corner coverage calibration, so every allocation carries a ``predicted_damage`` quote in calibration-KL units. Method-specific settings (path nodes, damage link, a deterministic DP solver, and a ``max_predicted_damage`` bound mode) ride in a new optional ``auto_quantize(method_options=...)`` argument validated by the selected method. AutoQuantize scoring methods are now registered in ``modelopt.torch.quantization.algorithms.AUTO_QUANTIZE_SEARCHERS``.
 - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint.
 
 *Megatron Framework (M-LM / M-Bridge)*

examples/hf_ptq/README.md

@@ -382,7 +382,7 @@ scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/n
 The recipe quantizes the less accuracy-sensitive layers with the more aggressive format (e.g. NVFP4) and
 keeps the more sensitive ones at higher precision (or unquantized), so the model meets the recipe's
 `effective_bits` target. To author your own, copy a shipped recipe and adjust `candidate_formats`,
-`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div`), `score_size`,
+`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `kl_div` / `aumann_shapley`), `score_size`,
 `module_search_spaces` (optional per-module candidate overrides), `disabled_layers` (excluded from
 the search), and `cost_excluded_layers` (kept out of the bit-budget accounting — e.g. VL vision
 towers). Recipes can splice a shared base `disabled_layers` set via `$import` (see
@@ -450,7 +450,7 @@ The example scripts above also have an additional flag `--tasks`, where the actu
 
 > *If GPU out-of-memory error is reported running the scripts, please try editing the scripts and reducing the max batch size to save GPU memory.*
 
-> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` method. The `kl_div` method does not require backpropagation.*
+> *NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the `gradient` or `aumann_shapley` methods (the latter is label-free but still backpropagates a KL loss). The `kl_div` method does not require backpropagation.*
 
 ## Real Quant
 

examples/hf_ptq/hf_ptq.py

@@ -495,7 +495,7 @@ def forward_step(model, batch):
             inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch
             return model(**inputs_)
 
-    elif inputs["method"] == "kl_div":
+    elif inputs["method"] in ("kl_div", "aumann_shapley"):
 
         def forward_step(model, batch):
             inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch
@@ -507,7 +507,8 @@ def forward_step(model, batch):
 
     else:
         raise ValueError(
-            f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient' or 'kl_div'"
+            f"Invalid auto_quantize method: {inputs['method']}. Must be 'gradient', 'kl_div', "
+            "or 'aumann_shapley'"
         )
 
     language_model, _ = mtq.auto_quantize(
@@ -1574,7 +1575,7 @@ def parse_args() -> argparse.Namespace:
         "--auto_quantize_method",
         type=str,
         default="gradient",
-        choices=["gradient", "kl_div"],
+        choices=["gradient", "kl_div", "aumann_shapley"],
         help="[Deprecated: use an AutoQuantize --recipe] Sensitivity scoring method.",
     )
     parser.add_argument(

modelopt/recipe/config.py

@@ -254,10 +254,12 @@ class AutoQuantizeConfig(ModeloptBaseConfig):
         description="Optional per-module overrides for candidate formats and BF16/no-quant "
         "selectability. Matching is performed after runtime-fusion grouping.",
     )
-    auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField(
+    auto_quantize_method: Literal["gradient", "kl_div", "aumann_shapley"] = ModeloptField(
         default="gradient",
         title="Sensitivity scoring method",
-        description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).",
+        description="'gradient' (Taylor + Fisher, needs labels), 'kl_div' (no labels), or "
+        "'aumann_shapley' (no labels; path-integral damage attributions with a predicted-damage "
+        "quote).",
     )
     score_size: int = ModeloptField(
         default=128,