Add normalized and group-boundary AutoQuant scoring

CHANGELOG.rst

@@ -19,6 +19,7 @@ Changelog
 
 **New Features**
 
+- Add normalized and group-boundary AutoQuant scoring. Set ``auto_quantize.constraints.score_model: per_element`` in an AutoQuantize recipe to normalize selector coefficients by represented weight elements, and set ``auto_quantize_method: group_recon`` with ``score_boundary: group`` to score projection recipes by normalized reconstruction error at their attention or MLP group output. ``quant_grouping_scheme`` independently controls whether self-attention q/k/v/o or linear-attention qkv/z/out share one recipe decision beyond mandatory runtime-fused groups. Shared-expert gate/up/down projections are grouped as one fused-MoE decision. The deprecated HF PTQ CLI flags map to the same recipe fields.
 - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 <https://arxiv.org/abs/2605.18810>`_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change).
 - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred.
 - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``.

examples/hf_ptq/README.md

@@ -352,18 +352,23 @@ Here is an example usage for `AutoQuantize` algorithm (Please see [auto_quantize
 `AutoQuantize` can be performed for Huggingface LLM models like [Qwen](https://huggingface.co/Qwen/Qwen3-8B) / [Nemotron](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16) as shown below:
 
 `AutoQuantize` is driven by an **AutoQuantize recipe** passed with `--recipe`. The recipe defines the
-candidate formats, the `effective_bits` target, cost model, scoring method, search-disabled layers, and
-cost-excluded layers — see [`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in
+candidate formats, the `effective_bits` target, cost model, scoring objective and boundary,
+search-disabled layers, and cost-excluded layers — see
+[`AutoQuantizeConfig`](../../modelopt/recipe/config.py). Shipped recipes live in
 [`modelopt_recipes/general/auto_quantize/`](../../modelopt_recipes/general/auto_quantize); model-specific
 recipes (carrying architecture-specific disabled layers — e.g. VL vision towers) live under
 `modelopt_recipes/huggingface/<model>/auto_quantize/`.
 
 > *Migration: prefer an AutoQuantize `--recipe`. The `--auto_quantize_bits`, `--auto_quantize_method`,
-> `--auto_quantize_score_size`, `--auto_quantize_cost_model`, and `--auto_quantize_active_moe_expert_ratio`
+> `--auto_quantize_score_size`, `--auto_quantize_score_model`,
+> `--auto_quantize_score_boundary`, `--auto_quantize_cost_model`, and
+> `--auto_quantize_active_moe_expert_ratio`
 > CLI flags are **deprecated but still work** — they are converted into an `AutoQuantizeConfig` on the fly
 > (with a `DeprecationWarning`) and will be removed in a future release. They map to recipe fields:
 > `--auto_quantize_bits` → `constraints.effective_bits`, `--auto_quantize_method` → `auto_quantize_method`,
-> `--auto_quantize_score_size` → `score_size`, `--auto_quantize_cost_model` → `constraints.cost_model`,
+> `--auto_quantize_score_size` → `score_size`, `--auto_quantize_score_model` →
+> `constraints.score_model`, `--auto_quantize_score_boundary` → `score_boundary`,
+> `--auto_quantize_cost_model` → `constraints.cost_model`,
 > `--auto_quantize_active_moe_expert_ratio` → `constraints.cost.active_moe_expert_ratio`, and the
 > `--qformat fp8,nvfp4` candidate list → `candidate_formats`. When converted, the shared base
 > `disabled_layers` and `cost_excluded_layers` patterns are appended automatically. `--auto_quantize_checkpoint`
@@ -381,9 +386,16 @@ 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`,
-`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
+`constraints.effective_bits`, `auto_quantize_method` (`gradient` / `group_recon` / `kl_div`),
+`constraints.score_model` (`raw` / `per_element`), `score_boundary` (`local` / `group`),
+`quant_grouping_scheme`, `score_size`, `disabled_layers` (excluded from the search), and
+`cost_excluded_layers` (kept out
+of the bit-budget accounting — e.g. VL vision towers). `group_recon` measures normalized
+reconstruction error at shared attention/MLP outputs; group scoring changes the measurement
+boundary but does not force those projections to share one recipe decision. Set
+`quant_grouping_scheme` to an attention-layer variant when the search should make one recipe
+decision for self-attention q/k/v/o and/or linear-attention qkv/z/out. Recipes can splice a
+shared base `disabled_layers` set via
 `$import` (see `modelopt_recipes/configs/auto_quantize/units/base_disabled_layers`).
 
 bf16 (no quantization) is always an implicit per-layer choice, so `candidate_formats` need only list
@@ -407,7 +419,8 @@ 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 when using the `gradient` method. Models without
+> backpropagation support can use `group_recon` or `kl_div`, which are forward-only.*
 
 ## Real Quant
 

examples/hf_ptq/hf_ptq.py

@@ -347,6 +347,8 @@ def _mtq_inputs_from_auto_quantize_config(aq_config, args: argparse.Namespace) -
         "disabled_layers": aq_config.disabled_layers,
         "kv_cache_quant_cfg": kv_cache_quant_cfg,
         "method": aq_config.auto_quantize_method,
+        "quant_grouping_scheme": aq_config.quant_grouping_scheme,
+        "score_boundary": aq_config.score_boundary,
         "score_size": aq_config.score_size,
     }
 
@@ -383,9 +385,12 @@ def _auto_quantize_config_from_cli(args: argparse.Namespace):
             effective_bits=args.auto_quantize_bits,
             cost_model=args.auto_quantize_cost_model,
             cost=cost,
+            score_model=args.auto_quantize_score_model,
         ),
         candidate_formats=[QuantizeConfig(**QUANT_CFG_CHOICES[q]) for q in args.qformat.split(",")],
         auto_quantize_method=args.auto_quantize_method,
+        quant_grouping_scheme=args.auto_quantize_grouping_scheme,
+        score_boundary=args.auto_quantize_score_boundary,
         score_size=args.auto_quantize_score_size,
         disabled_layers=disabled_layers,
         cost_excluded_layers=cost_excluded_layers,
@@ -439,7 +444,7 @@ def loss_func(output, data):
         def loss_func(output, data):
             return output.loss
 
-    if inputs["method"] == "gradient":
+    if inputs["method"] in {"gradient", "group_recon"}:
 
         def forward_step(model, batch):
             inputs_ = {k: v for k, v in batch.items() if k != "labels"} if is_base_model else batch
@@ -457,21 +462,24 @@ 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', 'group_recon', or 'kl_div'"
         )
 
     language_model, _ = mtq.auto_quantize(
         language_model,
         constraints=inputs["constraints"],
         data_loader=calib_dataloader,
         forward_step=forward_step,
-        loss_func=loss_func,
+        loss_func=loss_func if inputs["method"] == "gradient" else None,
         quantization_formats=inputs["quantization_formats"],
         num_calib_steps=len(calib_dataloader),
         num_score_steps=min(len(calib_dataloader), max(inputs["score_size"] // args.batch_size, 1)),
         verbose=True,
         disabled_layers=inputs["disabled_layers"],
         method=inputs["method"],
+        quant_grouping_scheme=inputs["quant_grouping_scheme"],
+        score_boundary=inputs["score_boundary"],
         checkpoint=args.auto_quantize_checkpoint,
     )
 
@@ -1476,7 +1484,7 @@ def parse_args() -> argparse.Namespace:
         "--auto_quantize_method",
         type=str,
         default="gradient",
-        choices=["gradient", "kl_div"],
+        choices=["gradient", "group_recon", "kl_div"],
         help="[Deprecated: use an AutoQuantize --recipe] Sensitivity scoring method.",
     )
     parser.add_argument(
@@ -1485,6 +1493,32 @@ def parse_args() -> argparse.Namespace:
         default=128,
         help="[Deprecated: use an AutoQuantize --recipe] Number of samples for sensitivity scoring.",
     )
+    parser.add_argument(
+        "--auto_quantize_score_model",
+        type=str,
+        default="raw",
+        choices=["raw", "per_element"],
+        help="[Deprecated: use an AutoQuantize --recipe] Selector score model.",
+    )
+    parser.add_argument(
+        "--auto_quantize_score_boundary",
+        type=str,
+        default=None,
+        choices=["local", "group"],
+        help="[Deprecated: use an AutoQuantize --recipe] Sensitivity score boundary.",
+    )
+    parser.add_argument(
+        "--auto_quantize_grouping_scheme",
+        type=str,
+        default="runtime_fused",
+        choices=[
+            "runtime_fused",
+            "runtime_fused+linear_attn_layer",
+            "runtime_fused+self_attn_layer",
+            "runtime_fused+linear_attn_layer+self_attn_layer",
+        ],
+        help="[Deprecated: use an AutoQuantize --recipe] Quantization decision grouping.",
+    )
     parser.add_argument(
         "--auto_quantize_cost_model",
         type=str,

examples/hf_ptq/scripts/huggingface_example.sh

@@ -111,6 +111,11 @@ if [ -n "$AUTO_QUANTIZE_BITS" ]; then
     PTQ_ARGS+=" --auto_quantize_bits=$AUTO_QUANTIZE_BITS "
     PTQ_ARGS+=" --auto_quantize_method=${AUTO_QUANTIZE_METHOD:-gradient} "
     PTQ_ARGS+=" --auto_quantize_score_size=${AUTO_QUANTIZE_SCORE_SIZE:-128} "
+    PTQ_ARGS+=" --auto_quantize_score_model=${AUTO_QUANTIZE_SCORE_MODEL:-raw} "
+    if [ -n "$AUTO_QUANTIZE_SCORE_BOUNDARY" ]; then
+        PTQ_ARGS+=" --auto_quantize_score_boundary=$AUTO_QUANTIZE_SCORE_BOUNDARY "
+    fi
+    PTQ_ARGS+=" --auto_quantize_grouping_scheme=${AUTO_QUANTIZE_GROUPING_SCHEME:-runtime_fused} "
     PTQ_ARGS+=" --auto_quantize_cost_model=${AUTO_QUANTIZE_COST_MODEL:-weight} "
     if [ -n "$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO" ]; then
         PTQ_ARGS+=" --auto_quantize_active_moe_expert_ratio=$AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO "

examples/hf_ptq/scripts/parser.sh

@@ -41,7 +41,7 @@ parse_options() {
     CALIB_WITH_IMAGES=false
 
   # Parse command-line options
-  ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@")
+  ARGS=$(getopt -o "" -l "model:,quant:,recipe:,kv_cache_quant:,tp:,pp:,sparsity:,awq_block_size:,calib:,calib_batch_size:,output:,batch:,tasks:,lm_eval_tasks:,lm_eval_limit:,simple_eval_tasks:,simple_eval_limit:,mmlu_limit:,trust_remote_code,use_seq_device_map,gpu_max_mem_percentage:,kv_cache_free_gpu_memory_fraction:,low_memory_mode,no-verbose,calib_dataset:,calib_seq:,auto_quantize_checkpoint:,auto_quantize_bits:,auto_quantize_method:,auto_quantize_score_size:,auto_quantize_score_model:,auto_quantize_score_boundary:,auto_quantize_grouping_scheme:,auto_quantize_cost_model:,auto_quantize_active_moe_expert_ratio:,moe_calib_experts_ratio:,cast_mxfp4_to_nvfp4,vlm,calib_with_images" -n "$0" -- "$@")
 
   eval set -- "$ARGS"
   while true; do
@@ -76,6 +76,9 @@ parse_options() {
       --auto_quantize_bits ) AUTO_QUANTIZE_BITS="$2"; shift 2;;
       --auto_quantize_method ) AUTO_QUANTIZE_METHOD="$2"; shift 2;;
       --auto_quantize_score_size ) AUTO_QUANTIZE_SCORE_SIZE="$2"; shift 2;;
+      --auto_quantize_score_model ) AUTO_QUANTIZE_SCORE_MODEL="$2"; shift 2;;
+      --auto_quantize_score_boundary ) AUTO_QUANTIZE_SCORE_BOUNDARY="$2"; shift 2;;
+      --auto_quantize_grouping_scheme ) AUTO_QUANTIZE_GROUPING_SCHEME="$2"; shift 2;;
       --auto_quantize_cost_model ) AUTO_QUANTIZE_COST_MODEL="$2"; shift 2;;
       --auto_quantize_active_moe_expert_ratio ) AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO="$2"; shift 2;;
       --moe_calib_experts_ratio ) MOE_CALIB_EXPERTS_RATIO="$2"; shift 2;;
@@ -180,6 +183,9 @@ parse_options() {
   echo "auto_quantize_bits: $AUTO_QUANTIZE_BITS"
   echo "auto_quantize_method: $AUTO_QUANTIZE_METHOD"
   echo "auto_quantize_score_size: $AUTO_QUANTIZE_SCORE_SIZE"
+  echo "auto_quantize_score_model: $AUTO_QUANTIZE_SCORE_MODEL"
+  echo "auto_quantize_score_boundary: $AUTO_QUANTIZE_SCORE_BOUNDARY"
+  echo "auto_quantize_grouping_scheme: $AUTO_QUANTIZE_GROUPING_SCHEME"
   echo "auto_quantize_cost_model: $AUTO_QUANTIZE_COST_MODEL"
   echo "auto_quantize_active_moe_expert_ratio: $AUTO_QUANTIZE_ACTIVE_MOE_EXPERT_RATIO"
   echo "moe_calib_experts_ratio: $MOE_CALIB_EXPERTS_RATIO"

modelopt/recipe/config.py

@@ -182,6 +182,12 @@ class AutoQuantizeConstraints(ModeloptBaseConfig):
         title="Cost-model parameters",
         description="Extra cost-model parameters; omit for the 'weight' cost model.",
     )
+    score_model: Literal["raw", "per_element"] = ModeloptField(
+        default="raw",
+        title="Selector score model",
+        description="'raw' uses sensitivity scores directly; 'per_element' normalizes each "
+        "score by the represented weight-element cost before budgeted selection.",
+    )
 
     @field_validator("effective_bits")
     @classmethod
@@ -206,10 +212,30 @@ class AutoQuantizeConfig(ModeloptBaseConfig):
         "(e.g. [fp8]) yields a {fp8, bf16} per-layer search.",
         validate_default=True,
     )
-    auto_quantize_method: Literal["gradient", "kl_div"] = ModeloptField(
+    auto_quantize_method: Literal["gradient", "group_recon", "kl_div"] = ModeloptField(
         default="gradient",
         title="Sensitivity scoring method",
-        description="'gradient' (Taylor + Fisher, needs labels) or 'kl_div' (no labels).",
+        description="'gradient' (Taylor + Fisher, needs labels), 'group_recon' (normalized "
+        "group-output reconstruction, no labels), or 'kl_div' (no labels).",
+    )
+    score_boundary: Literal["local", "group"] | None = ModeloptField(
+        default=None,
+        title="Sensitivity score boundary",
+        description="'local' scores each quantized module output; 'group' scores attention and "
+        "MoE projection perturbations at their shared attention/MLP output. Defaults to 'group' "
+        "for group_recon and 'local' otherwise.",
+    )
+    quant_grouping_scheme: Literal[
+        "runtime_fused",
+        "runtime_fused+linear_attn_layer",
+        "runtime_fused+self_attn_layer",
+        "runtime_fused+linear_attn_layer+self_attn_layer",
+    ] = ModeloptField(
+        default="runtime_fused",
+        title="Quantization decision grouping",
+        description="Search-decision grouping beyond mandatory runtime-fused groups. Attention-"
+        "layer options make all listed projections in one layer share a quantization recipe; "
+        "this is independent of score_boundary.",
     )
     score_size: int = ModeloptField(
         default=128,
@@ -248,6 +274,22 @@ def _at_least_one_candidate(cls, v: list[QuantizeConfig]) -> list[QuantizeConfig
             )
         return v
 
+    @model_validator(mode="after")
+    def _validate_scoring_configuration(self):
+        boundary = self.score_boundary or (
+            "group" if self.auto_quantize_method == "group_recon" else "local"
+        )
+        if self.auto_quantize_method == "group_recon" and boundary != "group":
+            raise ValueError("auto_quantize_method='group_recon' requires score_boundary='group'.")
+        if self.auto_quantize_method == "kl_div" and (
+            self.constraints.score_model != "raw" or boundary != "local"
+        ):
+            raise ValueError(
+                "auto_quantize_method='kl_div' requires constraints.score_model='raw' and "
+                "score_boundary='local'."
+            )
+        return self
+
 
 class ModelOptAutoQuantizeRecipe(ModelOptRecipeBase):
     """Our config class for AutoQuantize recipes."""

modelopt/torch/quantization/_auto_quantize_cost.py

@@ -26,7 +26,14 @@
 # constraint is supplied. The value is intentionally kept for backward compatibility.
 DEFAULT_AUTO_QUANTIZE_EFFECTIVE_BITS: Final = 4.8
 
-AUTO_QUANTIZE_CONSTRAINT_KEYS: Final = frozenset({"effective_bits", "cost_model", "cost"})
+AUTO_QUANTIZE_CONSTRAINT_KEYS: Final = frozenset(
+    {"effective_bits", "cost_model", "cost", "score_model"}
+)
+AUTO_QUANTIZE_SCORE_MODEL_RAW: Final = "raw"
+AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT: Final = "per_element"
+AUTO_QUANTIZE_SCORE_MODELS: Final = frozenset(
+    {AUTO_QUANTIZE_SCORE_MODEL_RAW, AUTO_QUANTIZE_SCORE_MODEL_PER_ELEMENT}
+)
 ACTIVE_MOE_EXPERT_RATIO_KEY: Final = "active_moe_expert_ratio"
 EXCLUDED_MODULE_NAME_PATTERNS_KEY: Final = "excluded_module_name_patterns"
 COST_MODEL_WEIGHT: Final = "weight"
@@ -235,8 +242,15 @@ def normalize_auto_quantize_constraints(
     if unexpected_constraint_keys:
         raise ValueError(
             f"Unsupported auto_quantize constraints: {unexpected_constraint_keys}. "
-            "Supported constraints are 'effective_bits', 'cost_model', and 'cost'."
+            "Supported constraints are 'effective_bits', 'cost_model', 'cost', and 'score_model'."
+        )
+
+    score_model = constraints.get("score_model", AUTO_QUANTIZE_SCORE_MODEL_RAW)
+    if score_model not in AUTO_QUANTIZE_SCORE_MODELS:
+        raise ValueError(
+            f"constraints['score_model'] must be one of {sorted(AUTO_QUANTIZE_SCORE_MODELS)}."
         )
+    constraints["score_model"] = score_model
 
     cost_model_name = constraints.get("cost_model", COST_MODEL_WEIGHT)
     if not isinstance(cost_model_name, str):