diff --git a/docs/source/reference/cli.rst b/docs/source/reference/cli.rst index f7e33c86c1..1829d151b3 100644 --- a/docs/source/reference/cli.rst +++ b/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 ======== diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index f3cb48019b..c7f285bec8 100644 --- a/olive/cli/capture_onnx.py +++ b/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"] diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index cf2a1835ba..a0b986d509 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ 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. diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index cdb40425d2..f2e1dbd17f 100644 --- a/test/cli/test_cli.py +++ b/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): diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index 98599649c2..64b71176fe 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/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"