Make quantization config parameters searchable

olive/engine/footprint.py

@@ -58,8 +58,19 @@ class FootprintNode(ConfigBase):
     end_time: float = 0
 
     def update(self, **kwargs):
+        # Callers pass kwargs by field alias (e.g. Footprint.record forwards `model_config=...`),
+        # matching how FootprintNode(**kwargs) constructs a node. Pydantic accepts aliases at
+        # construction time, but setattr only accepts real field names. Since `model_config` is
+        # reserved by Pydantic v2, the field is declared as `model_config_data` with
+        # alias="model_config", so a raw setattr(self, "model_config", v) raises
+        # "FootprintNode object has no field model_config". Map aliases back to field names first.
+        alias_to_field = {
+            field_info.alias: field_name
+            for field_name, field_info in type(self).model_fields.items()
+            if field_info.alias is not None
+        }
         for k, v in kwargs.items():
-            setattr(self, k, v)
+            setattr(self, alias_to_field.get(k, k), v)
 
 
 class Footprint:

olive/passes/pytorch/gptq.py

@@ -20,6 +20,7 @@
     prepare_model,
     run_layerwise_quantization,
 )
+from olive.search.search_parameter import Categorical
 
 if TYPE_CHECKING:
     from olive.hardware.accelerator import AcceleratorSpec
@@ -41,6 +42,7 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
             "damp_percent": PassConfigParam(
                 type_=float,
                 default_value=0.01,
+                search_defaults=Categorical([0.001, 0.01, 0.1]),
                 description="Damping factor for quantization. Default value is 0.01.",
             ),
             "desc_act": PassConfigParam(

olive/passes/pytorch/quant_utils.py

@@ -26,6 +26,7 @@
 from olive.passes.pass_config import PassConfigParam
 from olive.passes.pytorch.common import inherit_hf_from_hf
 from olive.passes.pytorch.train_utils import get_calibration_dataset, load_hf_base_model
+from olive.search.search_parameter import Boolean, Categorical
 
 if TYPE_CHECKING:
     from olive.model import HfModelHandler
@@ -40,21 +41,25 @@ def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigPara
         "bits": PassConfigParam(
             type_=PrecisionBits,
             default_value=PrecisionBits.BITS4,
+            search_defaults=Categorical([PrecisionBits.BITS2, PrecisionBits.BITS4, PrecisionBits.BITS8]),
             description="quantization bits. Default value is 4",
         ),
         "group_size": PassConfigParam(
             type_=int,
             default_value=128,
+            search_defaults=Categorical([-1, 16, 32, 64, 128]),
             description="Block size for quantization. Default value is 128.",
         ),
         "sym": PassConfigParam(
             type_=bool,
             default_value=False,
+            search_defaults=Boolean(),
             description="Symmetric quantization. Default value is False.",
         ),
         "lm_head": PassConfigParam(
             type_=bool,
             default_value=False,
+            search_defaults=Boolean(),
             description="Whether to quantize the language model head. Default value is False.",
         ),
         **(

olive/passes/pytorch/selective_mixed_precision.py

@@ -22,7 +22,7 @@
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 from olive.passes.pytorch.quant_utils import get_qkv_quantization_groups
 from olive.passes.pytorch.train_utils import get_calibration_dataset, kl_div_loss, load_hf_base_model
-from olive.search.search_parameter import Categorical
+from olive.search.search_parameter import Boolean, Categorical
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Sequence
@@ -655,28 +655,25 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
             "algorithm": PassConfigParam(
                 type_=SelectiveMixedPrecision.Algorithm,
                 required=False,
-                search_defaults=Categorical(
-                    [
-                        SelectiveMixedPrecision.Algorithm.K_QUANT_DOWN,
-                        SelectiveMixedPrecision.Algorithm.K_QUANT_MIXED,
-                        SelectiveMixedPrecision.Algorithm.K_QUANT_LAST,
-                    ]
-                ),
+                search_defaults=Categorical(list(SelectiveMixedPrecision.Algorithm)),
                 description="The algorithm to use for mixed precision.",
             ),
             "bits": PassConfigParam(
                 type_=PrecisionBits,
                 default_value=PrecisionBits.BITS4,
+                search_defaults=Categorical([PrecisionBits.BITS2, PrecisionBits.BITS4, PrecisionBits.BITS8]),
                 description="The default precision bits.",
             ),
             "group_size": PassConfigParam(
                 type_=int,
                 default_value=128,
+                search_defaults=Categorical([-1, 16, 32, 64, 128]),
                 description="The default group size. Only used for snr, iqe and kld_gradient algorithms.",
             ),
             "sym": PassConfigParam(
                 type_=bool,
                 default_value=False,
+                search_defaults=Boolean(),
                 description=(
                     "Whether to use symmetric quantization by default. Only used for snr, iqe and kld_gradient"
                     " algorithms."
@@ -685,11 +682,13 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
             "high_bits": PassConfigParam(
                 type_=PrecisionBits,
                 default_value=PrecisionBits.BITS8,
+                search_defaults=Categorical([PrecisionBits.BITS8, PrecisionBits.BITS16]),
                 description="The high precision bits for selected layers.",
             ),
             "high_group_size": PassConfigParam(
                 type_=int,
                 default_value=None,
+                search_defaults=Categorical([None, -1, 16, 32, 64, 128]),
                 description=(
                     "The group size for high precision layers. Only used for snr, iqe and kld_gradient algorithms. If"
                     " None, use group_size."
@@ -698,6 +697,7 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
             "high_sym": PassConfigParam(
                 type_=bool,
                 default_value=None,
+                search_defaults=Categorical([None, True, False]),
                 description=(
                     "Whether to use symmetric quantization for high precision layers. Only used for snr, iqe and"
                     " kld_gradient algorithms. If None, use sym."
@@ -706,6 +706,7 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
             "ratio": PassConfigParam(
                 type_=float,
                 default_value=None,
+                search_defaults=Categorical([0.5, 0.8, 0.9, 0.95]),
                 description=(
                     "The ratio of default precision parameters to total parameters. Only used for snr, iqe and"
                     " kld_gradient algorithms. Must be provided when using these algorithms."

olive/search/search_parameter.py

@@ -52,10 +52,10 @@ class Categorical(SearchParameter):
 
     """
 
-    def __init__(self, support: Union[list[str], list[int], list[float], list[bool]]):
+    def __init__(self, support: Union[list[str | None], list[int | None], list[float | None], list[bool | None]]):
         self.support = support
 
-    def get_support(self) -> Union[list[str], list[int], list[float], list[bool]]:
+    def get_support(self) -> Union[list[str | None], list[int | None], list[float | None], list[bool | None]]:
         """Get the support for the search parameter."""
         return self.support