Add multi build support

mcp/src/olive_mcp/worker.py

@@ -30,6 +30,11 @@ def _get_model_size_mb(model_path: str) -> float | None:
 
 def serialize_workflow_output(result):
     """Convert a WorkflowOutput (or None) into a JSON-serializable dict."""
+    if isinstance(result, dict):
+        return {
+            "status": "success",
+            "builds": {build_name: serialize_workflow_output(output) for build_name, output in result.items()},
+        }
     if result is None:
         return {"status": "success", "output_models": []}
 
@@ -240,11 +245,11 @@ def _handle_explore_passes(kwargs):
 
 def _handle_validate_config(kwargs):
     """Validate an Olive workflow config."""
-    from olive.workflows.run.config import RunConfig
+    from olive.workflows.run.builds import parse_run_config
 
     config = kwargs.get("config")
     try:
-        RunConfig.parse_file_or_obj(config)
+        parse_run_config(config)
         return {"valid": True, "message": "Config is valid."}
     except Exception as e:
         errors = []

olive/cache.py

@@ -8,6 +8,8 @@
 import os
 import re
 import shutil
+from contextlib import contextmanager
+from contextvars import ContextVar
 from copy import deepcopy
 from dataclasses import asdict, dataclass
 from pathlib import Path
@@ -28,6 +30,27 @@
 logger = logging.getLogger(__name__)
 
 SHARED_CACHE_PATTERN = r"https://([^.]+)\.blob\.core\.windows\.net/([^/]+)"
+_CACHE_DIR_CONTEXT: ContextVar[Optional[str]] = ContextVar("olive_cache_dir", default=None)
+_ISOLATED_CACHE_ENV: ContextVar[bool] = ContextVar("olive_isolated_cache_env", default=False)
+
+
+def get_cache_dir_from_env() -> Optional[str]:
+    """Get the cache directory from the current execution context or process environment."""
+    if _ISOLATED_CACHE_ENV.get():
+        return _CACHE_DIR_CONTEXT.get()
+    return os.environ.get("OLIVE_CACHE_DIR")
+
+
+@contextmanager
+def isolated_cache_env(cache_dir: Union[str, Path]):
+    """Scope the cache directory to the current context without changing the process environment."""
+    isolation_token = _ISOLATED_CACHE_ENV.set(True)
+    cache_dir_token = _CACHE_DIR_CONTEXT.set(str(cache_dir))
+    try:
+        yield
+    finally:
+        _CACHE_DIR_CONTEXT.reset(cache_dir_token)
+        _ISOLATED_CACHE_ENV.reset(isolation_token)
 
 
 def is_shared_cache_dir(s) -> bool:
@@ -172,7 +195,7 @@ def get_run_json(
     @classmethod
     def from_cache_env(cls) -> "OliveCache":
         """Create an OliveCache object from the cache directory environment variable."""
-        cache_dir = os.environ.get("OLIVE_CACHE_DIR")
+        cache_dir = get_cache_dir_from_env()
         if cache_dir is None:
             logger.debug("OLIVE_CACHE_DIR environment variable not set. Using default cache directory.")
             cache_dir = Path(DEFAULT_CACHE_DIR).resolve() / DEFAULT_WORKFLOW_ID
@@ -295,7 +318,11 @@ def get_cache_dir(self) -> Path:
 
     def set_cache_env(self):
         """Set environment variable for the cache directory."""
-        os.environ["OLIVE_CACHE_DIR"] = str(self.dirs.cache_dir)
+        cache_dir = str(self.dirs.cache_dir)
+        if _ISOLATED_CACHE_ENV.get():
+            _CACHE_DIR_CONTEXT.set(cache_dir)
+        else:
+            os.environ["OLIVE_CACHE_DIR"] = cache_dir
         logger.debug("Set OLIVE_CACHE_DIR: %s", self.dirs.cache_dir)
 
     def prepare_resources_for_local(self, config: Union[dict, ConfigBase]) -> Union[dict, ConfigBase]:
@@ -396,9 +423,8 @@ 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, or multimodal ORT GenAI packages from
-                # MobiusBuilder where components live in <root>/<component>/).
+                # Preserve directory structure for packages whose components live
+                # in separate subdirectories under a shared root.
                 source_path = Path(model_json_config["model_path"])
                 source_path_resolved = source_path.resolve()
                 if source_path.exists():

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

@@ -343,7 +343,7 @@ def _run_workflow(self):
                 mark_test_output_path(self.args.output_path)
                 save_discrepancy_check_results(workflow_output, self.args.output_path)
                 print(f"Test report saved at {self.args.output_path}")
-            elif not workflow_output.has_output_model():
+            if 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/run.py

@@ -62,6 +62,13 @@ def run(self):
         run_config = self.args.run_config
         if not isinstance(run_config, dict):
             run_config = load_config_file(run_config)
+        if "builds" in run_config and self.args.test not in (None, False):
+            raise ValueError("--test is not supported with multi-build run configurations.")
+        if "builds" in run_config and self.args.output_path is not None:
+            raise ValueError(
+                "--output_path is not supported with multi-build run configurations. "
+                "Set output_dir on each build instead."
+            )
         if input_model_config := get_input_model_config(self.args, required=False):
             print("Replacing input model config in run config")
             run_config["input_model"] = input_model_config
@@ -98,5 +105,18 @@ def run(self):
 
         if self.args.list_required_packages is True:
             print("Required packages listed!")
+        elif isinstance(workflow_output, dict):
+            self._print_build_outputs(run_config, workflow_output)
 
         return workflow_output
+
+    @staticmethod
+    def _print_build_outputs(run_config: dict, workflow_outputs: dict) -> None:
+        builds = run_config.get("builds") or {}
+        build_default = builds.get("_default") or {}
+        for build_name, workflow_output in workflow_outputs.items():
+            if workflow_output is None or not workflow_output.has_output_model():
+                print(f"Build {build_name!r}: no output model produced. Please check the log for details.")
+                continue
+            output_dir = (builds.get(build_name) or {}).get("output_dir") or build_default.get("output_dir")
+            print(f"Build {build_name!r}: model is saved under {output_dir}")

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, AutoModelForSeq2SeqLM, 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
@@ -194,16 +194,21 @@ def load_model_from_task(
     **kwargs,
 ) -> "PreTrainedModel":
     """Load huggingface model from task and model_name_or_path."""
-    from transformers.pipelines import check_task
-
-    task_results = check_task(task.replace("-with-past", ""))
-    assert isinstance(task_results, tuple)
-    if len(task_results) == 2:
-        targeted_task = task_results[0]
-    elif len(task_results) == 3:
-        targeted_task = task_results[1]
+    task_without_past = task.replace("-with-past", "")
+    if task_without_past == "text2text-generation":
+        class_tuple = (AutoModelForSeq2SeqLM,)
     else:
-        raise ValueError("unsupported transformers version")
+        from transformers.pipelines import check_task
+
+        task_results = check_task(task_without_past)
+        assert isinstance(task_results, tuple)
+        if len(task_results) == 2:
+            targeted_task = task_results[0]
+        elif len(task_results) == 3:
+            targeted_task = task_results[1]
+        else:
+            raise ValueError("unsupported transformers version")
+        class_tuple = targeted_task["pt"] or (AutoModel,)
 
     model_config = get_model_config(model_name_or_path, test_model_config=test_model_config, **kwargs)
     if getattr(model_config, "quantization_config", None):
@@ -229,7 +234,6 @@ def load_model_from_task(
             AUTO_QUANTIZATION_CONFIG_MAPPING["olive"] = OliveHfQuantizationConfig
             AUTO_QUANTIZER_MAPPING["olive"] = OliveHfQuantizer
 
-    class_tuple = targeted_task["pt"] or (AutoModel,)
     model = None
     for i, model_class in enumerate(class_tuple):
         try:
@@ -427,6 +431,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

@@ -3,7 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 import logging
-from typing import TYPE_CHECKING, Callable, Union
+from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import torch
 from torch import nn
@@ -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]):
@@ -250,6 +254,9 @@ def __init__(self, config: Union[PretrainedConfig, dict]):
 
         self._model = None
         self._layer_wrappers = None
+        self.olive_root_model: Optional[nn.Module] = None
+        self.olive_component_path: Optional[str] = None
+        self.olive_component_role: Optional[str] = None
 
     @property
     def model(self) -> "PreTrainedModel":
@@ -258,9 +265,13 @@ def model(self) -> "PreTrainedModel":
 
         return self._model
 
-    def set_model(self, model: "PreTrainedModel"):
+    def set_model(self, model: "PreTrainedModel", initialize_layer_wrappers: bool = True):
         self._model = model
-        self._layer_wrappers = [LayerWrapper(layer, self.model_type) for layer in self.get_layers(False)]
+        self._layer_wrappers = (
+            [LayerWrapper(layer, self.model_type) for layer in self.get_layers(False)]
+            if initialize_layer_wrappers
+            else []
+        )
 
     def get_embeds(self, return_name: bool = True):
         return get_submodules(self.model, self.EMBEDDINGS, self.model_type, return_name=return_name)
@@ -287,7 +298,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/quant/hf_utils.py

@@ -195,6 +195,30 @@ def create_quantized_module(module: nn.Linear | nn.Embedding, name: str) -> Quan
 
         replace_matching_submodules(model, should_quantize, create_quantized_module)
 
+        tied_weight_keys = getattr(model, "all_tied_weights_keys", None)
+        if tied_weight_keys:
+
+            def has_parameter(name: str) -> bool:
+                try:
+                    model.get_parameter(name)
+                    return True
+                except AttributeError:
+                    return False
+
+            # Transformers caches this mapping before the quantizer replaces float weights with quantized buffers.
+            model.all_tied_weights_keys = {
+                target: source
+                for target, source in tied_weight_keys.items()
+                if has_parameter(target) and has_parameter(source)
+            }
+
+        if isinstance(model.get_input_embeddings(), (QuantEmbedding, QuantLinear)) or isinstance(
+            model.get_output_embeddings(), (QuantEmbedding, QuantLinear)
+        ):
+            # Model-specific root initializers such as T5 dereference embedding.weight after child modules are loaded.
+            # Quantized embeddings have no float weight; child modules still retain their own initialization state.
+            model._is_hf_initialized = True  # pylint: disable=protected-access
+
         if self.quantization_config.tie_word_embeddings:
             # doing first time so that the weight load doesn't complain about missing weights
             tie_quant_word_embeddings(model)