From 2a3a502cb80e6974f2ee24d6cff43d967601bc10 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 1 Oct 2026 04:57:35 +0000 Subject: [PATCH 1/4] Support recipe-defined Mobius exports --- olive/passes/onnx/mobius_model_builder.py | 162 ++++++++++-------- test/passes/onnx/test_mobius_model_builder.py | 65 ++++++- 2 files changed, 158 insertions(+), 69 deletions(-) diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index d8c9a6f041..a6709d6168 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -8,6 +8,7 @@ import json import logging +from collections.abc import Callable from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar @@ -18,6 +19,7 @@ from olive.model.handler.composite import CompositeModelHandler from olive.passes import Pass from olive.passes.olive_pass import PassConfigParam +from olive.passes.pass_config import ParamCategory, get_user_script_data_config if TYPE_CHECKING: from olive.hardware.accelerator import AcceleratorSpec @@ -107,6 +109,7 @@ def is_accelerator_agnostic(cls, accelerator_spec: AcceleratorSpec) -> bool: @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: return { + **get_user_script_data_config(), "precision": PassConfigParam( type_=Precision, required=False, @@ -152,6 +155,24 @@ def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassCon "the model's components." ), ), + "exporter_func": PassConfigParam( + type_=str | Callable, + required=False, + default_value=None, + category=ParamCategory.OBJECT, + description=( + "Optional recipe-local function that exports a Mobius package " + "directly to output_model_path and returns a mapping containing " + "a non-empty 'components' list. The function receives model, " + "output_dir, precision, execution_provider, and exporter_config." + ), + ), + "exporter_config": PassConfigParam( + type_=dict, + required=False, + default_value=None, + description="Configuration forwarded to exporter_func.", + ), } def _run_for_config( @@ -167,13 +188,6 @@ def _run_for_config( validate_olive_onnx_export(model.get_hf_model_config().to_dict().get("quantization_config") or {}) - try: - from mobius import build - except ImportError as exc: - raise ImportError( - "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx" - ) from exc - # Map Olive EP to mobius EP. If unsupported/unknown, fall back to mobius default EP. requested_ep = self.accelerator_spec.execution_provider ep_str: str = EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT) @@ -187,7 +201,6 @@ def _run_for_config( dtype_str: str = _PRECISION_TO_DTYPE.get(config.precision, "f32") model_id: str = model.model_name_or_path - load_kwargs = model.get_load_kwargs() revision: str | None = load_kwargs.get("revision") trust_remote_code: bool = load_kwargs.get("trust_remote_code", False) @@ -213,82 +226,95 @@ def _run_for_config( output_dir = Path(output_model_path) output_dir.mkdir(parents=True, exist_ok=True) - - text_only_kwargs = {"text_only": True} if config.text_only else {} - revision_kwargs = {"revision": revision} if revision is not None else {} - pkg = build( - model_id, - dtype=dtype_str, - execution_provider=ep_str, - load_weights=True, - trust_remote_code=trust_remote_code, - **text_only_kwargs, - **revision_kwargs, - ) - - # Determine which package components to export. - all_keys = list(pkg.keys()) - if config.components_to_export is not None: - requested = set(config.components_to_export) - unknown = requested - set(all_keys) - if unknown: + custom_exporter = config.exporter_func is not None + if custom_exporter: + if config.components_to_export is not None or config.text_only: raise ValueError( - f"MobiusBuilder: components_to_export contains unknown component(s): {sorted(unknown)}. " - f"Available components from this model: {sorted(all_keys)}" + "MobiusBuilder: exporter_func cannot be combined with components_to_export or text_only." ) - package_keys = [k for k in all_keys if k in requested] - logger.info( - "MobiusBuilder: exporting subset of components %s (skipping %s)", - package_keys, - [k for k in all_keys if k not in requested], - ) - - def components_filter(name: str) -> bool: - return name in requested - else: - package_keys = all_keys - components_filter = None - - # ModelPackage.save() handles both single and multi-component layouts: - # single component → /model.onnx - # multi-component → //model.onnx for each key - pkg.save(str(output_dir), components=components_filter) - - # ORT GenAI config generation assumes every component in `pkg` was actually saved to - # disk (e.g. it unconditionally writes a "decoder/model.onnx" filename reference for - # multimodal packages). For a partial export (components_to_export set), that produces - # a genai_config.json/tokenizer set that references components we deliberately omitted - # — invalid artifacts, since mobius has no API to generate GenAI config for a subset of - # a package. Skip config generation entirely in that case and let the caller (e.g. a - # recipe combining this partial export with another tool's output) assemble the final - # genai_config.json itself. - if components_filter is not None: - logger.info( - "MobiusBuilder: components_to_export is set; skipping ORT GenAI config generation " - "since it cannot be scoped to a subset of components. The caller is responsible for " - "producing a genai_config.json that covers the full pipeline." + exporter = self._user_module_loader.load_object(config.exporter_func) + result = exporter( + model=model, + output_dir=output_dir, + precision=config.precision, + execution_provider=ep_str, + exporter_config=config.exporter_config or {}, ) - genai_artifacts = {} + if not isinstance(result, dict) or not isinstance(result.get("components"), list): + raise ValueError("MobiusBuilder: exporter_func must return a mapping with a 'components' list.") + package_keys = result["components"] + if not package_keys or not all(isinstance(name, str) and name for name in package_keys): + raise ValueError("MobiusBuilder: exporter_func returned invalid components.") + all_keys = package_keys + genai_artifacts = {path.name: str(path) for path in output_dir.iterdir() if path.is_file()} else: - # Generate ORT GenAI config artifacts (genai_config.json, tokenizer - # files, processor configs) alongside the ONNX models. - genai_artifacts = self._write_genai_config( - pkg, - str(output_dir), + try: + from mobius import build + except ImportError as exc: + raise ImportError( + "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx" + ) from exc + + text_only_kwargs = {"text_only": True} if config.text_only else {} + pkg = build( model_id, - ep_str, revision=revision, + dtype=dtype_str, + execution_provider=ep_str, + load_weights=True, trust_remote_code=trust_remote_code, + **text_only_kwargs, ) if config.genai_config: self._update_genai_config(output_dir, config.genai_config) + all_keys = list(pkg.keys()) + if config.components_to_export is not None: + requested = set(config.components_to_export) + unknown = requested - set(all_keys) + if unknown: + raise ValueError( + "MobiusBuilder: components_to_export contains unknown " + f"component(s): {sorted(unknown)}. Available components " + f"from this model: {sorted(all_keys)}" + ) + package_keys = [key for key in all_keys if key in requested] + logger.info( + "MobiusBuilder: exporting subset of components %s (skipping %s)", + package_keys, + [key for key in all_keys if key not in requested], + ) + + def components_filter(name: str) -> bool: + return name in requested + else: + package_keys = all_keys + components_filter = None + + pkg.save(str(output_dir), components=components_filter) + if components_filter is not None: + logger.info( + "MobiusBuilder: components_to_export is set; skipping ORT " + "GenAI config generation since it cannot be scoped to a " + "subset of components." + ) + genai_artifacts = {} + else: + genai_artifacts = self._write_genai_config( + pkg, + str(output_dir), + model_id, + ep_str, + revision=revision, + trust_remote_code=trust_remote_code, + ) + logger.info("MobiusBuilder: saved components %s to '%s'", package_keys, output_dir) # Use the single-component (root layout) path only when the model is # architecturally single-component. A multi-component model filtered # down to one component still uses component sub-directories on disk. - if len(all_keys) == 1: + if len(all_keys) == 1 and not custom_exporter: # Single-component model (most LLMs): return a plain ONNXModelHandler. onnx_path = output_dir / "model.onnx" if not onnx_path.exists(): diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index c213f5bfe0..95a08f02f9 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -78,11 +78,16 @@ def _make_hf_model(model_path: str, load_kwargs: dict | None = None, task: str | return model -def _make_pass(ep: str = ExecutionProvider.CPUExecutionProvider, text_only: bool | None = None) -> MobiusBuilder: +def _make_pass( + ep: str = ExecutionProvider.CPUExecutionProvider, + text_only: bool | None = None, + **extra_config, +) -> MobiusBuilder: accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ep) pass_config = {"precision": "fp32"} if text_only is not None: pass_config["text_only"] = text_only + pass_config.update(extra_config) return create_pass_from_dict( MobiusBuilder, pass_config, @@ -293,10 +298,68 @@ def test_default_config_params(): assert "precision" in config assert config["text_only"].default_value is False assert config["text_only"].required is False + assert config["exporter_func"].default_value is None + assert config["exporter_config"].default_value is None + assert "user_script" in config + assert "script_dir" in config assert "execution_provider" not in config assert "trust_remote_code" not in config +def test_recipe_exporter_function_creates_composite_package(tmp_path): + """A recipe callback can own multi-component package export.""" + + def exporter(**kwargs): + output = Path(kwargs["output_dir"]) + for name in ("backbone", "pointer_head"): + (output / name).mkdir(parents=True) + (output / name / "model.onnx").write_text("dummy") + (output / "component_manifest.json").write_text("{}") + assert kwargs["exporter_config"] == {"artifact_path": "artifact"} + return {"components": ["backbone", "pointer_head"]} + + p = _make_pass( + exporter_func=exporter, + exporter_config={"artifact_path": "artifact"}, + ) + result = p.run(_make_hf_model("Qwen/Qwen3.5-4B-Base"), tmp_path / "out") + + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == ["backbone", "pointer_head"] + assert result.model_attributes["additional_files"] == [str(tmp_path / "out" / "component_manifest.json")] + + +def test_recipe_exporter_function_requires_component_list(tmp_path): + """Invalid recipe callback results fail before handler construction.""" + p = _make_pass(exporter_func=lambda **_: {}) + with pytest.raises(ValueError, match="'components' list"): + p.run(_make_hf_model("Qwen/Qwen3-8B"), tmp_path / "out") + + +def test_recipe_exporter_function_loads_from_user_script(tmp_path): + """Recipe configs can resolve their exporter callback by name.""" + script = tmp_path / "user_script.py" + script.write_text( + """ +from pathlib import Path + +def export_package(**kwargs): + output = Path(kwargs["output_dir"]) + (output / "encoder").mkdir(parents=True) + (output / "encoder" / "model.onnx").write_text("dummy") + return {"components": ["encoder"]} +""" + ) + p = _make_pass( + exporter_func="export_package", + user_script=str(script), + ) + result = p.run(_make_hf_model("Qwen/Qwen3-8B"), tmp_path / "out") + + assert isinstance(result, CompositeModelHandler) + assert result.model_component_names == ["encoder"] + + def test_is_not_accelerator_agnostic(): """Pass must be EP-specific because it chooses fused ops based on the EP.""" accelerator_spec = AcceleratorSpec( From fc4afca2d3a5e50fcac85c823fb27b062005afac Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 1 Oct 2026 05:02:02 +0000 Subject: [PATCH 2/4] Preserve Mobius dependency validation --- olive/passes/onnx/mobius_model_builder.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index a6709d6168..9422ec552f 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -181,6 +181,13 @@ def _run_for_config( config: type[BasePassConfig], output_model_path: str, ) -> ONNXModelHandler | CompositeModelHandler: + try: + import mobius + except ImportError as exc: + raise ImportError( + "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx" + ) from exc + if not isinstance(model, HfModelHandler): raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.") @@ -248,15 +255,8 @@ def _run_for_config( all_keys = package_keys genai_artifacts = {path.name: str(path) for path in output_dir.iterdir() if path.is_file()} else: - try: - from mobius import build - except ImportError as exc: - raise ImportError( - "mobius-onnx is required to run MobiusBuilder. Install with: pip install mobius-onnx" - ) from exc - text_only_kwargs = {"text_only": True} if config.text_only else {} - pkg = build( + pkg = mobius.build( model_id, revision=revision, dtype=dtype_str, From bcbc570733602ada6aa1e17de74ab7177fef9897 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 1 Oct 2026 06:09:20 +0000 Subject: [PATCH 3/4] Preserve custom Mobius package artifacts --- olive/passes/onnx/mobius_model_builder.py | 9 +++--- test/passes/onnx/test_mobius_model_builder.py | 28 +++++++++++++++++-- 2 files changed, 31 insertions(+), 6 deletions(-) diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index 9422ec552f..5a5dbc2077 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -253,7 +253,8 @@ def _run_for_config( if not package_keys or not all(isinstance(name, str) and name for name in package_keys): raise ValueError("MobiusBuilder: exporter_func returned invalid components.") all_keys = package_keys - genai_artifacts = {path.name: str(path) for path in output_dir.iterdir() if path.is_file()} + component_paths = {output_dir / name for name in package_keys} + genai_artifacts = {path.name: str(path) for path in output_dir.iterdir() if path not in component_paths} else: text_only_kwargs = {"text_only": True} if config.text_only else {} pkg = mobius.build( @@ -331,9 +332,9 @@ def components_filter(name: str) -> bool: model_path=str(output_dir), onnx_file_name="model.onnx", model_attributes={ + **(model.model_attributes or {}), "mobius_package_keys": package_keys, "additional_files": additional_files, - **(model.model_attributes or {}), }, ) @@ -362,9 +363,9 @@ def components_filter(name: str) -> bool: model_path=str(component_dir), onnx_file_name="model.onnx", model_attributes={ + **(model.model_attributes or {}), "mobius_component": key, "additional_files": component_additional_files, - **(model.model_attributes or {}), }, ) ) @@ -374,6 +375,7 @@ def components_filter(name: str) -> bool: model_component_names=package_keys, model_path=str(output_dir), model_attributes={ + **(model.model_attributes or {}), "mobius_package_keys": package_keys, # Preserve the /model.onnx subdirectory layout so # ORT GenAI can resolve each component by its "filename" key. @@ -384,7 +386,6 @@ def components_filter(name: str) -> bool: # so they end up at the package root (alongside genai_config.json), # not duplicated into each / subdirectory. "additional_files": sorted(set(genai_artifacts.values())), - **(model.model_attributes or {}), }, ) diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index 95a08f02f9..eb3b80d781 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -314,6 +314,8 @@ def exporter(**kwargs): for name in ("backbone", "pointer_head"): (output / name).mkdir(parents=True) (output / name / "model.onnx").write_text("dummy") + (output / "assets").mkdir() + (output / "assets" / "metadata.json").write_text("{}") (output / "component_manifest.json").write_text("{}") assert kwargs["exporter_config"] == {"artifact_path": "artifact"} return {"components": ["backbone", "pointer_head"]} @@ -322,11 +324,33 @@ def exporter(**kwargs): exporter_func=exporter, exporter_config={"artifact_path": "artifact"}, ) - result = p.run(_make_hf_model("Qwen/Qwen3.5-4B-Base"), tmp_path / "out") + stale_file = tmp_path / "stale.txt" + stale_file.write_text("stale") + model = _make_hf_model("Qwen/Qwen3.5-4B-Base") + model.model_attributes = { + "additional_files": [str(stale_file)], + "mobius_package_keys": ["stale"], + "no_flatten": False, + "preserved": True, + } + result = p.run(model, tmp_path / "out") assert isinstance(result, CompositeModelHandler) assert result.model_component_names == ["backbone", "pointer_head"] - assert result.model_attributes["additional_files"] == [str(tmp_path / "out" / "component_manifest.json")] + assert result.model_attributes == { + "additional_files": [ + str(tmp_path / "out" / "assets"), + str(tmp_path / "out" / "component_manifest.json"), + str(tmp_path / "out" / "stale.txt"), + ], + "mobius_package_keys": ["backbone", "pointer_head"], + "no_flatten": True, + "preserved": True, + } + for name, component in result.get_model_components(): + assert component.model_attributes["additional_files"] == [] + assert component.model_attributes["mobius_component"] == name + assert component.model_attributes["preserved"] is True def test_recipe_exporter_function_requires_component_list(tmp_path): From 0cf2cc43124803bb2fb90a6323ca1b80be43c665 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 27 Oct 2026 19:31:36 +0000 Subject: [PATCH 4/4] Integrate custom Mobius exports with current builder Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- olive/passes/onnx/mobius_model_builder.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index 5a5dbc2077..ffda58bb50 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -257,17 +257,16 @@ def _run_for_config( genai_artifacts = {path.name: str(path) for path in output_dir.iterdir() if path not in component_paths} else: text_only_kwargs = {"text_only": True} if config.text_only else {} + revision_kwargs = {"revision": revision} if revision is not None else {} pkg = mobius.build( model_id, - revision=revision, dtype=dtype_str, execution_provider=ep_str, load_weights=True, trust_remote_code=trust_remote_code, **text_only_kwargs, + **revision_kwargs, ) - if config.genai_config: - self._update_genai_config(output_dir, config.genai_config) all_keys = list(pkg.keys()) if config.components_to_export is not None: @@ -309,6 +308,8 @@ def components_filter(name: str) -> bool: revision=revision, trust_remote_code=trust_remote_code, ) + if config.genai_config: + self._update_genai_config(output_dir, config.genai_config) logger.info("MobiusBuilder: saved components %s to '%s'", package_keys, output_dir)