Add active-MoE AutoQuant cost accounting (#1497)

### What does this PR do?

• Type of change: new feature

Adds an active_moe cost model for auto_quantize effective-bits search.
This lets AutoQuant account for routed MoE expert weights by active
decode weight traffic instead of total checkpoint weight
size, using active_moe_expert_ratio = num_experts_per_tok / num_experts.

The default behavior is unchanged: cost_model="weight" still counts all
quantizable weights equally.

  ### Usage

  import modelopt.torch.quantization as mtq

  model, search_state = mtq.auto_quantize(
      model,
      constraints={"effective_bits": 5.0},
      quantization_formats=[
          mtq.NVFP4_DEFAULT_CFG,
          mtq.FP8_DEFAULT_CFG,
      ],
      data_loader=calib_dataloader,
      forward_step=forward_step,
      loss_func=loss_func,
      cost_model="active_moe",
# Optional. If omitted, ModelOpt tries to infer this from model.config.
      active_moe_expert_ratio=2 / 64,
  )

  The HF PTQ example also exposes:

  --auto_quantize_cost_model active_moe \
  --auto_quantize_active_moe_expert_ratio 0.03125

  ### Testing

python -m pytest tests/unit/torch/quantization/test_autoquant.py -q -k
'active_moe or quant_recipe_hparam_cost_weight'
python -m pytest tests/unit/torch/quantization/test_autoquant.py -q -k
'not data_parallel_auto_quantize'

  Results:

  - 4 passed
  - 58 passed, 1 deselected

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ / ❌ / N/A <!--- If ❌, explain
why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A
<!--- Mandatory -->
- Did you write any new necessary tests?: ✅ / ❌ / N/A <!--- Mandatory
for new features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ / ❌ / N/A <!--- Only for new features, API changes, critical bug fixes
or backward incompatible changes. -->
- Did you get Claude approval on this PR?: ✅ / ❌ / N/A <!--- Run
`/claude review`. NVIDIA org members can self-trigger for complex
changes; orthogonal to CodeRabbit. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added active-MoE cost model option for auto-quantization with
configurable expert ratio; API and CLI accept cost_model and
active_moe_expert_ratio
  * Unified auto-quantize supports new quant format w4a16_nvfp4

* **Bug Fixes**
* Ensure labels are moved to the logits device for base models without
an lm_head
* CLI enforces valid expert-ratio range and requires active-MoE mode
when a ratio is provided

* **Tests**
* Added unit tests for active-MoE behavior, cost-weighting, ratio
handling, and search budget selection

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/NVIDIA/Model-Optimizer/pull/1497?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>

.github/workflows/gpu_tests.yml

@@ -42,7 +42,7 @@ jobs:
             timeout: 75
             container_image: nvcr.io/nvidia/pytorch:26.04-py3
           - example: gpu_megatron
-            timeout: 45
+            timeout: 60
             container_image: nvcr.io/nvidia/nemo:26.04
           - example: gpu_trtllm
             timeout: 30

CHANGELOG.rst

@@ -33,6 +33,7 @@ Changelog
 - DeepSeek PTQ (``examples/deepseek/ptq.py``) now defaults to native top-k calibration with post-hoc per-layer peer-max sync of expert ``input_quantizer.amax``; the all-experts path is preserved behind ``--calib_all_experts``.
 - Add NVFP4 W4A16 weight-only quantization (``w4a16_nvfp4``): FP4 weights with group_size=16, BF16 activations, no calibration forward pass required. Use ``mtq.W4A16_NVFP4_CFG`` or ``--qformat w4a16_nvfp4`` in ``hf_ptq.py``. vLLM deployment support is in progress.
 - Add Megatron Core export/import mapping for Qwen3-VL (``Qwen3VLForConditionalGeneration``) vision-language models. The mapping handles the ``model.language_model.`` weight prefix used by Qwen3-VL.
+- Add active-MoE cost accounting for ``mtq.auto_quantize`` effective-bits search. Set ``constraints={"effective_bits": ..., "cost_model": "active_moe", "cost": {"active_moe_expert_ratio": ...}}`` to weight routed MoE expert costs by active experts per token while keeping shared experts fully counted. The ``hf_ptq.py`` AutoQuant path exposes this via ``--auto_quantize_cost_model active_moe`` and ``--auto_quantize_active_moe_expert_ratio``.
 - Add ``DATASET_COMBOS`` to ``modelopt.torch.utils.dataset_utils`` — single ``--dataset`` tokens that fan out to multiple registered datasets; per-entry ``num_samples`` is split evenly across the members. Initial combos: ``cnn_nemotron_v2_mix`` (``cnn_dailymail`` + ``nemotron-post-training-dataset-v2``, used by ``hf_ptq.py`` when no ``--dataset`` is provided) and ``nemotron-post-training-v3`` (the seven ``nvidia/Nemotron-*`` SFT datasets added in #1498, mirroring the `nemotron-post-training-v3 collection <https://huggingface.co/collections/nvidia/nemotron-post-training-v3>`_). Combo names are listed by ``get_supported_datasets()`` and surfaced in ``--dataset`` help. ``get_dataset_dataloader`` rejects inputs that mix a combo with one of its member datasets (e.g. ``cnn_dailymail,cnn_nemotron_v2_mix``) to avoid double-sampling, and ``get_dataset_samples`` rejects combo names so callers route through the dataloader. ``hf_ptq.py`` default ``--calib_size`` is bumped from ``512`` to ``1024`` so the total calibration sample count under the new default combo matches the previous two-dataset fallback.
 - The ``nemotron-sft-agentic-v2`` registered dataset (added in #1498) now uses only the ``search`` split. The previously configured ``interactive_agent`` and ``tool_calling`` splits contain content-level defects (heterogeneous schema and a malformed JSON row, respectively) that cause pyarrow's streaming JSON reader to fail deterministically.
 - Add shared Megatron-Core calibration forward loop: ``modelopt.torch.utils.plugins.megatron_calibration.get_megatron_calibration_forward_loop`` produces the ``forward_loop`` callable expected by ``mtq.quantize`` / ``mtp.prune``. Replaces the bespoke calibration loops in Megatron-LM and Megatron-Bridge for quantization and pruning with a single canonical implementation.

examples/llm_ptq/hf_ptq.py

@@ -382,9 +382,18 @@ def auto_quantize(
             f"Invalid auto_quantize_method: {auto_quantize_method}. Must be 'gradient' or 'kl_div'"
         )
 
+    auto_quantize_constraints = {
+        "effective_bits": args.auto_quantize_bits,
+        "cost_model": args.auto_quantize_cost_model,
+    }
+    if args.auto_quantize_active_moe_expert_ratio is not None:
+        auto_quantize_constraints["cost"] = {
+            "active_moe_expert_ratio": args.auto_quantize_active_moe_expert_ratio
+        }
+
     language_model, _ = mtq.auto_quantize(
         language_model,
-        constraints={"effective_bits": args.auto_quantize_bits},
+        constraints=auto_quantize_constraints,
         data_loader=calib_dataloader,
         forward_step=forward_step,
         loss_func=loss_func,  # Only used for gradient-based method
@@ -1391,6 +1400,29 @@ def parse_args() -> argparse.Namespace:
             "(sensitivity scores, costs, etc.). Only used when auto_quantize_bits is specified."
         ),
     )
+    parser.add_argument(
+        "--auto_quantize_cost_model",
+        type=str,
+        default="weight",
+        choices=["weight", "active_moe"],
+        help=(
+            "Cost model for auto_quantize effective-bits accounting. 'weight' counts all "
+            "quantizable weights equally. 'active_moe' scales routed MoE expert weights by "
+            "--auto_quantize_active_moe_expert_ratio, or infers top_k/num_experts from model config."
+        ),
+    )
+    parser.add_argument(
+        "--auto_quantize_active_moe_expert_ratio",
+        type=float,
+        default=None,
+        help=(
+            "Routed MoE expert active ratio for --auto_quantize_cost_model active_moe. "
+            "For top-k MoE this is top_k / num_experts. If omitted, common model config "
+            "fields such as num_experts_per_tok and num_experts are used when available. "
+            "This only affects AutoQuant cost accounting and does not change calibration "
+            "routing; use --moe_calib_experts_ratio to control calibration expert coverage."
+        ),
+    )
     parser.add_argument(
         "--moe_calib_experts_ratio",
         type=float,
@@ -1424,6 +1456,18 @@ def parse_args() -> argparse.Namespace:
     args = parser.parse_args()
     if args.moe_calib_experts_ratio is not None and not (0.0 < args.moe_calib_experts_ratio <= 1.0):
         parser.error("--moe_calib_experts_ratio must be in the range (0.0, 1.0].")
+    if args.auto_quantize_active_moe_expert_ratio is not None and not (
+        0.0 < args.auto_quantize_active_moe_expert_ratio <= 1.0
+    ):
+        parser.error("--auto_quantize_active_moe_expert_ratio must be in the range (0.0, 1.0].")
+    if (
+        args.auto_quantize_cost_model == "weight"
+        and args.auto_quantize_active_moe_expert_ratio is not None
+    ):
+        parser.error(
+            "--auto_quantize_active_moe_expert_ratio requires "
+            "--auto_quantize_cost_model active_moe."
+        )
 
     if args.specdec_offline_dataset is not None and args.sparsity_fmt != "dense":
         parser.error("--specdec_offline_dataset is only supported with --sparsity_fmt dense (PTQ).")

modelopt/torch/quantization/_auto_quantize_cost.py

@@ -0,0 +1,221 @@
+# 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.
+
+"""Cost models for AutoQuantize effective-bits accounting."""
+
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any, Final
+
+import regex as re
+import torch.nn as nn
+
+# Default target used by historical AutoQuantize calls when no explicit effective-bits
+# 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"})
+ACTIVE_MOE_EXPERT_RATIO_KEY: Final = "active_moe_expert_ratio"
+COST_MODEL_WEIGHT: Final = "weight"
+COST_MODEL_ACTIVE_MOE: Final = "active_moe"
+
+_ROUTED_MOE_EXPERT_NAME_RE = re.compile(r"(^|\.)experts(\.|$)")
+_ACTIVE_MOE_TOP_K_ATTRS = (
+    "num_experts_per_tok",
+    "num_experts_per_token",
+    "moe_top_k",
+    "top_k",
+    "num_selected_experts",
+)
+_ACTIVE_MOE_NUM_EXPERTS_ATTRS = (
+    "num_experts",
+    "num_local_experts",
+    "n_routed_experts",
+    "moe_num_experts",
+    "num_routed_experts",
+)
+
+
+def _iter_model_configs(model: nn.Module):
+    seen = set()
+    for obj in (model, getattr(model, "model", None), getattr(model, "language_model", None)):
+        config = getattr(obj, "config", None)
+        if config is None or id(config) in seen:
+            continue
+        seen.add(id(config))
+        yield config
+        for nested_attr in ("text_config", "language_config"):
+            nested_config = getattr(config, nested_attr, None)
+            if nested_config is None or id(nested_config) in seen:
+                continue
+            seen.add(id(nested_config))
+            yield nested_config
+
+
+def _get_first_numeric_config_attr(config: Any, attr_names: tuple[str, ...]) -> float | None:
+    for attr_name in attr_names:
+        value = getattr(config, attr_name, None)
+        if isinstance(value, (int, float)) and not isinstance(value, bool):
+            return float(value)
+    return None
+
+
+def infer_active_moe_expert_ratio(model: nn.Module) -> float | None:
+    """Infer top-k / num-experts from a single model config object when possible."""
+    for config in _iter_model_configs(model):
+        num_active_experts = _get_first_numeric_config_attr(config, _ACTIVE_MOE_TOP_K_ATTRS)
+        num_experts = _get_first_numeric_config_attr(config, _ACTIVE_MOE_NUM_EXPERTS_ATTRS)
+        if num_active_experts is None or num_experts is None or num_experts <= 0:
+            continue
+        ratio = num_active_experts / num_experts
+        if ratio <= 0.0:
+            continue
+        return min(ratio, 1.0)
+    return None
+
+
+def is_routed_moe_module_name(name: str) -> bool:
+    """Return True for routed MoE expert modules, excluding shared experts."""
+    return "shared_expert" not in name and _ROUTED_MOE_EXPERT_NAME_RE.search(name) is not None
+
+
+class AutoQuantizeCostModel:
+    """Base class for AutoQuantize effective-bits cost accounting."""
+
+    name: str
+    supported_cost_keys: frozenset[str] = frozenset()
+
+    def normalize_cost_constraints(
+        self, model: nn.Module, cost_constraints: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Validate and normalize cost-model-specific constraints."""
+        unknown_cost_keys = set(cost_constraints) - self.supported_cost_keys
+        if unknown_cost_keys:
+            raise ValueError(f"Unsupported auto_quantize cost constraints: {unknown_cost_keys}.")
+        return cost_constraints
+
+    def module_cost_weight(
+        self, module_names: Sequence[str], cost_constraints: dict[str, Any]
+    ) -> float:
+        """Return the cost multiplier for a group of modules."""
+        return 1.0
+
+    def total_weight_size(
+        self,
+        named_modules: Iterable[tuple[str, nn.Module]],
+        is_auto_quantize_module: Callable[[nn.Module], bool],
+        cost_constraints: dict[str, Any],
+    ) -> float:
+        """Return the cost denominator for the effective-bits constraint."""
+        return sum(
+            module.weight.numel() * self.module_cost_weight([name], cost_constraints)
+            for name, module in named_modules
+            if is_auto_quantize_module(module)
+        )
+
+
+class WeightCostModel(AutoQuantizeCostModel):
+    """Count all quantizable weights equally."""
+
+    name = COST_MODEL_WEIGHT
+
+
+class ActiveMoECostModel(AutoQuantizeCostModel):
+    """Scale routed MoE expert weights by the active experts per-token ratio."""
+
+    name = COST_MODEL_ACTIVE_MOE
+    supported_cost_keys = frozenset({ACTIVE_MOE_EXPERT_RATIO_KEY})
+
+    def normalize_cost_constraints(
+        self, model: nn.Module, cost_constraints: dict[str, Any]
+    ) -> dict[str, Any]:
+        cost_constraints = super().normalize_cost_constraints(model, cost_constraints)
+        active_moe_expert_ratio = cost_constraints.get(ACTIVE_MOE_EXPERT_RATIO_KEY)
+        if active_moe_expert_ratio is None:
+            active_moe_expert_ratio = infer_active_moe_expert_ratio(model)
+            if active_moe_expert_ratio is None:
+                raise ValueError(
+                    "Could not infer active_moe_expert_ratio from model.config. "
+                    "Pass it via constraints['cost']['active_moe_expert_ratio']."
+                )
+
+        if not (
+            isinstance(active_moe_expert_ratio, (int, float))
+            and not isinstance(active_moe_expert_ratio, bool)
+            and 0.0 < active_moe_expert_ratio <= 1.0
+        ):
+            raise ValueError(
+                "constraints['cost']['active_moe_expert_ratio'] must be in (0.0, 1.0]."
+            )
+        cost_constraints[ACTIVE_MOE_EXPERT_RATIO_KEY] = float(active_moe_expert_ratio)
+        return cost_constraints
+
+    def module_cost_weight(
+        self, module_names: Sequence[str], cost_constraints: dict[str, Any]
+    ) -> float:
+        if any(is_routed_moe_module_name(n) for n in module_names):
+            return cost_constraints[ACTIVE_MOE_EXPERT_RATIO_KEY]
+        return 1.0
+
+
+_COST_MODELS: Final = {
+    COST_MODEL_WEIGHT: WeightCostModel(),
+    COST_MODEL_ACTIVE_MOE: ActiveMoECostModel(),
+}
+
+
+def get_auto_quantize_cost_model(name: str) -> AutoQuantizeCostModel:
+    """Return the registered AutoQuantize cost model."""
+    try:
+        return _COST_MODELS[name]
+    except KeyError as e:
+        raise ValueError(
+            f"Invalid constraints['cost_model']: {name}. Valid options are {tuple(_COST_MODELS)}."
+        ) from e
+
+
+def normalize_auto_quantize_constraints(
+    model: nn.Module, constraints: dict[str, Any] | None
+) -> dict[str, Any]:
+    """Validate and normalize AutoQuantize constraints."""
+    constraints = (
+        {"effective_bits": DEFAULT_AUTO_QUANTIZE_EFFECTIVE_BITS}
+        if constraints is None
+        else dict(constraints)
+    )
+    unexpected_constraint_keys = set(constraints) - AUTO_QUANTIZE_CONSTRAINT_KEYS
+    if unexpected_constraint_keys:
+        raise ValueError(
+            f"Unsupported auto_quantize constraints: {unexpected_constraint_keys}. "
+            "Supported constraints are 'effective_bits', 'cost_model', and 'cost'."
+        )
+
+    cost_model_name = constraints.get("cost_model", COST_MODEL_WEIGHT)
+    if not isinstance(cost_model_name, str):
+        raise ValueError("constraints['cost_model'] must be a string when provided.")
+    cost_model = get_auto_quantize_cost_model(cost_model_name)
+
+    cost_constraints = constraints.get("cost", {})
+    if cost_constraints is None:
+        cost_constraints = {}
+    if not isinstance(cost_constraints, dict):
+        raise ValueError("constraints['cost'] must be a dict when provided.")
+    cost_constraints = cost_model.normalize_cost_constraints(model, dict(cost_constraints))
+
+    constraints["cost_model"] = cost_model.name
+    if cost_constraints or cost_model.name == COST_MODEL_ACTIVE_MOE:
+        constraints["cost"] = cost_constraints
+    else:
+        constraints.pop("cost", None)
+    return constraints

modelopt/torch/quantization/algorithms.py

@@ -40,6 +40,14 @@ from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelSt
 
 from . import config as mtq_config
 from . import model_calib
+from ._auto_quantize_cost import (
+    ACTIVE_MOE_EXPERT_RATIO_KEY,
+    AUTO_QUANTIZE_CONSTRAINT_KEYS,
+    COST_MODEL_ACTIVE_MOE,
+    COST_MODEL_WEIGHT,
+    get_auto_quantize_cost_model,
+    normalize_auto_quantize_constraints,
+)
 from .config import QuantizeConfig, QuantizerAttributeConfig, QuantizerCfgEntry
 from .conversion import set_quantizer_by_cfg
 from .nn import QuantLinearConvBase, QuantModule, SequentialQuantizer, TensorQuantizer
@@ -206,6 +214,7 @@ class QuantRecipeHparam(Hparam):
         score_modules: list[nn.Module] | None = None,
         name: str | None = None,
         quant_module_names: list[str] | None = None,
+        cost_weight: float = 1.0,
     ) -> None:
         """Initializes Hparam with original value and choices."""
         choices = sorted({*(choices if choices else []), QuantRecipe(quant_cfg=None)})
@@ -213,6 +222,8 @@ class QuantRecipeHparam(Hparam):
 
         self.name = name
         self.quant_module_names = quant_module_names or []
+        assert cost_weight > 0.0, "cost_weight must be positive."
+        self.cost_weight = cost_weight
 
         self.quant_modules = list(set(quant_modules or []))
         self.score_modules = list(set(score_modules or self.quant_modules))
@@ -305,15 +316,18 @@ class QuantRecipeHparam(Hparam):
             total_score += importance.item()
         return total_score
 
-    def get_cost(self, recipe: QuantRecipe) -> float:
+    def get_cost(self, recipe: QuantRecipe, cost_weight: float | None = None) -> float:
         """Get the cost for a given recipe.
 
         The cost is the total weight size of the quantizable modules multiplied by
         the compression ratio of the recipe.
         """
+        cost_weight = self.cost_weight if cost_weight is None else cost_weight
         cost = 0
         for quant_module in self.quant_modules:
-            weight_size = _AutoQuantizeBaseSearcher._get_total_weight_size([quant_module])
+            weight_size = (
+                _AutoQuantizeBaseSearcher._get_total_weight_size([quant_module]) * cost_weight
+            )
             parallel_state = getattr(quant_module, "parallel_state", None)
 
             if parallel_state is None:
@@ -343,7 +357,7 @@ class QuantRecipeHparam(Hparam):
     @property
     def attrs(self) -> list[str]:
         """Return the attributes of the hparam for repr."""
-        return ["name", *super().attrs]
+        return ["name", "cost_weight", *super().attrs]
 
 
 class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
@@ -383,6 +397,9 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
             "disabled_layers": None,
             "verbose": is_master(),
             "checkpoint": None,
+            "cost_model": COST_MODEL_WEIGHT,
+            "cost": {},
+            "active_moe_expert_ratio": None,
         }
 
     @property
@@ -390,6 +407,10 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
         """Get the default state dict for AutoQuantize."""
         return {
             "method": self.method_name,
+            "cost_model": "weight",
+            "cost": {},
+            "active_moe_expert_ratio": None,
+            "cost_denominator": None,
             "candidate_stats": defaultdict(dict),
             "quantizer_states": {},
             "best": {"recipe": {}, "constraints": {}, "score": float("inf"), "is_satisfied": False},
@@ -547,6 +568,10 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
             quant_modules = [module for module, _, _, _ in module_info_list]
             disabled = any(disabled for _, _, disabled, _ in module_info_list)
             score_modules = [score_module for _, _, _, score_module in module_info_list]
+            quant_module_names = [name for _, name, _, _ in module_info_list]
+            cost_weight = self._cost_model.module_cost_weight(
+                quant_module_names, self.config["cost"]
+            )
 
             _quant_recipes = None if disabled else quant_recipes
             hparam = QuantRecipeHparam(
@@ -554,7 +579,8 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
                 quant_modules=quant_modules,
                 score_modules=score_modules,
                 name=str(group_key),
-                quant_module_names=[name for _, name, _, _ in module_info_list],
+                quant_module_names=quant_module_names,
+                cost_weight=cost_weight,
             )
 
             for module in quant_modules:
@@ -603,6 +629,7 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
             self.candidate_stats[name]["scores"] = scores
             self.candidate_stats[name]["costs"] = costs
             self.candidate_stats[name]["module_names"] = hparam.quant_module_names
+            self.candidate_stats[name]["cost_weight"] = hparam.cost_weight
 
     def _run_func(self, func, num_iters=1, desc=""):
         for i, data in tqdm(
@@ -621,16 +648,40 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
         from .utils import get_quantizer_state_dict, set_quantizer_state_dict
 
         super().before_search()
+        self.constraints = normalize_auto_quantize_constraints(self.model, self.constraints)
+        self.config["cost_model"] = self.constraints["cost_model"]
+        self.config["cost"] = self.constraints.get("cost", {})
+        self.config["active_moe_expert_ratio"] = self.config["cost"].get(
+            ACTIVE_MOE_EXPERT_RATIO_KEY
+        )
+        cost_model = get_auto_quantize_cost_model(self.config["cost_model"])
         restored_method = getattr(self, "method", None)
         if self.candidate_stats and restored_method not in (None, self.method_name):
             raise ValueError(
                 f"Checkpoint method '{restored_method}' does not match current method "
                 f"'{self.method_name}'. Use a different checkpoint path."
             )
+        restored_cost_model = getattr(self, "cost_model", "weight")
+        restored_active_moe_expert_ratio = getattr(self, "active_moe_expert_ratio", None)
+        if self.candidate_stats and (
+            restored_cost_model != self.config["cost_model"]
+            or restored_active_moe_expert_ratio != self.config["active_moe_expert_ratio"]
+        ):
+            raise ValueError(
+                "Checkpoint AutoQuantize cost model does not match current search config: "
+                f"checkpoint=({restored_cost_model}, {restored_active_moe_expert_ratio}), "
+                f"current=({self.config['cost_model']}, {self.config['active_moe_expert_ratio']}). "
+                "Use a different checkpoint path."
+            )
         self.method = self.method_name
+        self.cost_model = self.config["cost_model"]
+        self.cost = self.config["cost"]
+        self.active_moe_expert_ratio = self.config["active_moe_expert_ratio"]
+        self.cost_denominator = getattr(self, "cost_denominator", None)
 
         search_recipes = self._get_search_recipes(self.config["quantization_formats"])
         self._verify_constraint(search_recipes)
+        self._cost_model = cost_model
         self.insert_hparams_after_merge_rules(
             self.model, search_recipes, self.config["disabled_layers"]
         )
@@ -731,6 +782,12 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
         }
         return constraints, "weight_size_after_compression"
 
+    def _get_search_lower_bounds(self):
+        cost_model = getattr(self, "cost_model", getattr(self, "config", {}).get("cost_model"))
+        if cost_model == COST_MODEL_ACTIVE_MOE:
+            return [0.99, 0.90, None]
+        return [None, 0.99, 0.90]
+
     @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."""
@@ -738,14 +795,29 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
     def run_search(self):
         """Search for the best per-layer quantization configuration and return the best model and configuration."""
         verbose = self.config["verbose"]
-        assert len(self.constraints) == 1 and "effective_bits" in self.constraints, (
-            f"`constraints` must contain only 'effective_bits' constraint. "
-            f"Got {self.constraints.keys()}"
+        assert "effective_bits" in self.constraints and (
+            set(self.constraints) <= AUTO_QUANTIZE_CONSTRAINT_KEYS
+        ), (
+            "`constraints` must contain 'effective_bits' and may contain 'cost_model' and 'cost'. "
+            f"Got {self.constraints.keys()}."
         )
 
         compression = self._get_formatted_weight_compression_constraint()
-        total_weight_size = self._get_total_weight_size(self.model.modules())
+        total_weight_size = self._cost_model.total_weight_size(
+            self.model.named_modules(), self._is_auto_quantize_module, self.config["cost"]
+        )
+        self.cost_denominator = total_weight_size
         max_weight_size = total_weight_size * compression
+        if verbose:
+            print_rank_0(
+                "AutoQuantize cost model: "
+                f"{self.config['cost_model']}"
+                + (
+                    f" (active_moe_expert_ratio={self.config['active_moe_expert_ratio']})"
+                    if self.config["cost_model"] == COST_MODEL_ACTIVE_MOE
+                    else ""
+                )
+            )
 
         # Run the search with stats to get the best recipe and whether the constraints are satisfied
         best_recipe_info, is_satisfied = self.run_search_with_stats(max_weight_size, verbose)
@@ -1050,7 +1122,7 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher):
         """
         # TODO: Do this only for rank 0 in the respective pipeline group
 
-        for lower_bound in [None, 0.99, 0.90]:
+        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
@@ -1379,7 +1451,9 @@ def _resolve_best_recipe(search_state, constraints, verbose=False):
     effective_bits = constraints["effective_bits"]
     compression = effective_bits / 16.0
     candidate_stats = search_state["candidate_stats"]
-    total_weight_size = sum(s["costs"][-1] for s in candidate_stats.values())
+    total_weight_size = search_state.get("cost_denominator") or sum(
+        s["costs"][-1] for s in candidate_stats.values()
+    )
     max_weight_size = total_weight_size * compression
     method = search_state["method"]
 
@@ -1393,6 +1467,21 @@ def _resolve_best_recipe(search_state, constraints, verbose=False):
         )
 
     searcher.candidate_stats = candidate_stats
+    searcher.cost_model = search_state.get("cost_model", COST_MODEL_WEIGHT)
+    searcher.cost = search_state.get("cost", {})
+    searcher.active_moe_expert_ratio = search_state.get("active_moe_expert_ratio")
+    if (
+        searcher.cost_model == COST_MODEL_ACTIVE_MOE
+        and not searcher.cost
+        and searcher.active_moe_expert_ratio is not None
+    ):
+        searcher.cost = {ACTIVE_MOE_EXPERT_RATIO_KEY: searcher.active_moe_expert_ratio}
+    searcher.config = {
+        **searcher.default_search_config,
+        "cost_model": searcher.cost_model,
+        "cost": searcher.cost,
+        "active_moe_expert_ratio": searcher.active_moe_expert_ratio,
+    }
     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()}

modelopt/torch/quantization/model_quant.py

@@ -20,14 +20,14 @@ import inspect
 import os
 import warnings
 from collections.abc import Callable, Iterable
-from typing import Any
+from typing import Any, cast
 
 import torch
 import torch.nn as nn
 
 import modelopt.torch.quantization as mtq
 from modelopt.torch.opt import apply_mode
-from modelopt.torch.opt.searcher import ForwardLoop
+from modelopt.torch.opt.searcher import ConstraintsDict, ForwardLoop
 from modelopt.torch.opt.utils import forward_with_reshard
 from modelopt.torch.quantization.config import QuantizeConfig
 from modelopt.torch.quantization.conversion import (
@@ -268,7 +268,7 @@ _AUTO_QUANTIZE_SUPPORTED_ALGORITHMS = {
 
 def auto_quantize(
     model: nn.Module,
-    constraints: dict[str, float | str] = {"effective_bits": 4.8},
+    constraints: dict[str, Any] | None = None,
     quantization_formats: list[dict[str, Any] | str] = [
         mtq.NVFP4_AWQ_LITE_CFG,
         mtq.FP8_DEFAULT_CFG,
@@ -301,16 +301,26 @@ def auto_quantize(
 
     Args:
         model: A pytorch model with quantizer modules.
-        constraints: Constraints for the search. Currently we support only ``effective_bits``.
-            ``effective_bits`` specifies the effective number of bits for the quantized model.
+        constraints: Constraints for the search. ``effective_bits`` specifies the effective number
+            of bits for the quantized model and defaults to 4.8. ``cost_model`` selects the metric
+            used for the effective-bits constraint and currently supports ``"weight"`` (default)
+            and ``"active_moe"``. Additional cost-model parameters are provided through the nested
+            ``cost`` dict.
 
             Here is an example for valid ``effective_bits`` argument:
 
             .. code-block:: python
 
-                # For an effective quantization bits of 4.8
+                # For the default AutoQuantize effective-bits target
                 constraints = {"effective_bits": 4.8}
 
+                # For active-MoE accounting where 2 of 8 routed experts are active per token
+                constraints = {
+                    "effective_bits": 4.8,
+                    "cost_model": "active_moe",
+                    "cost": {"active_moe_expert_ratio": 0.25},
+                }
+
         quantization_formats: A list of quantization format config dictionaries or string names to search for.
             Each config dictionary should be valid as a ``config`` argument in
             :meth:`quantize <modelopt.torch.quantization.model_quant.quantize>`.
@@ -533,7 +543,8 @@ def auto_quantize(
     }
     # Disable all quantizers; AutoQuantize will enable the needed ones
     set_quantizer_by_cfg(model, [{"quantizer_name": "*", "enable": False}])
-    searcher.search(model, constraints, config=search_config)  # type: ignore[arg-type]
+    search_constraints = cast("ConstraintsDict", constraints or {})
+    searcher.search(model, search_constraints, config=search_config)
 
     return model, searcher.state_dict()
 

tests/unit/torch/quantization/test_autoquant.py

@@ -15,6 +15,7 @@
 
 import copy
 import io
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -23,7 +24,9 @@ from _test_utils.torch.quantization.models import SimpleConv, SimpleConvLinear,
 
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization._auto_quantize_cost import infer_active_moe_expert_ratio
 from modelopt.torch.quantization.algorithms import (
+    AutoQuantizeGradientSearcher,
     QuantRecipe,
     QuantRecipeHparam,
     estimate_quant_compression,
@@ -62,6 +65,31 @@ class TransformerBlock(torch.nn.Module):
         return torch.randn(1, 4, 32)
 
 
+class _AutoQuantMoeModel(torch.nn.Module):
+    def __init__(self, num_experts_attr="num_experts"):
+        super().__init__()
+        self.config = SimpleNamespace(text_config=SimpleNamespace(num_experts_per_tok=2))
+        setattr(self.config.text_config, num_experts_attr, 8)
+        self.mlp = torch.nn.Module()
+        self.mlp.experts = torch.nn.ModuleList()
+        for _ in range(2):
+            expert = torch.nn.Module()
+            expert.gate_proj = torch.nn.Linear(32, 32)
+            expert.up_proj = torch.nn.Linear(32, 32)
+            expert.down_proj = torch.nn.Linear(32, 32)
+            self.mlp.experts.append(expert)
+        self.mlp.shared_expert = torch.nn.Linear(32, 32)
+
+    def forward(self, x):
+        y = self.mlp.shared_expert(x)
+        for expert in self.mlp.experts:
+            y = y + expert.down_proj(expert.gate_proj(x) + expert.up_proj(x))
+        return y
+
+    def get_input(self):
+        return torch.randn(1, 4, 32)
+
+
 @pytest.mark.parametrize(
     ("quant_cfg", "other_quant_cfg", "is_less_than"),
     [
@@ -109,6 +137,89 @@ def test_quant_recipe_hparam():
     assert torch.allclose(output_test, output_ref)
 
 
+def test_quant_recipe_hparam_cost_weight():
+    model_test = mtq.quantize(torch.nn.Linear(4, 16), mtq.INT8_DEFAULT_CFG)
+    search_recipes = [QuantRecipe(mtq.INT8_DEFAULT_CFG)]
+    hparam = QuantRecipeHparam(
+        search_recipes,
+        quant_modules=[model_test],
+        quant_module_names=["layers.0.mlp.experts.0.down_proj"],
+        cost_weight=0.25,
+    )
+
+    dense_cost = hparam.get_cost(QuantRecipe(quant_cfg=None))
+    int8_cost = hparam.get_cost(QuantRecipe(mtq.INT8_DEFAULT_CFG))
+
+    assert dense_cost == pytest.approx(model_test.weight.numel() * 0.25)
+    assert int8_cost == pytest.approx(model_test.weight.numel() * 0.25 * 0.5)
+
+
+@pytest.mark.parametrize("num_experts_attr", ["num_experts", "num_local_experts"])
+def test_auto_quantize_active_moe_cost_model(num_experts_attr):
+    model = _AutoQuantMoeModel(num_experts_attr)
+
+    _, search_history = mtq.auto_quantize(
+        model,
+        constraints={"effective_bits": 6.0, "cost_model": "active_moe"},
+        quantization_formats=[mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, mtq.INT8_DEFAULT_CFG],
+        data_loader=[model.get_input() for _ in range(2)],
+        forward_step=lambda model, batch: model(batch),
+        loss_func=lambda output, data: output.sum(),
+        num_calib_steps=2,
+        num_score_steps=2,
+    )
+
+    assert search_history["cost_model"] == "active_moe"
+    assert search_history["active_moe_expert_ratio"] == pytest.approx(0.25)
+    weighted_no_quant_cost = sum(
+        stats["costs"][-1] for stats in search_history["candidate_stats"].values()
+    )
+    assert search_history["cost_denominator"] == pytest.approx(weighted_no_quant_cost)
+    routed_stats = [
+        stats
+        for stats in search_history["candidate_stats"].values()
+        if any("mlp.experts" in name for name in stats["module_names"])
+    ]
+    shared_stats = [
+        stats
+        for stats in search_history["candidate_stats"].values()
+        if any("mlp.shared_expert" in name for name in stats["module_names"])
+    ]
+    assert routed_stats
+    assert shared_stats
+    assert all(stats["cost_weight"] == pytest.approx(0.25) for stats in routed_stats)
+    assert all(stats["cost_weight"] == pytest.approx(1.0) for stats in shared_stats)
+    assert all("active_costs" not in stats for stats in search_history["candidate_stats"].values())
+
+
+def test_active_moe_ratio_requires_single_config_object():
+    model = torch.nn.Module()
+    model.config = SimpleNamespace(
+        num_experts_per_tok=2,
+        text_config=SimpleNamespace(num_experts=8),
+    )
+
+    assert infer_active_moe_expert_ratio(model) is None
+
+
+def test_active_moe_search_prefers_budget_lower_bound():
+    searcher = AutoQuantizeGradientSearcher()
+    searcher.config = {"cost_model": "active_moe"}
+    searcher.cost_model = "active_moe"
+    searcher.candidate_stats = {
+        "layers.0.mlp.quant_recipe": {
+            "formats": ["under_budget", "near_budget"],
+            "costs": [1.0, 4.95],
+            "scores": [0.0, 10.0],
+        }
+    }
+
+    best_recipes, is_satisfied = searcher.run_search_with_stats(5.0)
+
+    assert is_satisfied
+    assert best_recipes["layers.0.mlp.quant_recipe"]["format"] == "near_budget"
+
+
 # use this config to test custom quantization config
 INT8_CUSTOM_QUANT_TEST_CFG = {
     "quant_cfg": [