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/discrepancy_check.py

@@ -50,35 +50,49 @@ def _reconcile_genai_speech_output_names(genai_config: dict, actual_outputs: dic
     return _genai_speech_worker.reconcile_output_names(genai_config, actual_outputs)
 
 
+_SYMBOLIC_DIMENSION_DEFAULTS = {
+    "batch": 1,
+    "batch_size": 1,
+    "sequence": 8,
+    "sequence_len": 8,
+    "sequence_length": 8,
+    "past_sequence_len": 0,
+    "past_sequence_length": 0,
+    "total_sequence_length": 8,
+    "past_seq_len + seq_len": 8,
+}
+
+
+def _normalize_symbolic_dimension(dimension):
+    return dimension.rsplit(".", 1)[-1]
+
+
 def _infer_shape(dynamic_shape, known_values=None):
     # Use an empty past-KV cache (past_sequence_length=0) so the discrepancy check is a clean
     # prefill comparison.  The dummy dataloader passes ``past_key_values.<i>.key/value`` tensors,
     # but HuggingFace ``forward`` does not accept those dotted names as keyword arguments and
     # silently drops them, so the reference model would run without a cache while the ONNX model
     # would consume a (bogus, all-ones) cache -- producing a large, meaningless discrepancy.
     # Keeping the past length at 0 makes both models perform the same prefill over ``input_ids``.
-    default_values = {
-        "batch_size": 1,
-        "past_sequence_length": 0,
-        "sequence_length": 8,
-        "total_sequence_length": 8,
-    }
+    dimension_defaults = dict(_SYMBOLIC_DIMENSION_DEFAULTS)
     if known_values:
         # Shapes mix symbolic names and concrete ints, so only keep the symbolic entries;
         # otherwise the error message below would compare ints against strings.
-        default_values.update({key: value for key, value in known_values.items() if isinstance(key, str)})
+        dimension_defaults.update({key: value for key, value in known_values.items() if isinstance(key, str)})
     inferred_shape = []
-    for dim in dynamic_shape:
-        if isinstance(dim, int):
-            inferred_shape.append(dim)
+    for dimension in dynamic_shape:
+        if isinstance(dimension, int):
+            inferred_shape.append(dimension)
             continue
-        if dim not in default_values:
+        normalized_dimension = _normalize_symbolic_dimension(dimension)
+        dimension_value = dimension_defaults.get(dimension, dimension_defaults.get(normalized_dimension))
+        if dimension_value is None:
             raise KeyError(
-                f"Unsupported symbolic dimension '{dim}' in shape {dynamic_shape}. "
-                f"Known symbols are: {sorted(default_values)}. "
+                f"Unsupported symbolic dimension '{dimension}' in shape {dynamic_shape}. "
+                f"Known symbols are: {sorted(dimension_defaults)}. "
                 "Update OnnxDiscrepancyCheck to handle this new case."
             )
-        inferred_shape.append(default_values[dim])
+        inferred_shape.append(dimension_value)
     return tuple(inferred_shape)
 
 

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_discrepancy_check.py

@@ -62,6 +62,23 @@ def test_infer_shape_resolves_kv_cache_dim_from_known_values():
     assert inferred == (1, 8, 0, 16)
 
 
+def test_infer_shape_supports_mobius_namespaced_dimensions():
+    inferred = _infer_shape(
+        [
+            "component.model.batch",
+            "component.model.sequence_len",
+            "component.model.past_sequence_len",
+            "component.model.past_seq_len + seq_len",
+        ]
+    )
+    assert inferred == (1, 8, 0, 8)
+
+
+def test_infer_shape_prefers_exact_known_namespaced_value():
+    inferred = _infer_shape(["component.model.sequence_len"], {"component.model.sequence_len": 16})
+    assert inferred == (16,)
+
+
 def test_infer_shape_error_message_handles_mixed_known_symbol_keys():
     with pytest.raises(KeyError, match="Unsupported symbolic dimension 'mystery_dim'"):
         _infer_shape(["batch_size", "mystery_dim"], {"kv_cache_dim": 16, 8: 8})