Add multi build support

olive/cli/api.py

@@ -4,7 +4,7 @@
 # --------------------------------------------------------------------------
 import inspect
 from argparse import ArgumentParser, Namespace
-from typing import Any
+from typing import Any, Union
 
 from olive.cli.benchmark import BenchmarkCommand
 from olive.cli.capture_onnx import CaptureOnnxGraphCommand
@@ -300,15 +300,16 @@ def benchmark(model_name_or_path: str, **kwargs) -> WorkflowOutput:
     return _run_unified_command(BenchmarkCommand, **kwargs)
 
 
-def run(run_config: str, **kwargs) -> WorkflowOutput:
+def run(run_config: str, **kwargs) -> Union[WorkflowOutput, dict[str, WorkflowOutput]]:
     """Run a workflow.
 
     Args:
         run_config: Path to Olive workflow config
         **kwargs: All other CLI arguments supported by extract-adapters command
 
     Returns:
-        WorkflowOutput: Contains tuning results
+        WorkflowOutput for a single-pipeline workflow, or a ``dict[str, WorkflowOutput]`` keyed by build
+        name when the config declares ``builds``.
 
     """
     kwargs["run_config"] = run_config

olive/cli/base.py

@@ -151,8 +151,23 @@ def _run_workflow(self):
             workflow_output = olive_run(run_config)
             if getattr(self.args, "test", None) not in (None, False):
                 mark_test_output_path(self.args.output_path)
-                save_discrepancy_check_results(workflow_output, self.args.output_path)
-            if not workflow_output.has_output_model():
+                if not isinstance(workflow_output, dict):
+                    save_discrepancy_check_results(workflow_output, self.args.output_path)
+            if isinstance(workflow_output, dict):
+                # `builds` workflows return one WorkflowOutput per build keyed by build name.
+                builds_cfg = run_config.get("builds") or {}
+                build_default = builds_cfg.get("_default") or {}
+                for build_name, build_output in workflow_output.items():
+                    if build_output is None or not build_output.has_output_model():
+                        print(f"Build {build_name!r}: no output model produced. Please check the log for details.")
+                    else:
+                        build_output_dir = (
+                            (builds_cfg.get(build_name) or {}).get("output_dir")
+                            or build_default.get("output_dir")
+                            or self.args.output_path
+                        )
+                        print(f"Build {build_name!r}: model is saved under {build_output_dir}")
+            elif not workflow_output.has_output_model():
                 print("No output model produced. Please check the log for details.")
             else:
                 print(f"Model is saved at {self.args.output_path}")

olive/cli/capture_onnx.py

@@ -199,12 +199,14 @@ def run(self):
     def _get_run_config(self, tempdir: str) -> dict:
         config = deepcopy(TEMPLATE)
 
+        is_diffusers = is_valid_diffusers_model(self.args.model_name_or_path) if self.args.model_name_or_path else False
+
         if self.args.use_mobius_builder:
-            input_model_config = get_input_model_config(self.args)
+            if is_diffusers:
+                input_model_config = get_diffusers_input_model(self.args, self.args.model_name_or_path)
+            else:
+                input_model_config = get_input_model_config(self.args)
         else:
-            is_diffusers = (
-                is_valid_diffusers_model(self.args.model_name_or_path) if self.args.model_name_or_path else False
-            )
             if is_diffusers:
                 input_model_config = get_diffusers_input_model(self.args, self.args.model_name_or_path)
             else:

olive/common/hf/utils.py

@@ -9,7 +9,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from transformers import AutoConfig, AutoModel, AutoTokenizer, GenerationConfig
+from transformers import AutoConfig, AutoModel, AutoProcessor, AutoTokenizer, GenerationConfig
 
 from olive.common.hf.mappings import TASK_TO_PEFT_TASK_TYPE
 from olive.common.hf.mlflow import get_pretrained_name_or_path
@@ -342,6 +342,19 @@ def save_tokenizer(
     return tokenizer.save_pretrained(output_dir, **kwargs)
 
 
+def get_processor(model_name_or_path: str, **kwargs):
+    """Get HF model's processor if one exists."""
+    try:
+        return from_pretrained(AutoProcessor, model_name_or_path, "processor", **kwargs)
+    except (OSError, ValueError):
+        return None
+
+
+def save_processor(processor, output_dir: str, **kwargs) -> tuple[str]:
+    """Save input processor to output directory."""
+    return processor.save_pretrained(output_dir, **kwargs)
+
+
 def get_peft_task_type_from_task(task: str, fail_on_not_found=False) -> str:
     """Get peft task type from task."""
     peft_task_type = TASK_TO_PEFT_TASK_TYPE.get(task.replace("-with-past", ""), None)

olive/common/hf/wrapper.py

@@ -209,20 +209,23 @@ class ModelWrapper:
         "gptj": ["transformer.wte"],
         "opt": ["model.decoder.embed_tokens", "model.decoder.embed_positions"],
         "qwen": ["transformer.wte"],
+        "qwen3_vl_text": ["embed_tokens"],
     }
     # in newer transformers versions, there is one rotary embedding per model
     ROTARY_EMBEDDING = {
         "default": "model.rotary_emb",
         "falcon": "transformer.rotary_emb",
         "gpt_neox": "gpt_neox.rotary_emb",
         "qwen": "transformer.rotary_emb",
+        "qwen3_vl_text": "rotary_emb",
     }
     LM_HEAD = {"default": "lm_head"}
     PRE_HEAD_LAYERNORM = {
         "default": "model.norm",
         "gpt2": "transformer.ln_f",
         "lfm2": "model.embedding_norm",
         "qwen": "transformer.ln_f",
+        "qwen3_vl_text": "norm",
     }
     LAYERS = {
         "default": "model.layers",
@@ -233,6 +236,7 @@ class ModelWrapper:
         "gptj": "transformer.h",
         "opt": "model.decoder.layers",
         "qwen": "transformer.h",
+        "qwen3_vl_text": "layers",
     }
 
     def __init__(self, config: Union[PretrainedConfig, dict]):
@@ -287,7 +291,7 @@ def get_layer_wrappers(self):
 
     def maybe_untie_word_embeddings(self):
         """Untie the word embeddings if they are tied."""
-        if self.config.tie_word_embeddings:
+        if getattr(self.config, "tie_word_embeddings", False):
             self.config.tie_word_embeddings = False
             self.model.config.tie_word_embeddings = False
 

olive/common/mobius_utils.py

@@ -0,0 +1,101 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+"""Helpers for obtaining a model's component plan from mobius.
+
+Mobius owns the per-architecture knowledge of which components a model exposes (e.g. a VLM's
+``decoder`` / ``vision_encoder`` / ``embedding``), how each maps back to a submodule, and the role
+of each component. Olive consumes that plan to drive per-component builds without re-implementing the
+architecture-specific logic.
+
+``mobius-ai`` is imported lazily so Olive keeps working when it is not installed; only the code paths
+that actually need a component plan for a Hugging Face model require it.
+"""
+
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ComponentInfo:
+    """A single component returned by a component source.
+
+    Attributes:
+        name: Stable, user-facing component name used in ``builds.components``.
+        kind: Component role/kind (e.g. ``decoder``, ``vision_encoder``). Optional; used for
+            pass/component compatibility validation.
+        source_path: Dotted submodule path locating the component inside the full model
+            (e.g. ``model.language_model``). Used to slice the component for PyTorch-stage passes.
+
+    """
+
+    name: str
+    kind: Optional[str] = None
+    source_path: Optional[str] = None
+    metadata: dict = field(default_factory=dict)
+
+    @classmethod
+    def coerce(cls, data: "ComponentInfo | dict | object") -> "ComponentInfo":
+        """Normalize a component from any source into an Olive :class:`ComponentInfo`.
+
+        Accepts an existing Olive ``ComponentInfo`` (returned as-is), a mapping following the
+        component contract, or a duck-typed object exposing ``name``/``kind``/``source_path``
+        attributes (e.g. a ``mobius`` ``ComponentInfo`` dataclass).
+        """
+        if isinstance(data, cls):
+            return data
+        if isinstance(data, dict):
+            source = data.get("source") or {}
+            return cls(
+                name=data["name"],
+                kind=data.get("kind"),
+                source_path=data.get("source_path") or source.get("path"),
+                metadata={k: v for k, v in data.items() if k not in ("name", "kind", "source", "source_path")},
+            )
+        return cls(
+            name=data.name,
+            kind=getattr(data, "kind", None),
+            source_path=getattr(data, "source_path", None),
+        )
+
+
+def inspect_components(
+    model_name_or_path: str,
+    task: Optional[str] = None,
+    trust_remote_code: bool = False,
+) -> list[ComponentInfo]:
+    """Return the component plan for a Hugging Face model by querying mobius.
+
+    Args:
+        model_name_or_path: Hugging Face model id or local path.
+        task: Optional task hint passed to mobius.
+        trust_remote_code: Whether to trust remote code when mobius loads the config.
+
+    Returns:
+        A list of :class:`ComponentInfo`. An empty list means the model is single-component
+        (no separable components).
+
+    Raises:
+        ImportError: If ``mobius-ai`` is not installed.
+
+    """
+    try:
+        import mobius
+    except ImportError as exc:
+        raise ImportError(
+            "mobius-ai is required to resolve model components for a Hugging Face model. "
+            "Install with: pip install mobius-ai"
+        ) from exc
+
+    raw_components = mobius.inspect_components(
+        model_name_or_path,
+        task=task,
+        trust_remote_code=trust_remote_code,
+    )
+    components = [ComponentInfo.coerce(c) for c in raw_components]
+    logger.debug("mobius.inspect_components(%s) -> %s", model_name_or_path, [c.name for c in components])
+    return components

olive/engine/config.py

@@ -19,6 +19,9 @@
 # list of all pruned configs
 PRUNED_CONFIGS = (FAILED_CONFIG, INVALID_CONFIG)
 
+# sentinel key inside `builds` that holds partial defaults applied to all sibling builds
+BUILD_DEFAULT_KEY = "_default"
+
 
 class EngineConfig(ConfigBase):
     model_config = ConfigDict(extra="forbid")
@@ -90,3 +93,88 @@ class RunPassConfig(AbstractPassConfig):
             " If not provided, use the engine's evaluator."
         ),
     )
+
+
+class BuildConfigPartial(ConfigBase):
+    """Partial build configuration.
+
+    All fields are optional. Used as the schema for the ``_default`` sentinel inside ``builds``
+    and as the unmerged form of every sibling entry before defaults are applied.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    components: Optional[list[str]] = Field(
+        None,
+        description=(
+            "Names of input model components this build operates on. Each name must match an entry in the input"
+            " model's ``model_component_names``. When omitted, the build runs on the full input model."
+            " When a single name is given, the build receives the unwrapped component handler instead of a one-element"
+            " composite."
+        ),
+    )
+    pipeline: Optional[list[str]] = Field(
+        None,
+        description=(
+            "Ordered list of pass names (referencing entries in the top-level ``passes`` dict) that form this build's"
+            " pipeline."
+        ),
+    )
+    output_dir: Optional[str] = Field(
+        None,
+        description="Directory where this build's final model artifacts get saved.",
+    )
+    host: Optional[Union[SystemConfig, str]] = Field(
+        None,
+        description=(
+            "Host system override for this build. If a string, must refer to a system config under ``systems``."
+            " If omitted, the engine's host is used."
+        ),
+    )
+    target: Optional[Union[SystemConfig, str]] = Field(
+        None,
+        description=(
+            "Target system override for this build. If a string, must refer to a system config under ``systems``."
+            " If omitted, the engine's target is used."
+        ),
+    )
+    evaluator: Optional[Union[OliveEvaluatorConfig, str]] = Field(
+        None,
+        description=(
+            "Evaluator override for this build. If a string, must refer to an evaluator config under ``evaluators``."
+            " If omitted, the engine's evaluator is used."
+        ),
+    )
+    search_strategy: Optional[Union[SearchStrategyConfig, bool]] = Field(
+        None,
+        description="Search strategy override for this build. If omitted, the engine's search strategy is used.",
+    )
+
+
+class BuildConfig(BuildConfigPartial):
+    """Full build configuration after defaults have been merged.
+
+    ``pipeline`` and ``output_dir`` are required post-merge; the other fields remain optional and
+    fall back to the engine-level configuration when not provided.
+    """
+
+    pipeline: list[str] = Field(
+        ...,
+        description=(
+            "Ordered list of pass names (referencing entries in the top-level ``passes`` dict) that form this build's"
+            " pipeline."
+        ),
+    )
+    output_dir: str = Field(
+        ...,
+        description="Directory where this build's final model artifacts get saved.",
+    )
+
+
+def merge_build_default(default_partial: dict, sibling: dict) -> dict:
+    """Merge ``_default`` partial values into a sibling build dict.
+
+    Sibling values fully override default values (no deep merge). Returns a new dict; inputs are
+    not mutated.
+    """
+    return {**{k: v for k, v in default_partial.items() if v is not None}, **sibling}