Add mobius ep export support

olive/cli/capture_onnx.py

@@ -176,6 +176,12 @@ def register_subcommand(parser: ArgumentParser):
             required=False,
             help="Extra key-value pairs options to pass to the model builder. e.g., 'int4_is_symmetric=true,int4_op_types_to_quantize=MatMul/Gemm'.",
         )
+        mb_group.add_argument(
+            "--execution_provider",
+            type=str,
+            required=False,
+            help=("Mobius execution provider profile, such as 'openvino'. Only used with --use_mobius_builder."),
+        )
 
         sub_parser.add_argument(
             "--use_ort_genai",
@@ -247,6 +253,10 @@ def _get_run_config(self, tempdir: str) -> dict:
             del config["passes"]["c"]
             del config["passes"]["m"]
             to_replace.append((("passes", "b", "precision"), self.args.precision))
+            if self.args.execution_provider:
+                to_replace.append((("passes", "b", "execution_provider"), self.args.execution_provider))
+        elif self.args.execution_provider:
+            raise ValueError("--execution_provider is only supported with --use_mobius_builder.")
         elif is_diffusers_model:
             del config["passes"]["m"]
             del config["passes"]["b"]

olive/model/handler/openvino.py

@@ -12,6 +12,20 @@
 from olive.resource_path import OLIVE_RESOURCE_ANNOTATIONS, create_resource_path
 
 
+def create_openvino_core():
+    """Create an OpenVINO Core with extensions needed to deserialize OpenVINO IR."""
+    try:
+        import openvino as ov
+    except ImportError:
+        raise ImportError("Please install olive-ai[openvino] to use OpenVINO model") from None
+
+    core = ov.Core()
+    group_query_attention_extension = getattr(getattr(ov, "op", None), "_GroupQueryAttentionExtension", None)
+    if group_query_attention_extension:
+        core.add_extension(group_query_attention_extension())
+    return core
+
+
 @model_handler_registry("OpenVINOModel")
 class OpenVINOModelHandler(OliveModelHandler):
     """OpenVINO model handler.
@@ -54,11 +68,7 @@ def model_config(self) -> dict[str, str]:
         }
 
     def load_model(self, rank: int = None, cache_model: bool = True):
-        try:
-            import openvino as ov
-        except ImportError:
-            raise ImportError("Please install olive-ai[openvino] to use OpenVINO model") from None
-        core = ov.Core()
+        core = create_openvino_core()
         return core.read_model(self.model_config["model"])
 
     @property
@@ -73,11 +83,7 @@ def prepare_session(
         execution_providers: Union[str, list[str]] = None,
         rank: Optional[int] = None,
     ):
-        try:
-            import openvino as ov
-        except ImportError:
-            raise ImportError("Please install olive-ai[openvino] to use OpenVINO model") from None
-        core = ov.Core()
+        core = create_openvino_core()
         if inference_settings and inference_settings.get("device_name"):
             device = inference_settings["device_name"]
         elif device == Device.INTEL_MYRIAD:

olive/passes/onnx/mobius_model_builder.py

@@ -111,6 +111,15 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                     "quantization pass (e.g. OnnxMatMulNBits) after this pass."
                 ),
             ),
+            "execution_provider": PassConfigParam(
+                type_=str,
+                required=False,
+                default_value=None,
+                description=(
+                    "Mobius execution provider profile to use, such as 'openvino'. "
+                    "When omitted, the profile is inferred from the Olive accelerator."
+                ),
+            ),
             "components_to_export": PassConfigParam(
                 type_=list[str],
                 required=False,
@@ -143,16 +152,19 @@ def _run_for_config(
         if not isinstance(model, HfModelHandler):
             raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.")
 
-        # Map Olive EP to mobius EP. If unsupported/unknown, fall back to mobius default EP.
-        requested_ep = self.accelerator_spec.execution_provider
-        ep_str: str = EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT)
-        if ep_str == self.MobiusEP.DEFAULT:
-            logger.warning(
-                "MobiusBuilder: execution provider '%s' on accelerator '%s' is not explicitly supported; "
-                "falling back to mobius default EP.",
-                requested_ep,
-                self.accelerator_spec.accelerator_type,
-            )
+        if config.execution_provider:
+            ep_str = config.execution_provider
+        else:
+            # Map Olive EP to mobius EP. If unsupported/unknown, fall back to mobius default EP.
+            requested_ep = self.accelerator_spec.execution_provider
+            ep_str = EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT)
+            if ep_str == self.MobiusEP.DEFAULT:
+                logger.warning(
+                    "MobiusBuilder: execution provider '%s' on accelerator '%s' is not explicitly supported; "
+                    "falling back to mobius default EP.",
+                    requested_ep,
+                    self.accelerator_spec.accelerator_type,
+                )
 
         dtype_str: str = _PRECISION_TO_DTYPE.get(config.precision, "f32")
         model_id: str = model.model_name_or_path

olive/passes/openvino/compression.py

@@ -13,6 +13,7 @@
 from olive.data.config import DataConfig
 from olive.hardware.accelerator import AcceleratorSpec, Device
 from olive.model.handler import CompositeModelHandler, HfModelHandler, ONNXModelHandler, OpenVINOModelHandler
+from olive.model.handler.openvino import create_openvino_core
 from olive.passes import Pass
 from olive.passes.openvino.ov_utils import (
     IgnoreScopeTypeEnum,
@@ -802,7 +803,7 @@ def _run_openvino_pass(
             raise ImportError("Please install openvino to use OpenVINO weight compression") from None
 
         # load the OpenVINO model
-        core = ov.Core()
+        core = create_openvino_core()
         model_config = model.model_config
         loaded_model = core.read_model(model_config["model"])
 

olive/passes/openvino/encapsulation.py

@@ -13,6 +13,7 @@
 from olive.common.utils import hardlink_copy_dir, hardlink_copy_file
 from olive.hardware.accelerator import AcceleratorSpec, Device
 from olive.model import ONNXModelHandler, OpenVINOModelHandler
+from olive.model.handler.openvino import create_openvino_core
 from olive.passes import Pass
 from olive.passes.openvino.ov_utils import create_genai_config
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
@@ -141,7 +142,7 @@ def _run_single_target(
         else:
             ov_version = ov.get_version()
 
-        core = ov.Core()
+        core = create_openvino_core()
         model_name_path = Path(model.model_path) / (f"{model_name}.xml")
         weight_name_path = Path(model.model_path) / (f"{model_name}.bin")
 

olive/passes/openvino/io_update.py

@@ -9,6 +9,7 @@
 from olive.common.utils import hardlink_copy_dir, hardlink_copy_file
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import OpenVINOModelHandler
+from olive.model.handler.openvino import create_openvino_core
 from olive.passes import Pass
 from olive.passes.pass_config import BasePassConfig, PassConfigParam, get_user_script_data_config
 
@@ -103,7 +104,7 @@ def _run_pass(
 
         model_name = model.model_config["model_name"]
 
-        core = ov.Core()
+        core = create_openvino_core()
         model_name_path = Path(model.model_path) / (f"{model_name}.xml")
         weight_name_path = Path(model.model_path) / (f"{model_name}.bin")
 

test/cli/test_cli.py

@@ -555,6 +555,49 @@ def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort
     assert mock_run.call_count == 1
 
 
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+def test_capture_onnx_command_use_mobius_builder_execution_provider(_, mock_run, tmp_path):
+    cli_main(
+        [
+            "capture-onnx-graph",
+            "-m",
+            "dummy-model-id",
+            "-o",
+            str(tmp_path / "output_dir"),
+            "--use_mobius_builder",
+            "--precision",
+            "fp32",
+            "--execution_provider",
+            "openvino",
+        ]
+    )
+
+    config = mock_run.call_args.args[0]
+    assert config["passes"]["b"]["execution_provider"] == "openvino"
+    assert config["systems"]["local_system"]["accelerators"][0] == {
+        "device": "cpu",
+        "execution_providers": ["CPUExecutionProvider"],
+    }
+
+
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+def test_capture_onnx_command_execution_provider_requires_mobius_builder(_, __, tmp_path):
+    with pytest.raises(ValueError, match="only supported with --use_mobius_builder"):
+        cli_main(
+            [
+                "capture-onnx-graph",
+                "-m",
+                "dummy-model-id",
+                "-o",
+                str(tmp_path / "output_dir"),
+                "--execution_provider",
+                "openvino",
+            ]
+        )
+
+
 @patch("olive.workflows.run")
 @patch("huggingface_hub.repo_exists", return_value=True)
 def test_capture_onnx_command_use_mobius_builder_rejects_int4(_, __, tmp_path):

test/model/test_openvino_model.py

@@ -0,0 +1,36 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import sys
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from olive.model.handler.openvino import create_openvino_core
+
+
+def test_create_openvino_core_registers_group_query_attention_extension():
+    core = MagicMock()
+    extension = MagicMock()
+    openvino = SimpleNamespace(
+        Core=MagicMock(return_value=core),
+        op=SimpleNamespace(_GroupQueryAttentionExtension=extension),
+    )
+
+    with patch.dict(sys.modules, {"openvino": openvino}):
+        result = create_openvino_core()
+
+    assert result is core
+    extension.assert_called_once_with()
+    core.add_extension.assert_called_once_with(extension.return_value)
+
+
+def test_create_openvino_core_without_group_query_attention_extension():
+    core = MagicMock()
+    openvino = SimpleNamespace(Core=MagicMock(return_value=core), op=SimpleNamespace())
+
+    with patch.dict(sys.modules, {"openvino": openvino}):
+        result = create_openvino_core()
+
+    assert result is core
+    core.add_extension.assert_not_called()

test/passes/onnx/test_mobius_model_builder.py

@@ -181,13 +181,14 @@ def __exit__(self, *args):
 
 
 def test_default_config_params():
-    """MobiusBuilder must declare precision, and must not declare execution_provider or trust_remote_code."""
+    """MobiusBuilder must declare precision and an optional execution provider override."""
     accelerator_spec = AcceleratorSpec(
         accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
     )
     config = MobiusBuilder._default_config(accelerator_spec)  # pylint: disable=protected-access
     assert "precision" in config
-    assert "execution_provider" not in config
+    assert config["execution_provider"].default_value is None
+    assert config["execution_provider"].required is False
     assert "trust_remote_code" not in config
 
 
@@ -241,6 +242,26 @@ def test_single_component_returns_onnx_handler(tmp_path):
     assert call_kwargs["dtype"] == "f32"
 
 
+def test_execution_provider_config_overrides_accelerator(tmp_path):
+    """An explicit Mobius profile overrides the profile inferred from the Olive accelerator."""
+    out = tmp_path / "out"
+    pkg = _fake_pkg(["model"], out)
+    accelerator_spec = AcceleratorSpec(
+        accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+    )
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp32", "execution_provider": "openvino"},
+        disable_search=True,
+        accelerator_spec=accelerator_spec,
+    )
+
+    with _patch_build(pkg) as mock_build:
+        p.run(_make_hf_model("org/model"), out)
+
+    assert mock_build.call_args.kwargs["execution_provider"] == "openvino"
+
+
 def test_model_onnx_exists_after_run(tmp_path):
     """The saved model.onnx file must exist on disk."""
     out = tmp_path / "out"