Refine quantization config schema handling

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>

Author: shengliangxu

modelopt/torch/quantization/algorithms.py

@@ -104,39 +104,40 @@ def estimate_quant_compression_for_quantizer(quantizer_attr_cfg):
 
 
 QuantRecipeConfig = str | Mapping[str, Any] | QuantizeConfig | None
+QuantizationFormatConfig = QuantRecipeConfig
+NamedQuantRecipeConfig = tuple[str | Mapping[str, Any] | QuantizeConfig, str]
 
 
 class QuantRecipe(CustomHPType):
     """A subclass of QuantizeConfig enabling auto_quantize specific configurations.
 
     Args:
         quant_cfg: str, QuantizeConfig, mapping, or None. A mapping is used for custom quantization formats.
-        name: name for custom quantization formats. Only used if quantization format is a custom
-            format not available in :mod:`modelopt.torch.quantization.config`.
+        name: Required display/search name when ``quant_cfg`` is not ``None``. Must be
+            ``None`` when ``quant_cfg=None``, which uses the built-in ``"NONE"`` recipe name.
     """
 
     def __init__(self, quant_cfg: QuantRecipeConfig = None, name: str | None = None):
         """Initialize the QuantRecipe with the quantization configuration."""
-        name = self.get_auto_name_for_config(quant_cfg) or name
-
         if quant_cfg is None:
+            if name is not None:
+                raise ValueError("name must be None when quant_cfg is None")
+            name = "NONE"
             self.config = mtq_config.QuantizeConfig(
                 quant_cfg=[mtq_config.QuantizerCfgEntry(quantizer_name="*", enable=False)]
             )
         else:
+            if name is None:
+                raise ValueError("name must be provided when quant_cfg is not None")
             if isinstance(quant_cfg, str):
                 assert hasattr(mtq_config, quant_cfg), f"Unknown quantization format {quant_cfg}"
                 quant_cfg = getattr(mtq_config, quant_cfg)
-            elif not isinstance(quant_cfg, QuantizeConfig) and name is None:
-                raise ValueError("name must be provided for custom quantization formats")
 
             self.config = (
                 quant_cfg.model_copy(deep=True)
                 if isinstance(quant_cfg, QuantizeConfig)
                 else mtq_config.QuantizeConfig.model_validate(quant_cfg)
             )
-            if name is None:
-                raise ValueError("name must be provided for custom quantization formats")
 
         # Disable KV Cache quantization
         # Currently KV Cache quantization is enabled for some quantization formats and disabled for others
@@ -211,6 +212,31 @@ def fold_pqs_to_weights(model):
                         model_calib._apply_weight_pre_quant_scale(module, weight_pqs)
 
 
+def _validate_named_auto_quantize_formats(
+    quantization_formats: Any,
+):
+    """Validate the internal AutoQuantize format protocol."""
+    error_msg = (
+        "`quantization_formats` must be a list of (quant_cfg, name) tuples. "
+        "Normalize public inputs before calling the AutoQuantize searcher."
+    )
+    if not isinstance(quantization_formats, list):
+        raise TypeError(error_msg)
+
+    for entry in quantization_formats:
+        if not isinstance(entry, tuple) or len(entry) != 2:
+            raise TypeError(error_msg)
+
+        quant_cfg, name = entry
+        if quant_cfg is None or not isinstance(quant_cfg, str | QuantizeConfig | Mapping):
+            raise TypeError(
+                "Each named quantization format must contain a string, mapping, or "
+                "QuantizeConfig as the first tuple item."
+            )
+        if not isinstance(name, str) or not name:
+            raise TypeError("Each named quantization format must provide a non-empty name.")
+
+
 class QuantRecipeHparam(Hparam):
     """An Hparam for quantization recipes.
 
@@ -231,7 +257,7 @@ def __init__(
         quant_module_names: list[str] | None = None,
     ) -> None:
         """Initializes Hparam with original value and choices."""
-        choices = sorted({*(choices if choices else []), QuantRecipe(quant_cfg=None)})
+        choices = sorted({*(choices or []), QuantRecipe(quant_cfg=None)})
         super().__init__(choices, original=choices[0])
 
         self.name = name
@@ -398,7 +424,10 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
     def default_search_config(self):
         """Get the default config for the searcher."""
         return {
-            "quantization_formats": ["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"],
+            "quantization_formats": [
+                ("NVFP4_DEFAULT_CFG", "NVFP4_DEFAULT_CFG"),
+                ("FP8_DEFAULT_CFG", "FP8_DEFAULT_CFG"),
+            ],
             "data_loader": None,
             "num_calib_steps": 512,
             "num_score_steps": 128,
@@ -428,6 +457,7 @@ def sanitize_search_config(self, config: SearchConfig | None) -> SearchConfig:
         assert config["forward_step"] is not None, (
             "`forward_step` must be provided for `auto_quantize`."
         )
+        _validate_named_auto_quantize_formats(config["quantization_formats"])
         return config
 
     def load_search_checkpoint(self) -> bool:
@@ -440,13 +470,11 @@ def _is_auto_quantize_module(module):
         ) and isinstance(module, QuantModule)
 
     @staticmethod
-    def _get_search_recipes(quantization_formats):
+    def _get_search_recipes(quantization_formats: Sequence[NamedQuantRecipeConfig]):
         return sorted(
             {
-                QuantRecipe(quant_cfg=q[0], name=q[1])
-                if isinstance(q, tuple)
-                else QuantRecipe(quant_cfg=q)
-                for q in quantization_formats
+                QuantRecipe(quant_cfg=quant_cfg, name=name)
+                for quant_cfg, name in quantization_formats
             }
         )
 

modelopt/torch/quantization/model_quant.py

@@ -19,7 +19,7 @@
 import inspect
 import os
 import warnings
-from collections.abc import Callable, Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping, Sequence
 from typing import Any
 
 import torch
@@ -36,7 +36,13 @@
 )
 from modelopt.torch.utils import atomic_print
 
-from .algorithms import AutoQuantizeGradientSearcher, AutoQuantizeKLDivSearcher, QuantRecipe
+from .algorithms import (
+    AutoQuantizeGradientSearcher,
+    AutoQuantizeKLDivSearcher,
+    NamedQuantRecipeConfig,
+    QuantizationFormatConfig,
+    QuantRecipe,
+)
 from .algorithms import get_auto_quantize_config as _get_auto_quantize_config
 from .config import QuantizeAlgoCfgType
 from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg
@@ -270,10 +276,38 @@ def forward_loop(model) -> None:
 }
 
 
+def _normalize_auto_quantize_formats(
+    quantization_formats: Sequence[QuantizationFormatConfig],
+) -> list[NamedQuantRecipeConfig]:
+    """Normalize public auto_quantize format inputs into named search entries."""
+    processed_quantization_formats: list[NamedQuantRecipeConfig] = []
+    for i, quant_cfg in enumerate(quantization_formats):
+        if quant_cfg is None:
+            continue
+        if isinstance(quant_cfg, tuple):
+            raise TypeError(
+                "Named quantization format tuples are internal to AutoQuantize search; "
+                "pass raw configs to auto_quantize()."
+            )
+
+        name = QuantRecipe.get_auto_name_for_config(quant_cfg)
+        if name is None:
+            name = f"CUSTOM_{i}"
+            warnings.warn(
+                "Received custom quantization formats for search, auto_quantize results "
+                f"may not be optimal. This config will be displayed as {name}"
+            )
+
+        processed_quantization_formats.append((quant_cfg, name))
+
+    assert len(processed_quantization_formats) > 0, "`quantization_formats` should not be empty"
+    return processed_quantization_formats
+
+
 def auto_quantize(
     model: nn.Module,
     constraints: dict[str, float | str] = {"effective_bits": 4.8},
-    quantization_formats: list[QuantizeConfig | Mapping[str, Any] | str | None] = [
+    quantization_formats: list[QuantizationFormatConfig] = [
         mtq.NVFP4_AWQ_LITE_CFG,
         mtq.FP8_DEFAULT_CFG,
     ],
@@ -319,6 +353,7 @@ def auto_quantize(
             Each config dictionary should be valid as a ``config`` argument in
             :meth:`quantize <modelopt.torch.quantization.model_quant.quantize>`.
             The supported quantization format names are as listed by :attr:`modelopt.torch.quantization.config.choices`.
+            Custom configs without a built-in name are assigned ``CUSTOM_<index>`` display names internally.
 
             Internally we always add "do not quantize" as a choice. Therefore, it is possible that a layer is
             not quantized by any of the quantization formats.
@@ -484,21 +519,7 @@ def forward_backward_step(model, batch) -> None:
         might not be readily deployable to TensorRT-LLM yet.
 
     """
-    processed_quantization_formats = []
-    for i, quant_cfg in enumerate(quantization_formats):
-        if quant_cfg is None:
-            continue
-
-        name = QuantRecipe.get_auto_name_for_config(quant_cfg)
-        if name is None:
-            name = f"CUSTOM_{i}"
-            warnings.warn(
-                f"Received custom quantization formats for search, auto_quantize results may not be optimal. "
-                f"This config will be displayed as {name}"
-            )
-        processed_quantization_formats.append((quant_cfg, name))
-
-    assert len(processed_quantization_formats) > 0, "`quantization_formats` should not be empty"
+    processed_quantization_formats = _normalize_auto_quantize_formats(quantization_formats)
 
     for quant_cfg, name in processed_quantization_formats:
         algo = QuantRecipe(quant_cfg, name=name).config.algorithm

tests/unit/torch/quantization/test_autoquant.py

@@ -24,11 +24,13 @@
 import modelopt.torch.opt as mto
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.algorithms import (
+    AutoQuantizeGradientSearcher,
     QuantRecipe,
     QuantRecipeHparam,
     estimate_quant_compression,
 )
 from modelopt.torch.quantization.config import _base_disable_all, _default_disabled_quantizer_cfg
+from modelopt.torch.quantization.model_quant import _normalize_auto_quantize_formats
 from modelopt.torch.utils import safe_load
 from modelopt.torch.utils.distributed import DistributedProcessGroup
 
@@ -62,6 +64,11 @@ def get_input(self):
         return torch.randn(1, 4, 32)
 
 
+def _recipe(quant_cfg):
+    name = None if quant_cfg is None else QuantRecipe.get_auto_name_for_config(quant_cfg)
+    return QuantRecipe(quant_cfg, name=name)
+
+
 @pytest.mark.parametrize(
     ("quant_cfg", "other_quant_cfg", "is_less_than"),
     [
@@ -71,30 +78,85 @@ def get_input(self):
     ],
 )
 def test_quant_recipe(quant_cfg, other_quant_cfg, is_less_than):
-    qr_this = QuantRecipe(quant_cfg)
-    qr_other = QuantRecipe(other_quant_cfg)
+    qr_this = _recipe(quant_cfg)
+    qr_other = _recipe(other_quant_cfg)
     assert (qr_this < qr_other) == is_less_than
 
-    qr_this_duplicate = QuantRecipe(quant_cfg)
+    qr_this_duplicate = _recipe(quant_cfg)
     assert qr_this_duplicate in {qr_this}
 
 
-def test_quant_recipe_custom_quantize_config_requires_name():
-    custom_cfg = mtq.QuantizeConfig(
+def _custom_quantize_config(path):
+    return mtq.QuantizeConfig(
         quant_cfg=[
             mtq.QuantizerCfgEntry(
-                quantizer_name="*weight_quantizer",
+                quantizer_name=path,
                 cfg=mtq.QuantizerAttributeConfig(num_bits=8, axis=None),
             )
         ]
     )
 
+
+def test_quant_recipe_custom_quantize_config_requires_name():
+    custom_cfg = _custom_quantize_config("*custom_weight_quantizer")
+
     with pytest.raises(ValueError, match="name must be provided"):
         QuantRecipe(custom_cfg)
 
     assert str(QuantRecipe(custom_cfg, name="custom_cfg")).startswith("custom_cfg(")
 
 
+def test_quant_recipe_none_requires_no_name():
+    assert str(QuantRecipe(quant_cfg=None)).startswith("NONE(")
+
+    with pytest.raises(ValueError, match="name must be None"):
+        QuantRecipe(quant_cfg=None, name="NONE")
+
+
+def test_quant_recipe_honors_explicit_name():
+    assert str(QuantRecipe(mtq.INT8_DEFAULT_CFG, name="int8_alias")).startswith("int8_alias(")
+
+
+def test_auto_quantize_search_config_requires_named_formats():
+    custom_a = _custom_quantize_config("*custom_weight_quantizer_a")
+    custom_b = _custom_quantize_config("*custom_weight_quantizer_b")
+    searcher = AutoQuantizeGradientSearcher()
+
+    with pytest.warns(UserWarning) as records:
+        quantization_formats = _normalize_auto_quantize_formats([custom_a, custom_b])
+
+    assert quantization_formats == [(custom_a, "CUSTOM_0"), (custom_b, "CUSTOM_1")]
+    assert any("CUSTOM_0" in str(record.message) for record in records)
+    assert any("CUSTOM_1" in str(record.message) for record in records)
+
+    config = searcher.sanitize_search_config(
+        {
+            "quantization_formats": quantization_formats,
+            "data_loader": [torch.randn(1)],
+            "forward_step": lambda model, data: data,
+            "loss_func": lambda output, data: output.sum(),
+        }
+    )
+    assert config["quantization_formats"] == quantization_formats
+
+    with pytest.raises(TypeError, match="Named quantization format tuples are internal"):
+        _normalize_auto_quantize_formats([(custom_a, "custom_a")])
+
+    with pytest.raises(TypeError, match="must be a list of"):
+        searcher.sanitize_search_config(
+            {
+                "quantization_formats": [custom_a],
+                "data_loader": [torch.randn(1)],
+                "forward_step": lambda model, data: data,
+                "loss_func": lambda output, data: output.sum(),
+            }
+        )
+
+    recipes = AutoQuantizeGradientSearcher._get_search_recipes(config["quantization_formats"])
+    assert {str(recipe).split("(", 1)[0] for recipe in recipes} == {"CUSTOM_0", "CUSTOM_1"}
+    assert len(set(recipes)) == 2
+
+
 def test_quant_recipe_hparam():
     model_test = torch.nn.Linear(4, 16)
     model_ref = torch.nn.Linear(4, 16)
@@ -104,20 +166,20 @@ def test_quant_recipe_hparam():
     model_ref = mtq.quantize(model_ref, mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
 
     search_recipes = [
-        QuantRecipe(mtq.INT8_DEFAULT_CFG),
-        QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG),
+        _recipe(mtq.INT8_DEFAULT_CFG),
+        _recipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG),
     ]
     hparam = QuantRecipeHparam(
         search_recipes,
         quant_modules=[model_test],
     )
     model_test._register_hparam("quant_recipe", hparam)
-    assert model_test.quant_recipe == QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    assert model_test.quant_recipe == _recipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
     assert model_test.get_hparam("quant_recipe").choices == sorted(
         [*search_recipes, QuantRecipe(quant_cfg=None)]
     )
 
-    model_test.quant_recipe = QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
+    model_test.quant_recipe = _recipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG)
     inputs = torch.randn(1, 4, 4)
     output_test = model_test(inputs)
     output_ref = model_ref(inputs)
@@ -244,7 +306,7 @@ def test_auto_quantize_disabled_layers_no_poison():
 
     assert not best_model.mlp.input_quantizer.is_enabled
     hparam = best_model.attn.q_proj.get_hparam("quant_recipe")
-    assert QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG) in hparam.choices
+    assert _recipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG) in hparam.choices
 
 
 INT4INT8_AWQ_CFG = {