Update based on review feedback

Author: shaahji

examples/llama2/llama2_auto_opt.json

@@ -71,6 +71,7 @@
     },
     "auto_optimizer_config": {
         "precision": "fp32",
+        "use_model_builder": true,
         "train_data_config": "wikitext2_train",
         "calibration_data_config": "transformer_token_dummy_data",
         "accelerator": { "accelerator_type": "cpu", "execution_provider": "CPUExecutionProvider" },

olive/auto_optimizer/auto_optimizer.py

@@ -157,7 +157,7 @@ def register_subcommand(parser: ArgumentParser):
             nargs="*",
             default=None,
             help=(
-                "Dictionary of name to precision. Has to be even number of entreis with even "
+                "Dictionary of name to precision. Has to be even number of entries with even "
                 "entries being the keys and odd entries being the values. "
                 'Required only when output precision is "fp16" and MixedPrecisionOverrides pass is enabled.'
             ),
@@ -167,6 +167,10 @@ def register_subcommand(parser: ArgumentParser):
             "--use_ort_genai", action="store_true", help="Use OnnxRuntime generate() API to run the model"
         )
 
+        sub_parser.add_argument(
+            "--surgeries", type=str, nargs="*", default=None, help="List of graph surgeries to apply."
+        )
+
         add_search_options(sub_parser)
         add_remote_options(sub_parser)
         add_shared_cache_options(sub_parser)
@@ -293,6 +297,7 @@ def _get_passes_config(self, config: dict[str, Any], olive_config: OlivePackageC
             (("to_fixed_shape", "dim_param"), self.args.dynamic_to_fixed_shape_dim_param),
             (("to_fixed_shape", "dim_value"), self.args.dynamic_to_fixed_shape_dim_value),
             (("mixed_precision_overrides", "overrides_config"), mixed_precision_overrides_config),
+            (("surgeries", "surgeries"), [{"surgeon": surgeon} for surgeon in self.args.surgeries]),
         ]
         for keys, value in to_replace:
             if value is not None:
@@ -343,6 +348,10 @@ def _get_passes_config(self, config: dict[str, Any], olive_config: OlivePackageC
             # Remove QDQ encoding pass if not required
             passes_to_remove.add("mnb_to_qdq")
 
+        if not self.args.surgeries:
+            # Remove surgeon if user hasn't provided any surgeries
+            passes_to_remove.add("surgeries")
+
         # remove passes that are incompatible with the selected precision, provider, or device
         for pass_name in list(passes_config.keys()):
             pass_run_config = passes_config[pass_name]
@@ -439,6 +448,7 @@ def _get_passes_config(self, config: dict[str, Any], olive_config: OlivePackageC
             ("bnb4", {"type": "OnnxBnb4Quantization", "precision": Precision.NF4}),
             # post processing passes
             ("mnb_to_qdq", {"type": "MatMulNBitsToQDQ"}),
+            ("surgeries", {"type": "GraphSurgeries", "surgeries": {}}),
             ("split_model", {"type": "SplitModel"}),
             ("extract_adapters", {"type": "ExtractAdapters"}),
         ]

olive/cli/auto_opt.py

@@ -608,26 +608,5 @@
         "tf": [ "tensorflow==1.15.0" ],
         "torch-tensorrt": [ "torch-tensorrt" ],
         "tune-session-params": [ "psutil" ]
-    },
-    "auto_opt_passes": {
-        "capture": [ "CaptureSplitInfo" ],
-        "pt_finetune": [ "DoRA", "LoftQ", "LoHa", "LoKr", "LoRA", "QLoRA" ],
-        "pt_quantize": [ "QuaRot", "SpinQuant", "AutoAWQQuantizer", "GptqQuantizer" ],
-        "conversion": [ "OnnxConversion", "ModelBuilder" ],
-        "peephole": [ "OnnxPeepholeOptimizer" ],
-        "transformers": [ "OrtTransformersOptimization" ],
-        "io_converter": [ "OnnxIODataTypeConverter" ],
-        "prepare_qnn": [ "DynamicToFixedShape", "QNNPreprocess", "MixedPrecisionOverrides" ],
-        "onnx_quantize": [
-            "OnnxBnb4Quantization",
-            "OnnxMatMul4Quantizer",
-            "OnnxDynamicQuantization",
-            "OnnxStaticQuantization",
-            "NVModelOptQuantization",
-            "IncDynamicQuantization",
-            "IncStaticQuantization"
-        ],
-        "onnx_finetune": [ "OrtSessionParamsTuning" ],
-        "post_process": [ "MatMulNBitsToQDQ", "SplitModel", "ExtractAdapters" ]
     }
 }

olive/olive_config.json

@@ -64,3 +64,19 @@ def get_pass_module_config(self, pass_type: str) -> PassModuleConfig:
     def is_onnx_module(self, pass_type: str) -> bool:
         pass_module = self.get_pass_module_config(pass_type)
         return pass_module.module_path.startswith("olive.passes.onnx")
+
+    def is_openvino_module(self, pass_type: str) -> bool:
+        pass_module = self.get_pass_module_config(pass_type)
+        return pass_module.module_path.startswith("olive.passes.openvino")
+
+    def is_pytorch_module(self, pass_type: str) -> bool:
+        pass_module = self.get_pass_module_config(pass_type)
+        return pass_module.module_path.startswith("olive.passes.pytorch")
+
+    def is_qnn_module(self, pass_type: str) -> bool:
+        pass_module = self.get_pass_module_config(pass_type)
+        return pass_module.module_path.startswith("olive.passes.qnn")
+
+    def is_snpe_module(self, pass_type: str) -> bool:
+        pass_module = self.get_pass_module_config(pass_type)
+        return pass_module.module_path.startswith("olive.passes.snpe")

olive/package_config.py

@@ -194,11 +194,7 @@ def default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConf
         return config
 
     @classmethod
-    def validate_config(
-        cls,
-        config: type[BasePassConfig],
-        accelerator_spec: AcceleratorSpec,
-    ) -> bool:
+    def validate_config(cls, config: type[BasePassConfig], accelerator_spec: AcceleratorSpec) -> bool:
         """Validate the input config for the pass."""
         return True
 

olive/passes/olive_pass.py

@@ -138,6 +138,19 @@ def _verify_elem_type(self, elem_type):
                 "for details."
             )
 
+    @classmethod
+    def validate_config(cls, config: type[BasePassConfig], accelerator_spec: AcceleratorSpec) -> bool:
+        if not super().validate_config(config, accelerator_spec):
+            return False
+
+        if config.target_dtype == onnx.TensorProto.FLOAT16 and (
+            accelerator_spec.execution_provider in {"JsExecutionProvider", "WebGpuExecutionProvider"}
+        ):
+            logger.info("Web execution providers don't support fp16.")
+            return False
+
+        return True
+
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:

olive/passes/onnx/io_datatype_converter.py

@@ -266,6 +266,17 @@ class OnnxPeepholeOptimizer(Pass):
     def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]:
         return get_external_data_config()
 
+    @classmethod
+    def validate_config(cls, config: type[BasePassConfig], accelerator_spec: AcceleratorSpec) -> bool:
+        if not super().validate_config(config, accelerator_spec):
+            return False
+
+        if accelerator_spec.execution_provider == "QNNExecutionProvider":
+            logger.info("QNNExecutionProvider doesn't support optimized model.")
+            return False
+
+        return True
+
     def _run_for_config(
         self, model: ONNXModelHandler, config: type[BasePassConfig], output_model_path: str
     ) -> ONNXModelHandler:

olive/passes/onnx/peephole_optimizer.py

@@ -136,11 +136,7 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
         return config
 
     @classmethod
-    def validate_config(
-        cls,
-        config: type[BasePassConfig],
-        accelerator_spec: AcceleratorSpec,
-    ) -> bool:
+    def validate_config(cls, config: type[BasePassConfig], accelerator_spec: AcceleratorSpec) -> bool:
         if not super().validate_config(config, accelerator_spec):
             return False
 
@@ -157,6 +153,10 @@ def validate_config(
             if accelerator_spec.execution_provider == "CPUExecutionProvider":
                 logger.info("CPUExecutionProvider does not support float16 very well, please avoid to use float16.")
                 return False
+
+        if accelerator_spec.execution_provider == "QNNExecutionProvider":
+            logger.info("QNNExecutionProvider doesn't support optimized model.")
+            return False
         if not config.float16 and config.use_gqa:
             logger.info("use_gqa is only supported when float16 is True.")
             return False

olive/passes/onnx/transformer_optimization.py

@@ -2,26 +2,21 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
-from pathlib import Path
 
 import pytest
-import yaml
 
 from olive.auto_optimizer import AutoOptimizer, AutoOptimizerConfig
-from olive.auto_optimizer.template_mapping import get_pass_flows_by_accelerator_ep_precision
-from olive.evaluator.metric import AccuracySubType
-from olive.evaluator.olive_evaluator import OliveEvaluatorConfig
-from olive.hardware import DEFAULT_CPU_ACCELERATOR, DEFAULT_GPU_CUDA_ACCELERATOR, DEFAULT_GPU_TRT_ACCELERATOR
+from olive.constants import Precision
+from olive.hardware import DEFAULT_CPU_ACCELERATOR, DEFAULT_GPU_CUDA_ACCELERATOR
 from olive.model import ModelConfig
-from test.unit_test.utils import get_accuracy_metric, get_glue_huggingface_data_config
 
 # pylint: disable=attribute-defined-outside-init
 
 
 class TestAutoOptimizer:
     @pytest.fixture(autouse=True)
     def setup(self):
-        self.input_model_config = ModelConfig(
+        self.model_config = ModelConfig(
             type="PyTorchModel",
             config={
                 "hf_config": {
@@ -30,116 +25,68 @@ def setup(self):
                 }
             },
         )
-        self.data_configs = [get_glue_huggingface_data_config()]
 
     @pytest.mark.parametrize(
-        ("accelerator_spec", "auto_optimizer_config", "expected_cuda_fp16", "expected_trt_fp16"),
+        ("optimizer_config", "expected_pass_types"),
         [
             (
-                # running on gpu-cuda, enable cuda fp16, disable trt fp16
-                DEFAULT_GPU_CUDA_ACCELERATOR,
-                None,
-                True,
-                False,
+                AutoOptimizerConfig(
+                    precision=Precision.FP16,
+                    accelerator=DEFAULT_CPU_ACCELERATOR,
+                    finetune=False,
+                ),
+                {
+                    "CaptureSplitInfo",
+                    "ExtractAdapters",
+                    "MatMulNBitsToQDQ",
+                    "ModelBuilder",
+                    "OnnxIODataTypeConverter",
+                    "QuaRot",
+                    "SpinQuant",
+                    "SplitModel",
+                },
             ),
             (
-                # running on gpu-trt, disable cuda fp16, enable trt fp16
-                DEFAULT_GPU_TRT_ACCELERATOR,
-                None,
-                False,
-                True,
+                AutoOptimizerConfig(
+                    precision=Precision.FP32,
+                    accelerator=DEFAULT_CPU_ACCELERATOR,
+                    finetune=False,
+                ),
+                {
+                    "CaptureSplitInfo",
+                    "ExtractAdapters",
+                    "MatMulNBitsToQDQ",
+                    "ModelBuilder",
+                    "OnnxIODataTypeConverter",
+                    "QuaRot",
+                    "SpinQuant",
+                    "SplitModel",
+                },
             ),
-        ],
-    )
-    def test_regulate_fp16(self, accelerator_spec, auto_optimizer_config, expected_cuda_fp16, expected_trt_fp16):
-        metrics = [get_accuracy_metric(AccuracySubType.ACCURACY_SCORE, goal_type="max-degradation")]
-        for metric in metrics:
-            metric.data_config = self.data_configs[0]
-        evaluator_config = OliveEvaluatorConfig(metrics=metrics)
-        auto_optimizer = AutoOptimizer(
-            input_model_config=self.input_model_config,
-            evaluator_config=evaluator_config,
-            accelerator_spec=accelerator_spec,
-            auto_optimizer_config=auto_optimizer_config,
-            data_configs=self.data_configs,
-        )
-
-        pass_config, _ = auto_optimizer.suggest()
-        trans_opt_name = "OrtTransformerOptimization_cuda_fp16" if expected_cuda_fp16 else "OrtTransformersOptimization"
-        session_params_opt_name = "OrtSessionParamsTuning_trt_fp16" if expected_trt_fp16 else "OrtSessionParamsTuning"
-        assert pass_config[trans_opt_name]["config"]["float16"] == expected_cuda_fp16
-        assert pass_config[session_params_opt_name]["config"]["enable_cuda_graph"] == expected_cuda_fp16
-        assert pass_config[session_params_opt_name]["config"]["trt_fp16_enable"] == expected_trt_fp16
-
-    @pytest.mark.parametrize(
-        ("metrics_configs", "accelerator_spec", "auto_optimizer_config", "expected_pass_flows"),
-        [
             (
-                [{"args": [AccuracySubType.ACCURACY_SCORE], "kwargs": {"goal_type": "max-degradation"}}],
-                DEFAULT_CPU_ACCELERATOR,
-                None,
-                [
-                    ["OnnxConversion", "OrtTransformersOptimization", "OrtSessionParamsTuning"],
-                    ["OnnxConversion", "OrtTransformersOptimization", "OnnxQuantization", "OrtSessionParamsTuning"],
-                    ["OnnxConversion", "OrtTransformersOptimization", "IncQuantization", "OrtSessionParamsTuning"],
-                    ["OnnxConversion", "OrtTransformersOptimization", "OnnxMatMul4Quantizer", "OrtSessionParamsTuning"],
-                    ["ModelBuilder_fp32", "OrtSessionParamsTuning"],
-                    ["ModelBuilder_int4", "OrtSessionParamsTuning"],
-                    ["ModelBuilder_int8", "OrtSessionParamsTuning"],
-                    ["ModelBuilder_fp16", "OrtSessionParamsTuning"],
-                ],
-            ),
-            (
-                # cannot tolerate accuracy drop, then skip quantization
-                [
-                    {
-                        "args": [AccuracySubType.ACCURACY_SCORE],
-                        "kwargs": {"goal_type": "max-degradation", "goal_value": 0},
-                    }
-                ],
-                DEFAULT_CPU_ACCELERATOR,
-                AutoOptimizerConfig(precisions=["fp32"]),
-                [
-                    ["OnnxConversion", "OrtTransformersOptimization", "OrtSessionParamsTuning"],
-                    ["ModelBuilder_fp32", "OrtSessionParamsTuning"],
-                ],
-            ),
-            (
-                # running on gpu-cuda, skip quantization
-                [{"args": [AccuracySubType.ACCURACY_SCORE], "kwargs": {"goal_type": "max-degradation"}}],
-                DEFAULT_GPU_CUDA_ACCELERATOR,
-                AutoOptimizerConfig(precisions=["fp16"], excluded_passes=["ModelBuilder"]),
-                [
-                    ["OnnxConversion", "OrtTransformerOptimization_cuda_fp16", "OrtSessionParamsTuning"],
-                    ["OnnxConversion", "OrtTransformersOptimization", "OrtMixedPrecision", "OrtSessionParamsTuning"],
-                ],
+                AutoOptimizerConfig(
+                    precision=Precision.FP16,
+                    accelerator=DEFAULT_GPU_CUDA_ACCELERATOR,
+                    finetune=False,
+                    excluded_passes=["ModelBuilder"]
+                ),
+                {
+                    "CaptureSplitInfo",
+                    "ExtractAdapters",
+                    "MatMulNBitsToQDQ",
+                    "OnnxIODataTypeConverter",
+                    "QuaRot",
+                    "SpinQuant",
+                    "SplitModel",
+                },
             ),
         ],
     )
-    def test_regulate_pass(self, metrics_configs, accelerator_spec, auto_optimizer_config, expected_pass_flows):
-        metrics = [get_accuracy_metric(*mc["args"], **mc["kwargs"]) for mc in metrics_configs]
-        for metric in metrics:
-            metric.data_config = self.data_configs[0]
-        evaluator_config = OliveEvaluatorConfig(metrics=metrics)
-        auto_optimizer = AutoOptimizer(
-            input_model_config=self.input_model_config,
-            evaluator_config=evaluator_config,
-            accelerator_spec=accelerator_spec,
-            auto_optimizer_config=auto_optimizer_config,
-            data_configs=self.data_configs,
-        )
-
-        pass_config, pass_flows = auto_optimizer.suggest()
-        assert pass_config, "Expect pass_config to be populated by auto optimizer"
-        assert sorted(pass_flows) == sorted(expected_pass_flows)
-
-    def test_pass_flows_generation_opt_level_0(self):
-        pass_flows_map = Path(__file__).parent / "mock_data" / "available_pass_flows.yaml"
-        with pass_flows_map.open() as f:
-            pass_flows_map = yaml.safe_load(f)["mapping"]
+    def test_generate_run_passes_configs(self, optimizer_config, expected_pass_types):
+        auto_optimizer = AutoOptimizer(model_config=self.model_config, optimizer_config=optimizer_config)
+        pass_configs = auto_optimizer.generate_run_passes_configs()
+        assert pass_configs, "Expect pass_configs to be populated by auto optimizer"
 
-        for k, pf in pass_flows_map.items():
-            k_list = k.split("_")
-            accelerator, ep, precision = k_list[0], k_list[1], k_list[2]
-            rls_pf = get_pass_flows_by_accelerator_ep_precision(0, accelerator, ep, precision)
-            assert sorted(rls_pf) == sorted(pf)
+        actual_pass_types = {pc.type for _, pcs in pass_configs.items() for pc in pcs}
+        expected_pass_types = {pt.lower() for pt in expected_pass_types}
+        assert sorted(actual_pass_types) == sorted(expected_pass_types)