feat: Aumann-Shapley sensitivity scoring method for auto_quantize

Adds method='aumann_shapley': label-free sensitivity scoring via Aumann-Shapley
path-integral damage attributions (KL divergence against the model's own
unquantized outputs, or the fixed_quantization_config baseline when supplied),
with a measured-corner coverage calibration so every allocation carries a
predicted_damage quote in calibration units with recorded validity, anchored
to reproduce the measured corner. Scores all candidate formats in one
reference forward, one corner forward, and one fwd+bwd per (format, path
node) per batch using the same local-replay mechanism as the gradient method;
a KL loss requires path integration because its gradient is exactly zero at
the unquantized point. This is an efficient implementation of the estimator in
https://arxiv.org/abs/2607.12266, validated empirically against it;
implementation details are documented in the module docstring.

Method-specific settings ride in a new optional auto_quantize(method_options=)
dict, validated against each searcher's declared method_options_keys so core
inputs cannot be overridden: num_path_nodes, damage_link, a deterministic
grid-approximate DP solver alternative to the LP, and max_predicted_damage
(minimize weight cost subject to predicted damage <= bound, conservatively
rounded and mutually exclusive with an effective_bits constraint). Internal
format tables are keyed by QuantRecipe.checkpoint_signature so identical
custom formats under different auto-generated names resolve to one format; a
scoring signature in the search state rejects checkpoint resumes that would
change what stored scores mean while allowing solver-only re-solves.

The hardcoded method dispatch becomes a registry (AUTO_QUANTIZE_SEARCHERS) so
methods register themselves; gradient/kl_div behavior is unchanged (existing
suite passes as-is). Vocab-sharded (Megatron-TP) losses raise
NotImplementedError pending an autograd-correct vocab-parallel log-softmax.

Tests: method parametrizations extended in test_autoquant.py (21 new cases);
test_autoquant_shapley.py pins config parity with the standard builder
(dict-for-dict), the path-integral completeness diagnostic, corner anchoring
under incomplete attributions, damage-bound certification, solver optimality
contracts against brute force, custom-format identity, heterogeneous-ladder
flagging, exact-zero and tiny-attribution inversion behavior,
scoring-signature resume guards, and method-option validation.

Signed-off-by: Joshua Hill <joshua.hill@baseten.co>

Author: joshua-hill

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,