Support Mobius extra options in capture-onnx-graph

docs/source/reference/cli.rst

@@ -47,6 +47,23 @@ Capture ONNX graph using PyTorch Exporter or Model Builder from the Huggingface
     :prog: olive
     :path: capture-onnx-graph
 
+With ``--use_mobius_builder``, ``--extra_mb_options`` forwards comma-separated
+``key=value`` settings to ``mobius.build(extra_options=...)``. Values are parsed
+as JSON scalars or strings; ``null`` becomes Python ``None``. For example:
+
+.. code-block:: powershell
+
+    olive capture-onnx-graph `
+        --model_name_or_path gemma4_quantized_hf `
+        --use_mobius_builder --precision fp32 `
+        --extra_mb_options "vision_patch_capacity=2520,vision_batch_size=1" `
+        --output_path gemma4_fixed_capacity_onnx
+
+Supported options depend on the installed Mobius version. This example enables
+fixed-capacity Gemma4 vision; omitting the options preserves the existing export
+behavior. For ``--use_model_builder``, ``--extra_mb_options`` retains its existing
+Model Builder parsing and validation.
+
 Run Pass
 ========
 

olive/cli/capture_onnx.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import argparse
+import json
 from argparse import ArgumentParser
 from copy import deepcopy
 
@@ -30,6 +31,23 @@ def parse_dim_dict(s):
         raise argparse.ArgumentTypeError("Format must be key=value,... with positive integers as values") from exc
 
 
+def parse_mobius_options(value: str) -> dict:
+    """Parse typed Mobius options without depending on the ORT GenAI model builder."""
+    options = {}
+    for item in value.split(","):
+        key, separator, raw_value = item.partition("=")
+        key, raw_value = key.strip(), raw_value.strip()
+        if not separator or not key or not raw_value:
+            raise ValueError(f"Invalid Mobius option {item!r}. Expected key=value.")
+        if key in options:
+            raise ValueError(f"Duplicate Mobius option {key!r}.")
+        try:
+            options[key] = json.loads(raw_value)
+        except json.JSONDecodeError:
+            options[key] = raw_value
+    return options
+
+
 class CaptureOnnxGraphCommand(BaseOliveCLICommand):
     @staticmethod
     def register_subcommand(parser: ArgumentParser):
@@ -183,7 +201,12 @@ def register_subcommand(parser: ArgumentParser):
             "--extra_mb_options",
             type=str,
             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'.",
+            help=(
+                "Extra comma-separated key=value options for Model Builder or Mobius Builder. "
+                "Mobius values are parsed as JSON scalars or strings, e.g., "
+                "'vision_patch_capacity=2520,vision_batch_size=1'. "
+                "Model Builder example: 'int4_is_symmetric=true,int4_op_types_to_quantize=MatMul/Gemm'."
+            ),
         )
 
         sub_parser.add_argument(
@@ -276,6 +299,8 @@ def _get_run_config(self, tempdir: str) -> dict:
             to_replace.append((("passes", "b", "precision"), self.args.precision))
             if self.args.int4_accuracy_level is not None:
                 to_replace.append((("passes", "b", "matmul_nbits_accuracy_level"), self.args.int4_accuracy_level))
+            if self.args.extra_mb_options:
+                to_replace.append((("passes", "b", "extra_options"), parse_mobius_options(self.args.extra_mb_options)))
         elif is_diffusers_model:
             del config["passes"]["m"]
             del config["passes"]["b"]

olive/passes/onnx/mobius_model_builder.py

@@ -141,6 +141,17 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon
                     "Requires a Mobius version supporting matmul_nbits_accuracy_level."
                 ),
             ),
+            "extra_options": PassConfigParam(
+                type_=dict,
+                required=False,
+                default_value=None,
+                description=(
+                    "Optional model-specific build options passed to mobius.build(extra_options=...). "
+                    "For example, {'vision_patch_capacity': 2520, 'vision_batch_size': 1}. "
+                    "Supported keys and values are validated by the installed Mobius version. "
+                    "These affect graph construction, unlike genai_config runtime overrides."
+                ),
+            ),
             "genai_config": PassConfigParam(
                 type_=dict,
                 required=False,
@@ -230,6 +241,7 @@ def _run_for_config(
             if config.matmul_nbits_accuracy_level is not None
             else {}
         )
+        extra_options_kwargs = {"extra_options": config.extra_options} if config.extra_options else {}
         pkg = build(
             model_id,
             dtype=dtype_str,
@@ -239,6 +251,7 @@ def _run_for_config(
             **text_only_kwargs,
             **revision_kwargs,
             **accuracy_kwargs,
+            **extra_options_kwargs,
         )
 
         # Determine which package components to export.

test/cli/test_cli.py

@@ -779,6 +779,74 @@ def test_capture_onnx_command_shared_accuracy_level_rejects_invalid(mock_run, ac
     mock_run.assert_not_called()
 
 
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+@pytest.mark.parametrize(
+    ("options", "expected"),
+    [
+        ("vision_patch_capacity=2520", {"vision_patch_capacity": 2520}),
+        (
+            "vision_patch_capacity=2520,vision_batch_size=1",
+            {"vision_patch_capacity": 2520, "vision_batch_size": 1},
+        ),
+        (
+            "vision_patch_capacity=2520,vision_batch_size=null",
+            {"vision_patch_capacity": 2520, "vision_batch_size": None},
+        ),
+    ],
+)
+def test_capture_onnx_mobius_forwards_typed_extra_options(_, mock_run, tmp_path, options, expected):
+    with patch("olive.cli.base.BaseOliveCLICommand._parse_extra_options", side_effect=AssertionError):
+        cli_main(
+            [
+                "capture-onnx-graph",
+                "-m",
+                "dummy-model-id",
+                "-o",
+                str(tmp_path / "output"),
+                "--use_mobius_builder",
+                "--precision",
+                "fp32",
+                "--extra_mb_options",
+                options,
+                "--int4_accuracy_level",
+                "4",
+            ]
+        )
+    config = mock_run.call_args[0][0]
+    assert config["passes"]["b"]["extra_options"] == expected
+    assert config["passes"]["b"]["matmul_nbits_accuracy_level"] == 4
+    assert config["systems"]["local_system"]["accelerators"][0]["execution_providers"] == ["CPUExecutionProvider"]
+
+
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+@pytest.mark.parametrize(
+    ("options", "message"),
+    [
+        ("vision_patch_capacity", "Expected key=value"),
+        ("=2520", "Expected key=value"),
+        ("vision_batch_size=", "Expected key=value"),
+        ("vision_patch_capacity=2520,vision_patch_capacity=72", "Duplicate Mobius option"),
+    ],
+)
+def test_capture_onnx_mobius_rejects_malformed_extra_options(_, mock_run, tmp_path, options, message):
+    with pytest.raises(ValueError, match=message):
+        cli_main(
+            [
+                "capture-onnx-graph",
+                "-m",
+                "dummy-model-id",
+                "-o",
+                str(tmp_path / "output"),
+                "--use_mobius_builder",
+                "--extra_mb_options",
+                options,
+            ]
+        )
+    mock_run.assert_not_called()
+
+
 @patch("olive.workflows.run")
 @patch("olive.cli.capture_onnx.get_input_model_config")
 @pytest.mark.parametrize("exporter_flags", [[], ["--use_dynamo_exporter"]])
@@ -803,6 +871,27 @@ def test_capture_onnx_command_shared_accuracy_level_requires_builder(
     mock_run.assert_not_called()
 
 
+@patch("olive.workflows.run")
+@patch("huggingface_hub.repo_exists", return_value=True)
+def test_capture_onnx_model_builder_preserves_extra_options_parser(_, mock_run, tmp_path):
+    expected = {"filename": "model=part.onnx"}
+    with patch("olive.cli.base.BaseOliveCLICommand._parse_extra_options", return_value=expected) as parser:
+        cli_main(
+            [
+                "capture-onnx-graph",
+                "-m",
+                "dummy-model-id",
+                "-o",
+                str(tmp_path / "output"),
+                "--use_model_builder",
+                "--extra_mb_options",
+                "filename=model=part.onnx",
+            ]
+        )
+    parser.assert_called_once_with(["filename=model=part.onnx"])
+    assert mock_run.call_args[0][0]["passes"]["m"]["extra_options"] == expected
+
+
 @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/passes/onnx/test_mobius_model_builder.py

@@ -418,6 +418,30 @@ def test_text_only_true_forwarded_to_mobius_build(tmp_path):
     assert mock_build.call_args.kwargs["text_only"] is True
 
 
+@pytest.mark.parametrize(
+    "extra_options",
+    [
+        {"vision_patch_capacity": 2520, "vision_batch_size": 1},
+        {"vision_patch_capacity": 2520, "vision_batch_size": None},
+    ],
+)
+def test_extra_options_forwarded_to_mobius_build(tmp_path, extra_options):
+    out = tmp_path / "out"
+    pkg = _fake_pkg(["decoder", "vision_encoder", "embedding"], out)
+    p = create_pass_from_dict(
+        MobiusBuilder,
+        {"precision": "fp32", "extra_options": extra_options},
+        disable_search=True,
+        accelerator_spec=AcceleratorSpec(
+            accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider
+        ),
+    )
+    with _patch_build(pkg) as mock_build:
+        result = p.run(_make_hf_model("org/model"), out)
+    assert mock_build.call_args.kwargs["extra_options"] == extra_options
+    assert isinstance(result, CompositeModelHandler)
+
+
 def test_hf_task_not_forwarded_to_mobius_build(tmp_path):
     """Mobius auto-detects its task instead of receiving the HfModelHandler task."""
     out = tmp_path / "out"