Enable MobiusBuilder in 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_builder", "dynamo_exporter", "torchscript_exporter", "optimum_exporter").
 - `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`.

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_builder", "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
@@ -256,6 +257,16 @@ def _validate_arguments(self):
         if self.args.exporter is None and self.args.modality == "text":
             self.args.exporter = "model_builder"
 
+        if self.args.exporter == "mobius_builder" and self.args.precision not in (
+            Precision.FP32,
+            Precision.FP16,
+            Precision.BF16,
+        ):
+            raise ValueError(
+                f"MobiusBuilder supports precisions fp32/fp16/bf16; got '{self.args.precision}'. "
+                "For INT4, capture in fp32/fp16/bf16 first and run a quantization pass afterwards."
+            )
+
         if self.args.modality not in ["text"]:
             raise ValueError(f"Unsupported modality: {self.args.modality}. Only 'text' is supported for optimization.")
 
@@ -336,6 +347,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 +517,19 @@ 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."""
+        provider = ExecutionProvider(self.args.provider)
+        return (
+            self.is_hf_model
+            and provider != ExecutionProvider.OpenVINOExecutionProvider
+            and self.args.exporter == "mobius_builder"
+        )
+
+    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)
@@ -564,7 +592,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_builder"]
 
     def _get_onnx_peephole_optimizer_pass_config(self) -> dict[str, Any]:
         """Return pass dictionary for OnnxPeepholeOptimizer pass."""
@@ -634,7 +662,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."""

test/cli/test_cli.py

@@ -840,6 +840,13 @@ def test_optimize_cli_pass_list(mock_repo_exists, mock_run, tmp_path):
             None,
             "CUDAExecutionProvider",
         ],
+        [
+            "optimize",
+            "--precision fp16 --provider CUDAExecutionProvider --exporter mobius_builder",
+            "MobiusBuilder",
+            None,
+            "CUDAExecutionProvider",
+        ],
         [
             "optimize",
             (