MobiusBuilder + cache: produce loadable ORT GenAI composite packages (microsoft#2472)

Three fixes that together let `MobiusBuilder`-produced multimodal ORT
GenAI packages round-trip through Olive without corruption.

## Problem 1: Composite output flattened, sidecars dropped

`mobius.build()` writes a multi-component ORT GenAI package laid out as:

```
<output_dir>/
├── decoder/model.onnx
├── embedding/model.onnx
├── vision_encoder/model.onnx
├── audio_encoder/model.onnx
├── genai_config.json
├── tokenizer.json / tokenizer_config.json / chat_template.jinja
├── image_processor.json
└── audio_feature_extraction.json
```

ORT GenAI reads `genai_config.json` from the root and resolves each
component's relative `filename` (e.g. `decoder/model.onnx`) underneath.

By default Olive's cache **flattens** `CompositeModelHandler` outputs to
top-level `<component_name>.onnx` and **drops sidecars**
(`olive/cache.py:394-407`). Result:

```
output/
├── decoder.onnx, decoder.onnx.data, embedding.onnx, vision_encoder.onnx, audio_encoder.onnx
└── (no genai_config.json, no tokenizer*, no image_processor.json)
```

Breaks ORT GenAI loading entirely.

### Fix (commit 1: `15f7c0f0`)
Set `no_flatten=True` on the `CompositeModelHandler.model_attributes`.
Olive's cache then uses the `shutil.copytree(source_path,
actual_output_dir)` branch, preserving the `<component>/model.onnx`
subdirectory layout and copying every sidecar file alongside.

## Problem 2: Sidecars duplicated into each component dir after a
downstream pass

When a downstream pass (e.g. `OnnxKQuantQuantization`) runs on the
composite, Olive splits it and runs the pass per component
(`olive_pass.py:230-241`). `_carry_forward_additional_files` then copies
each input component's `additional_files` into the output component's
directory.

The original implementation attached the shared package sidecars to
**every** component's `additional_files`. After the quant pass:

```
cuda/int4/models/
├── decoder/{model.onnx, model.onnx.data, genai_config.json, tokenizer.json, image_processor.json, ...}
├── embedding/{model.onnx, model.onnx.data, genai_config.json, tokenizer.json, ...}
├── vision_encoder/...
├── audio_encoder/...
└── (no genai_config.json at root!)
```

Sidecars live in every subdirectory but the root has none — still broken
for ORT GenAI.

### Fix (commit 2: `3b2ac939`)
Attach shared package sidecars to the **CompositeModelHandler**'s
`model_attributes['additional_files']`, not each component's. The
composite's `model_path` *is* the package root, so the cross-pass
carry-forward / `no_flatten` copytree puts them at the root.
Per-component `additional_files` now only carries files that genuinely
live inside the component dir.

## Problem 3: `model_config.json` records nonexistent component paths

`Cache.save_model()`'s `no_flatten` branch (`olive/cache.py:404-407`)
rewrites every component's `model_path` to the composite root while
leaving `onnx_file_name` untouched. For models whose components live
under `<root>/<component>/model.onnx` (both `MobiusBuilder` packages and
`optimum`-exported diffusers pipelines from
`conversion.py:829-831,863`), the saved `model_config.json` then points
each component at `<root>/model.onnx`, which does not exist. Reloading
the saved composite fails.

### Fix (commit 3: `14a15462`)
Rebase each component's `model_path` by its position relative to the
original source dir. A component at `<source>/decoder/` now lands at
`<actual_output_dir>/decoder/`, with `onnx_file_name` unchanged.
Components that fall outside the source tree fall back to the package
root.

## Verification

Two recipes from microsoft/olive-recipes#381:

**`cpu/fp32/config.json` (MobiusBuilder only):**
```
cpu/fp32/models/
├── audio_encoder/model.onnx
├── decoder/model.onnx
├── embedding/model.onnx
├── vision_encoder/model.onnx
├── audio_feature_extraction.json
├── chat_template.jinja
├── genai_config.json
├── image_processor.json
├── tokenizer.json
└── tokenizer_config.json
```

**`cuda/int4/config.json` (MobiusBuilder → OnnxKQuantQuantization):**
```
cuda/int4/models/
├── audio_encoder/{model.onnx, model.onnx.data}
├── decoder/{model.onnx, model.onnx.data}
├── embedding/{model.onnx, model.onnx.data}
├── vision_encoder/{model.onnx, model.onnx.data}
├── audio_feature_extraction.json
├── chat_template.jinja
├── genai_config.json
├── image_processor.json
├── tokenizer.json
└── tokenizer_config.json
```

`model_config.json` records the correct component paths
(`decoder/model.onnx`, `embedding/model.onnx`,
`vision_encoder/model.onnx`,
`audio_encoder/model.onnx`), matching the on-disk layout.

`ruff check` and `ruff format --check` both pass on the changed files.

## Notes

- `no_flatten=True` reuses the existing pattern from the
diffusers/optimum integration (`olive/cache.py:398`).
- `_carry_forward_additional_files` already works for any model whose
`model_path` is a directory (`olive_pass.py:266` `is_dir()` branch), so
`CompositeModelHandler` (whose `model_path` is the package root dir) is
supported without any framework changes.
- The cache-layer component-path rebasing also fixes a latent bug in the
existing optimum/diffusers `no_flatten` consumer; nobody hit it because
nobody had reloaded those saved configs.
- No public API change.

---------

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

Author: justinchuby

olive/cache.py

@@ -396,14 +396,57 @@ def save_model(
             model_attributes = model_json_config.get("model_attributes") or {}
 
             if model_attributes.get("no_flatten"):
-                # Preserve directory structure (e.g., for diffusers models exported by optimum)
+                # Preserve directory structure (e.g., for diffusers models
+                # exported by optimum, or multimodal ORT GenAI packages from
+                # MobiusBuilder where components live in <root>/<component>/).
                 source_path = Path(model_json_config["model_path"])
+                source_path_resolved = source_path.resolve()
                 if source_path.exists():
                     shutil.copytree(source_path, actual_output_dir, dirs_exist_ok=overwrite)
 
-                # Update component paths to point to new location
+                def _rebase_additional_files(config: dict, fallback_dir: Path):
+                    model_attributes = config.get("model_attributes") or {}
+                    additional_files = model_attributes.get("additional_files") or []
+                    if not additional_files:
+                        return
+
+                    rebased_additional_files = []
+                    for additional_file in additional_files:
+                        source_additional_file = Path(additional_file)
+                        try:
+                            relative = source_additional_file.resolve().relative_to(source_path_resolved)
+                            output_additional_file = actual_output_dir / relative
+                        except ValueError:
+                            output_additional_file = fallback_dir / source_additional_file.name
+                            if source_additional_file.exists() and not output_additional_file.exists():
+                                output_additional_file.parent.mkdir(parents=True, exist_ok=True)
+                                shutil.copy2(source_additional_file, output_additional_file)
+                        rebased_additional_files.append(str(output_additional_file))
+
+                    model_attributes["additional_files"] = rebased_additional_files
+                    config["model_attributes"] = model_attributes
+
+                # Rewrite each component's model_path so it points into the
+                # new output location while preserving the component's
+                # relative position underneath the package root. Without
+                # rebasing component paths the saved model_config.json
+                # cannot be loaded (and onnx_file_name is left untouched,
+                # so we must not collapse component subdirs into the root).
+                _rebase_additional_files(model_json_config, actual_output_dir)
                 for component in model_json_config["model_components"]:
-                    component["config"]["model_path"] = str(actual_output_dir)
+                    component_config = component["config"]
+                    component_model_path = component_config.get("model_path")
+                    if component_model_path:
+                        try:
+                            relative = Path(component_model_path).resolve().relative_to(source_path_resolved)
+                        except ValueError:
+                            # Component path is not under the composite root;
+                            # fall back to placing it at the package root.
+                            relative = Path()
+                        component_config["model_path"] = str(actual_output_dir / relative)
+                    else:
+                        component_config["model_path"] = str(actual_output_dir)
+                    _rebase_additional_files(component_config, Path(component_config["model_path"]))
                 model_json_config["model_path"] = str(actual_output_dir)
             else:
                 copied_components = []

olive/passes/onnx/mobius_model_builder.py

@@ -201,7 +201,9 @@ def _run_for_config(
             )
 
         # Multi-component model (VLMs, encoder-decoders, diffusion pipelines):
-        # mobius saves each component to <output_dir>/<key>/model.onnx.
+        # mobius saves each component to <output_dir>/<key>/model.onnx with shared
+        # sidecar files (genai_config.json, tokenizer.json, image_processor.json,
+        # audio_feature_extraction.json) at output_dir root.
         components = []
         for key in package_keys:
             component_dir = output_dir / key
@@ -211,18 +213,20 @@ def _run_for_config(
                     f"MobiusBuilder: expected output file not found: {onnx_path}. "
                     f"mobius.build() may have failed silently for component '{key}'."
                 )
-            additional_files = sorted(
+            # Per-component additional files: only files that live inside the
+            # component's own directory. Shared sidecars (genai_config, tokenizer,
+            # image_processor) are attached to the composite handler below so
+            # they land in the output root, not duplicated in every component.
+            component_additional_files = sorted(
                 {str(fp) for fp in component_dir.iterdir()} - {str(onnx_path), str(onnx_path) + ".data"}
             )
-            # Include ORT GenAI artifacts from root output_dir (shared across components)
-            additional_files = sorted(set(additional_files) | set(genai_artifacts.values()))
             components.append(
                 ONNXModelHandler(
                     model_path=str(component_dir),
                     onnx_file_name="model.onnx",
                     model_attributes={
                         "mobius_component": key,
-                        "additional_files": additional_files,
+                        "additional_files": component_additional_files,
                         **(model.model_attributes or {}),
                     },
                 )
@@ -234,6 +238,15 @@ def _run_for_config(
             model_path=str(output_dir),
             model_attributes={
                 "mobius_package_keys": package_keys,
+                # Preserve the <component>/model.onnx subdirectory layout so
+                # ORT GenAI can resolve each component by its "filename" key.
+                # Without this, Olive's cache flattens components to top-level
+                # <name>.onnx files and breaks GenAI loading.
+                "no_flatten": True,
+                # Shared package-level sidecars carried via the composite handler
+                # so they end up at the package root (alongside genai_config.json),
+                # not duplicated into each <component>/ subdirectory.
+                "additional_files": sorted(set(genai_artifacts.values())),
                 **(model.model_attributes or {}),
             },
         )

test/passes/onnx/test_mobius_model_builder.py

@@ -220,12 +220,22 @@ def test_genai_artifacts_in_single_component(tmp_path):
 
 
 def test_genai_artifacts_in_multi_component(tmp_path):
-    """ORT GenAI artifacts must be included in all components of multi-component models."""
+    """ORT GenAI artifacts must be attached at composite-level, not duplicated per component."""
     out = tmp_path / "out"
     out.mkdir(parents=True, exist_ok=True)
     keys = ["model", "vision", "embedding"]
     pkg = _fake_pkg(keys, out)
 
+    def _save_with_component_sidecar(directory: str, **_kwargs):
+        out_dir = Path(directory)
+        for key in keys:
+            component_dir = out_dir / key
+            component_dir.mkdir(parents=True, exist_ok=True)
+            (component_dir / "model.onnx").write_text("dummy")
+        (out_dir / "vision" / "vision_local.txt").write_text("vision")
+
+    pkg.save.side_effect = _save_with_component_sidecar
+
     # Mock genai artifact files
     genai_config = str(out / "genai_config.json")
     image_processor = str(out / "image_processor.json")
@@ -240,11 +250,19 @@ def test_genai_artifacts_in_multi_component(tmp_path):
         result = p.run(_make_hf_model("microsoft/phi-4-vision"), out)
 
     assert isinstance(result, CompositeModelHandler)
-    # Verify all components include genai artifacts
-    for component in result.model_components:
+    composite_additional_files = result.model_attributes.get("additional_files", [])
+    assert genai_config in composite_additional_files
+    assert image_processor in composite_additional_files
+
+    # Shared GenAI sidecars should not be duplicated into each component.
+    components = list(result.model_components)
+    for component in components:
         additional_files = component.model_attributes.get("additional_files", [])
-        assert genai_config in additional_files
-        assert image_processor in additional_files
+        assert genai_config not in additional_files
+        assert image_processor not in additional_files
+
+    vision = components[result.model_component_names.index("vision")]
+    assert str(out / "vision" / "vision_local.txt") in vision.model_attributes.get("additional_files", [])
 
 
 # ---------------------------------------------------------------------------

test/test_cache.py

@@ -350,6 +350,73 @@ def test_save_model_with_custom_onnx_filename(self, tmp_path):
         with open(output_json_path) as f:
             assert expected_output_path == json.load(f)["config"]["model_path"]
 
+    def test_save_model_no_flatten_rebases_component_and_additional_file_paths(self, tmp_path):
+        # setup
+        model_id = "composite_model"
+        cache = CacheConfig(cache_dir=tmp_path / "cache").create_cache()
+
+        source_dir = tmp_path / "source_model"
+        decoder_dir = source_dir / "decoder"
+        embedding_dir = source_dir / "embedding"
+        decoder_dir.mkdir(parents=True, exist_ok=True)
+        embedding_dir.mkdir(parents=True, exist_ok=True)
+        (decoder_dir / "model.onnx").write_text("decoder")
+        (embedding_dir / "model.onnx").write_text("embedding")
+        (decoder_dir / "decoder_local.txt").write_text("decoder local")
+        (source_dir / "genai_config.json").write_text("{}")
+        (source_dir / "tokenizer.json").write_text("{}")
+
+        model_json = {
+            "type": "compositemodel",
+            "config": {
+                "model_path": str(source_dir),
+                "model_component_names": ["decoder", "embedding"],
+                "model_components": [
+                    {
+                        "type": "onnxmodel",
+                        "config": {
+                            "model_path": str(decoder_dir),
+                            "onnx_file_name": "model.onnx",
+                            "model_attributes": {"additional_files": [str(decoder_dir / "decoder_local.txt")]},
+                        },
+                    },
+                    {
+                        "type": "onnxmodel",
+                        "config": {"model_path": str(embedding_dir), "onnx_file_name": "model.onnx"},
+                    },
+                ],
+                "model_attributes": {
+                    "no_flatten": True,
+                    "additional_files": [str(source_dir / "genai_config.json"), str(source_dir / "tokenizer.json")],
+                },
+            },
+        }
+        with cache.get_model_json_path(model_id).open("w") as f:
+            json.dump(model_json, f)
+
+        output_dir = tmp_path / "output"
+        output_json = cache.save_model(model_id, output_dir, True)
+
+        # assert copied layout
+        assert (output_dir / "decoder" / "model.onnx").exists()
+        assert (output_dir / "embedding" / "model.onnx").exists()
+        assert (output_dir / "genai_config.json").exists()
+        assert (output_dir / "tokenizer.json").exists()
+
+        # assert rewritten config paths
+        assert output_json["config"]["model_path"] == str(output_dir)
+        decoder_component = output_json["config"]["model_components"][0]["config"]
+        assert decoder_component["model_path"] == str(output_dir / "decoder")
+        assert decoder_component["onnx_file_name"] == "model.onnx"
+        assert decoder_component["model_attributes"]["additional_files"] == [
+            str(output_dir / "decoder" / "decoder_local.txt")
+        ]
+
+        assert output_json["config"]["model_attributes"]["additional_files"] == [
+            str(output_dir / "genai_config.json"),
+            str(output_dir / "tokenizer.json"),
+        ]
+
 
 class TestSharedCache:
     @pytest.fixture(autouse=True)