Add Mobius exporter support to olive optimize

docs/source/reference/python_api.md

@@ -39,7 +39,7 @@ Optimize the input model with comprehensive pass scheduling.
 - `act_precision` (str, optional): Activation precision for quantization.
 - `num_split` (int, optional): Number of splits for model splitting.
 - `memory` (int, optional): Available device memory in MB.
-- `exporter` (str, optional): Exporter to use ("model_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter").
+- `exporter` (str, optional): Exporter to use ("model_builder", "mobius", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"). Mobius requires `mobius-onnx` and supports fp32, fp16, and bf16 export; apply quantization as a subsequent pass.
 - `dim_param` (str, optional): Dynamic parameter names for dynamic to fixed shape conversion.
 - `dim_value` (str, optional): Fixed dimension values for dynamic to fixed shape conversion.
 - `use_qdq_format` (bool): Use QDQ format for quantization. Defaults to `False`.

mcp/src/olive_mcp/packages.py

@@ -126,6 +126,8 @@ def _resolve_packages(command: str, provider: str | None = None, **kwargs) -> li
             # ModelBuilder pass requires onnxruntime-genai (undeclared in olive_config.json)
             genai_pkg = PROVIDER_TO_GENAI.get(provider or "CPUExecutionProvider", "onnxruntime-genai")
             extra_packages.append(genai_pkg)
+        elif exporter == "mobius":
+            extra_packages.append("mobius-onnx")
         elif exporter == "optimum_exporter":
             extras.add("optimum")
 

mcp/src/olive_mcp/tools.py

@@ -258,7 +258,8 @@ async def optimize(
         device: Target device - "cpu", "gpu", or "npu". Auto-detected from provider if omitted.
         precision: Target precision - "fp32", "fp16", "int4", "int8", etc.
         act_precision: Activation precision for quantization (optional).
-        exporter: Model exporter - "model_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter".
+        exporter: Model exporter - "model_builder", "mobius", "dynamo_exporter", "torchscript_exporter",
+            or "optimum_exporter". Mobius requires mobius-onnx and supports fp32, fp16, and bf16 export.
         use_qdq_format: Use QDQ format for quantization instead of QOperator.
         num_split: Number of splits for model splitting.
         memory: Available device memory in MB.

olive/cli/optimize.py

@@ -113,7 +113,7 @@ def register_subcommand(parser: ArgumentParser):
         sub_parser.add_argument(
             "--exporter",
             type=str,
-            choices=["model_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"],
+            choices=["model_builder", "mobius", "dynamo_exporter", "torchscript_exporter", "optimum_exporter"],
             help="Exporter to use for model conversion (optional).",
         )
 
@@ -198,6 +198,7 @@ def __init__(self, parser: ArgumentParser, args: Namespace, unknown_args: Option
         self.enable_gptq = False
         self.enable_capture_split_info = False
         self.enable_model_builder = False
+        self.enable_mobius_builder = False
         self.enable_onnx_conversion = False
         self.enable_optimum_openvino_conversion = False
         self.enable_dynamic_to_fixed_shape = False
@@ -259,6 +260,9 @@ def _validate_arguments(self):
         if self.args.modality not in ["text"]:
             raise ValueError(f"Unsupported modality: {self.args.modality}. Only 'text' is supported for optimization.")
 
+        if self.args.exporter == "mobius" and self.args.precision not in ("fp32", "fp16", "bf16"):
+            raise ValueError("MobiusBuilder supports fp32, fp16, and bf16. Export first, then apply quantization.")
+
         if self.args.provider == ExecutionProvider.CPUExecutionProvider and self.args.device in ["gpu", "npu"]:
             raise ValueError(
                 f"Invalid combination of provider {self.args.provider} and device {self.args.device}. "
@@ -336,6 +340,10 @@ def _build_passes_config(self) -> dict[str, Any]:
         if self.enable_model_builder:
             passes_config["model_builder"] = self._get_model_builder_pass_config()
 
+        self.enable_mobius_builder = self._enable_mobius_builder_pass()
+        if self.enable_mobius_builder:
+            passes_config["mobius_builder"] = self._get_mobius_builder_pass_config()
+
         self.enable_onnx_conversion = self._enable_onnx_conversion_pass()
         if self.enable_onnx_conversion:
             passes_config["onnx_conversion"] = self._get_onnx_conversion_pass_config()
@@ -502,6 +510,14 @@ def _get_model_builder_pass_config(self) -> dict[str, Any]:
 
         return config
 
+    def _enable_mobius_builder_pass(self) -> bool:
+        """Return true if condition to add MobiusBuilder pass is met."""
+        return self.is_hf_model and self.args.exporter == "mobius"
+
+    def _get_mobius_builder_pass_config(self) -> dict[str, Any]:
+        """Return pass dictionary for MobiusBuilder pass."""
+        return {"type": "MobiusBuilder", "precision": Precision(self.args.precision).value}
+
     def _enable_onnx_conversion_pass(self) -> bool:
         """Return true if condition to add OnnxConversion pass is met."""
         provider = ExecutionProvider(self.args.provider)
@@ -518,7 +534,11 @@ def _get_onnx_conversion_pass_config(self) -> dict[str, Any]:
     def _enable_optimum_openvino_conversion_pass(self) -> bool:
         """Return true if condition to add OptimumOpenvinoConversion pass is met."""
         provider = ExecutionProvider(self.args.provider)
-        return self.is_hf_model and provider == ExecutionProvider.OpenVINOExecutionProvider
+        return (
+            self.is_hf_model
+            and provider == ExecutionProvider.OpenVINOExecutionProvider
+            and self.args.exporter != "mobius"
+        )
 
     def _get_optimum_openvino_conversion_pass_config(self) -> dict[str, Any]:
         """Return pass dictionary for OptimumOpenvinoConversion pass."""
@@ -564,7 +584,7 @@ def _get_openvino_io_update_pass_config(self) -> dict[str, Any]:
 
     def _enable_onnx_peephole_optimizer_pass(self) -> bool:
         """Return true if condition to add OnnxPeepholeOptimizer pass is met."""
-        return not self.is_hf_model or self.args.exporter != "model_builder"
+        return not self.is_hf_model or self.args.exporter not in ("model_builder", "mobius")
 
     def _get_onnx_peephole_optimizer_pass_config(self) -> dict[str, Any]:
         """Return pass dictionary for OnnxPeepholeOptimizer pass."""
@@ -634,7 +654,7 @@ def _get_onnx_blockwise_rtn_quantization_pass_config(self) -> dict[str, Any]:
     def _enable_onnx_float_to_float16_pass(self) -> bool:
         """Return true if condition to add OnnxFloatToFloat16 pass is met."""
         precision = Precision(self.args.precision)
-        return precision == Precision.FP16 and not self.enable_model_builder
+        return precision == Precision.FP16 and not (self.enable_model_builder or self.enable_mobius_builder)
 
     def _get_onnx_float_to_float16_pass_config(self) -> dict[str, Any]:
         """Return pass dictionary for OnnxFloatToFloat16 pass."""

olive/passes/onnx/mobius_model_builder.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar
 
+from olive.common.hf.utils import has_test_model_weights, is_test_model_dir
 from olive.common.utils import StrEnumBase
 from olive.constants import Precision
 from olive.hardware.constants import EXECUTION_PROVIDER_TO_MOBIUS_EP, ExecutionProvider
@@ -177,6 +178,16 @@ def _run_for_config(
         revision: str | None = load_kwargs.get("revision")
         trust_remote_code: bool = load_kwargs.get("trust_remote_code", False)
 
+        if model.test_model_config:
+            if not model.test_model_path:
+                raise ValueError("MobiusBuilder requires test_model_path in test mode.")
+
+            if not is_test_model_dir(model.test_model_path) or not has_test_model_weights(model.test_model_path):
+                model.load_model(cache_model=False)
+
+            model_id = str(Path(model.test_model_path).resolve())
+            revision = None
+
         logger.info(
             "MobiusBuilder: building '%s' (ep=%s, dtype=%s)",
             model_id,

test/cli/test_cli.py

@@ -863,6 +863,13 @@ def test_optimize_cli_pass_list(mock_repo_exists, mock_run, tmp_path):
             None,
             "CUDAExecutionProvider",
         ],
+        [
+            "optimize",
+            "--precision fp16 --exporter mobius --provider CUDAExecutionProvider",
+            "MobiusBuilder",
+            None,
+            "CUDAExecutionProvider",
+        ],
         [
             "optimize",
             (
@@ -915,6 +922,55 @@ def test_optimize_cli_pass_list(mock_repo_exists, mock_run, tmp_path):
         )
 
 
+@pytest.mark.parametrize("precision", ["fp32", "fp16", "bf16"])
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+def test_optimize_cli_mobius_exporter_supported_precisions(_, mock_run, precision, tmp_path):
+    output_dir = tmp_path / precision
+
+    cli_main(
+        [
+            "optimize",
+            "-m",
+            "dummy_model",
+            "--exporter",
+            "mobius",
+            "--precision",
+            precision,
+            "--dry_run",
+            "-o",
+            str(output_dir),
+        ]
+    )
+
+    config = json.loads((output_dir / "config.json").read_text())
+    assert list(config["passes"]) == ["mobius_builder"]
+    assert config["passes"]["mobius_builder"] == {"type": "MobiusBuilder", "precision": precision}
+    mock_run.assert_not_called()
+
+
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+def test_optimize_cli_mobius_exporter_rejects_quantized_precision(_, mock_run, tmp_path):
+    with pytest.raises(ValueError, match="MobiusBuilder supports fp32, fp16, and bf16"):
+        cli_main(
+            [
+                "optimize",
+                "-m",
+                "dummy_model",
+                "--exporter",
+                "mobius",
+                "--precision",
+                "int4",
+                "--dry_run",
+                "-o",
+                str(tmp_path / "output"),
+            ]
+        )
+
+    mock_run.assert_not_called()
+
+
 @patch("olive.workflows.run")
 @patch("huggingface_hub.repo_exists", return_value=True)
 def test_benchmark_command_hfmodel(_, mock_run, tmp_path):

test/mcp/test_packages.py

@@ -0,0 +1,53 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import importlib
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+# pylint: disable=protected-access
+
+
+def _load_packages_module(monkeypatch):
+    constants = types.ModuleType("olive_mcp.constants")
+    constants.PROVIDER_TO_EXTRAS = {}
+    constants.PROVIDER_TO_GENAI = {"CPUExecutionProvider": "onnxruntime-genai"}
+    constants.Command = types.SimpleNamespace(
+        OPTIMIZE="optimize",
+        QUANTIZE="quantize",
+        FINETUNE="finetune",
+        CAPTURE_ONNX_GRAPH="capture-onnx-graph",
+        BENCHMARK="benchmark",
+        DIFFUSION_LORA="diffusion-lora",
+    )
+    monkeypatch.setitem(sys.modules, "olive_mcp.constants", constants)
+
+    packages_path = Path(__file__).parents[2] / "mcp" / "src" / "olive_mcp" / "packages.py"
+    spec = importlib.util.spec_from_file_location("olive_mcp_packages_test", packages_path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Unable to load MCP packages from {packages_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_resolve_optimize_packages_includes_mobius(monkeypatch):
+    packages = _load_packages_module(monkeypatch)
+
+    resolved = packages._resolve_packages(packages.Command.OPTIMIZE, exporter="mobius")
+
+    assert "olive-ai[cpu]" in resolved
+    assert "mobius-onnx" in resolved
+
+
+def test_resolve_optimize_packages_keeps_default_model_builder(monkeypatch):
+    packages = _load_packages_module(monkeypatch)
+
+    resolved = packages._resolve_packages(packages.Command.OPTIMIZE)
+
+    assert "olive-ai[cpu]" in resolved
+    assert "onnxruntime-genai" in resolved
+    assert "mobius-onnx" not in resolved

test/passes/onnx/test_mobius_model_builder.py

@@ -66,10 +66,20 @@ def mock_hf_config():
         yield
 
 
-def _make_hf_model(model_path: str, load_kwargs: dict | None = None, task: str | None = None) -> HfModelHandler:
+def _make_hf_model(
+    model_path: str,
+    load_kwargs: dict | None = None,
+    task: str | None = None,
+    test_model_config: dict | None = None,
+    test_model_path: str | None = None,
+) -> HfModelHandler:
     model_kwargs = {"model_path": model_path}
     if task is not None:
         model_kwargs["task"] = task
+    if test_model_config is not None:
+        model_kwargs["test_model_config"] = test_model_config
+    if test_model_path is not None:
+        model_kwargs["test_model_path"] = test_model_path
     model = HfModelHandler(**model_kwargs)
     if load_kwargs:
         # Patch get_load_kwargs on the instance to return the given kwargs.
@@ -260,6 +270,80 @@ def test_single_component_returns_onnx_handler(tmp_path):
     assert call_kwargs["dtype"] == "f32"
 
 
+def test_mobius_builder_uses_huggingface_id_for_normal_export(tmp_path):
+    model_id = "org/full-model"
+    pkg = _fake_pkg(["model"], tmp_path / "out")
+
+    with _patch_build(pkg) as mock_build:
+        _make_pass().run(_make_hf_model(model_id), tmp_path / "out")
+
+    assert mock_build.call_args.args[0] == model_id
+
+
+def test_mobius_builder_uses_saved_test_model_path(tmp_path):
+    test_model_path = tmp_path / "saved_test_model"
+    test_model_path.mkdir()
+    model = _make_hf_model(
+        "org/full-model",
+        {"revision": "abc123", "trust_remote_code": True},
+        test_model_config={"hidden_layers": 2},
+        test_model_path=str(test_model_path),
+    )
+    pkg = _fake_pkg(["model"], tmp_path / "out")
+
+    with (
+        _patch_build(pkg) as mock_build,
+        patch("olive.passes.onnx.mobius_model_builder.is_test_model_dir", return_value=True),
+        patch("olive.passes.onnx.mobius_model_builder.has_test_model_weights", return_value=True),
+        patch.object(model, "load_model") as mock_load_model,
+        patch.object(MobiusBuilder, "_write_genai_config", return_value={}) as mock_write_genai_config,
+    ):
+        _make_pass().run(model, tmp_path / "out")
+
+    resolved_test_model_path = str(test_model_path.resolve())
+    assert mock_build.call_args.args[0] == resolved_test_model_path
+    assert mock_build.call_args.kwargs["revision"] is None
+    assert mock_build.call_args.kwargs["trust_remote_code"] is True
+    mock_write_genai_config.assert_called_once_with(
+        pkg,
+        str(tmp_path / "out"),
+        resolved_test_model_path,
+        "cpu",
+        revision=None,
+        trust_remote_code=True,
+    )
+    mock_load_model.assert_not_called()
+
+
+def test_mobius_builder_materializes_missing_test_weights(tmp_path):
+    test_model_path = tmp_path / "config_only_test_model"
+    test_model_path.mkdir()
+    model = _make_hf_model(
+        "org/full-model",
+        test_model_config={"hidden_layers": 2},
+        test_model_path=str(test_model_path),
+    )
+    pkg = _fake_pkg(["model"], tmp_path / "out")
+
+    with (
+        _patch_build(pkg) as mock_build,
+        patch("olive.passes.onnx.mobius_model_builder.is_test_model_dir", return_value=True),
+        patch("olive.passes.onnx.mobius_model_builder.has_test_model_weights", return_value=False),
+        patch.object(model, "load_model") as mock_load_model,
+    ):
+        _make_pass().run(model, tmp_path / "out")
+
+    mock_load_model.assert_called_once_with(cache_model=False)
+    assert mock_build.call_args.args[0] == str(test_model_path.resolve())
+
+
+def test_mobius_builder_requires_test_model_path(tmp_path):
+    model = _make_hf_model("org/full-model", test_model_config={"hidden_layers": 2})
+
+    with pytest.raises(ValueError, match="requires test_model_path in test mode"):
+        _make_pass().run(model, tmp_path / "out")
+
+
 def test_text_only_default_omits_mobius_build_kwarg(tmp_path):
     """The default remains compatible with Mobius versions that predate text_only."""
     out = tmp_path / "out"