Make quantization config parameters searchable

olive/engine/engine.py

@@ -49,7 +49,7 @@ class Engine:
 
     def __init__(
         self,
-        olive_config: OlivePackageConfig = None,
+        olive_config: OlivePackageConfig | None = None,
         workflow_id: str = DEFAULT_WORKFLOW_ID,
         search_strategy: Optional[Union[dict[str, Any], SearchStrategyConfig]] = None,
         host: Optional[Union[dict[str, Any], "SystemConfig"]] = None,
@@ -61,7 +61,7 @@ def __init__(
     ):
         self.olive_config = olive_config or OlivePackageConfig.load_default_config()
         self.workflow_id = workflow_id
-        self.search_strategy = SearchStrategy(search_strategy) if search_strategy else None
+        self.search_strategy: SearchStrategy | None = SearchStrategy(search_strategy) if search_strategy else None
 
         # default host
         host = host or {"type": SystemType.Local}
@@ -332,6 +332,8 @@ def _run_no_search(
             failed_pass = pass_flow[len(model_ids)]
             logger.warning("Flow %s is pruned due to failed or invalid config for pass '%s'", pass_flow, failed_pass)
             return
+        else:
+            logger.debug("Signal: %s, %s", signal, model_ids)
 
         if signal is not None and not self.skip_saving_artifacts:
             results_path = artifacts_dir / "metrics.json"
@@ -422,6 +424,8 @@ def _run_search(
         )
         self.search_strategy.initialize(search_space_config, input_model_id, search_space_objectives)
 
+        logger.info("Search space contains %d search points ...", self.search_strategy.max_samples)
+
         for sample in self.search_strategy:  # pylint: disable=not-an-iterable
             self._compute_search_pass_configs(accelerator_spec, sample)
 
@@ -446,6 +450,8 @@ def _run_search(
                         exc_info=True,
                     )
 
+                logger.info("Signal: %s, %s, %s", signal, model_ids, sample.search_point)
+
             # record feedback signal
             self.search_strategy.record_feedback_signal(sample.search_point.index, signal, model_ids, should_prune)
 
@@ -643,7 +649,6 @@ def _run_passes(
             else:
                 logger.info("Run model evaluation for the final model...")
                 signal = self._evaluate_model(model_config, model_id, evaluator_config, accelerator_spec)
-            logger.debug("Signal: %s, %s", signal, model_ids)
         else:
             signal = None
             logger.warning("Skipping evaluation as model was pruned")

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