diff --git a/mcp/src/olive_mcp/worker.py b/mcp/src/olive_mcp/worker.py index 8dbad076e0..9a3dd496cd 100644 --- a/mcp/src/olive_mcp/worker.py +++ b/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 = [] diff --git a/olive/cache.py b/olive/cache.py index 7f27a50de1..86a6cd164e 100644 --- a/olive/cache.py +++ b/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 //). + # 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(): diff --git a/olive/cli/api.py b/olive/cli/api.py index cae2963264..fbbcdcb90f 100644 --- a/olive/cli/api.py +++ b/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,7 +300,7 @@ 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: @@ -308,7 +308,8 @@ def run(run_config: str, **kwargs) -> WorkflowOutput: **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 diff --git a/olive/cli/base.py b/olive/cli/base.py index a64a9bf473..85cf0ed316 100644 --- a/olive/cli/base.py +++ b/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}") diff --git a/olive/cli/run.py b/olive/cli/run.py index a0a91e8115..3111066452 100644 --- a/olive/cli/run.py +++ b/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}") diff --git a/olive/common/hf/utils.py b/olive/common/hf/utils.py index 69c8b3b7af..a1c3ac5634 100644 --- a/olive/common/hf/utils.py +++ b/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) diff --git a/olive/common/hf/wrapper.py b/olive/common/hf/wrapper.py index 8bde44ef58..b73ef50b82 100644 --- a/olive/common/hf/wrapper.py +++ b/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,6 +209,7 @@ 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 = { @@ -216,6 +217,7 @@ class ModelWrapper: "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 = { @@ -223,6 +225,7 @@ class ModelWrapper: "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 diff --git a/olive/common/quant/hf_utils.py b/olive/common/quant/hf_utils.py index c1e4c5065b..a586e19138 100644 --- a/olive/common/quant/hf_utils.py +++ b/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) diff --git a/olive/model/config/model_config.py b/olive/model/config/model_config.py index d6eac90043..c30d27811a 100644 --- a/olive/model/config/model_config.py +++ b/olive/model/config/model_config.py @@ -5,6 +5,7 @@ import logging from copy import deepcopy from pathlib import Path +from typing import Optional from pydantic import Field, field_validator @@ -43,6 +44,188 @@ def create_model(self): cls = get_model_handler(self.type) return cls(**self.config) + def get_components(self) -> Optional[list[str]]: + """Return the component names that builds can target, or None for a single-component model. + + * ``CompositeModel`` -> its configured ``model_component_names``, or, when the config only + points at a directory of per-component ONNX subfolders, the discovered subfolder names. + * ``HfModel`` -> the components reported by ``mobius.inspect_components``, or ``None`` when + Mobius reports no components. + * ``DiffusersModel`` -> its exportable components. + * Anything else -> ``None`` (single-component model). + """ + if self.type == "compositemodel": + names = list(self.config.get("model_component_names") or []) + if names: + return names + return [name for name, _ in self._discover_composite_components()] + if self.type == "hfmodel": + return self._get_hf_components() or None + if self.type == "diffusersmodel": + return self._get_diffusers_components() or None + return None + + def _discover_composite_components(self) -> list[tuple[str, str]]: + """Discover ``(name, onnx_relpath)`` from a directory-based composite, or empty list.""" + from olive.model.utils.onnx_utils import discover_onnx_components + + model_path = self.config.get("model_path") + if not model_path or not Path(str(model_path)).is_dir(): + return [] + return discover_onnx_components(str(model_path)) + + def _get_hf_components(self) -> list[str]: + """Return component names for an HfModel by querying Mobius.""" + from olive.common.mobius_utils import inspect_components + + model_path = self.config.get("model_path") + if not model_path: + return [] + load_kwargs = self.config.get("load_kwargs") or {} + model_attributes = self.config.get("model_attributes") or {} + components = inspect_components( + model_path, + task=model_attributes.get("mobius_task"), + trust_remote_code=bool(load_kwargs.get("trust_remote_code")), + ) + return [component.name for component in components] + + def _get_diffusers_components(self) -> list[str]: + """Return the exportable component names for a DiffusersModel, or empty list. + + Reads the variant's component layout from the handler (which only inspects config files, + not weights). When the config is already scoped to a subset via ``components``, returns that + subset; the top-level input model carries no such scope, so ``builds`` sees the full set. + """ + model_path = self.config.get("model_path") + if not model_path: + return [] + handler = self.create_model() + return [str(c) for c in handler.get_exportable_components()] + + def select_components(self, names: list[str]) -> "ModelConfig": + """Return a new ModelConfig holding only the named components. + + * ``CompositeModel`` -> the unwrapped child component ``ModelConfig`` when exactly one name + is given, otherwise a new ``CompositeModel`` ``ModelConfig`` containing the subset (in the + requested order). Directory-based composites are discovered first. + * ``HfModel`` -> a copy tagged with the selected component's Mobius role and source paths. + * ``DiffusersModel`` -> a copy scoped to the selected exportable components. + + Raises ``ValueError`` if any name is missing from the available components. + """ + if self.type == "hfmodel": + return self._select_hf_component(names) + if self.type == "diffusersmodel": + return self._select_diffusers_components(names) + if self.type != "compositemodel": + raise ValueError( + "select_components is only supported on CompositeModel, HfModel, or DiffusersModel " + f"input configs (got type {self.type!r})." + ) + if not names: + raise ValueError("select_components requires a non-empty list of names.") + component_names = list(self.config.get("model_component_names") or []) + model_components = list(self.config.get("model_components") or []) + if not component_names: + discovered = self._discover_composite_components() + if not discovered: + raise ValueError( + "CompositeModel config has no model_components and model_path is not a directory of " + "per-component ONNX subfolders." + ) + component_names = [name for name, _ in discovered] + model_path = self.config.get("model_path") + model_components = [ + {"type": "ONNXModel", "config": {"model_path": str(model_path), "onnx_file_name": onnx_rel}} + for _, onnx_rel in discovered + ] + if len(component_names) != len(model_components): + raise ValueError("CompositeModel config has mismatched model_components and model_component_names lengths.") + missing = [n for n in names if n not in component_names] + if missing: + raise ValueError(f"Unknown component name(s) {missing}. Available components: {list(component_names)}.") + component_map = dict(zip(component_names, model_components)) + selected = [deepcopy(component_map[n]) for n in names] + if len(selected) == 1: + # Unwrap to the child handler config, inheriting the composite's shared model_attributes so a + # single-component build keeps parent context (matches CompositeModelHandler.select_components). + child = selected[0] + parent_attributes = self.config.get("model_attributes") or {} + if isinstance(child, ModelConfig): + merged = {**parent_attributes, **(child.config.get("model_attributes") or {})} + if merged: + child.config["model_attributes"] = merged + return child + child_config = dict(child.get("config") or {}) + merged = {**parent_attributes, **(child_config.get("model_attributes") or {})} + if merged: + child_config["model_attributes"] = merged + return ModelConfig.model_validate({**child, "config": child_config}) + new_config = { + **{k: v for k, v in self.config.items() if k not in ("model_components", "model_component_names")}, + "model_components": selected, + "model_component_names": list(names), + } + return ModelConfig(type=self.type, config=new_config) + + def _select_hf_component(self, names: list[str]) -> "ModelConfig": + """Select one HfModel component for component-scoped PyTorch optimization.""" + if not names: + raise ValueError("select_components requires a non-empty list of names.") + if len(names) != 1: + raise ValueError( + f"HfModel components must be optimized one at a time; got {names}. Use a separate build per component." + ) + + from olive.common.mobius_utils import inspect_components + + model_path = self.config.get("model_path") + load_kwargs = self.config.get("load_kwargs") or {} + model_attributes = self.config.get("model_attributes") or {} + components = inspect_components( + model_path, + task=model_attributes.get("mobius_task"), + trust_remote_code=bool(load_kwargs.get("trust_remote_code")), + ) + components_by_name = {component.name: component for component in components} + missing = [name for name in names if name not in components_by_name] + if missing: + raise ValueError(f"Unknown component name(s) {missing}. Available components: {list(components_by_name)}.") + component = components_by_name[names[0]] + if not component.source_paths and len(components) > 1: + raise ValueError( + f"Component {component.name!r} has no runtime source paths, so Olive cannot safely " + "scope PyTorch passes to it." + ) + + new_config = deepcopy(self.config) + attributes = dict(new_config.get("model_attributes") or {}) + attributes["component_name"] = component.name + if component.role is not None: + attributes["component_role"] = component.role + if component.source_paths: + attributes["component_source_paths"] = list(component.source_paths) + new_config["model_attributes"] = attributes + return ModelConfig(type=self.type, config=new_config) + + def _select_diffusers_components(self, names: list[str]) -> "ModelConfig": + """Scope a DiffusersModel to the named exportable components. + + Returns a copy of the config with ``components`` set to the requested subset (in the + variant's canonical order). The scoped handler exports only those components, so a build's + conversion pass produces just that component's ONNX while subsequent passes map over it. + """ + if not names: + raise ValueError("select_components requires a non-empty list of names.") + available = self._get_diffusers_components() + missing = [n for n in names if n not in available] + if missing: + raise ValueError(f"Unknown component name(s) {missing}. Available components: {available}.") + new_config = deepcopy(self.config) + new_config["components"] = [name for name in available if name in set(names)] + return ModelConfig(type=self.type, config=new_config) + def get_model_id(self): for v in self.config.values(): if callable(v): diff --git a/olive/model/handler/composite.py b/olive/model/handler/composite.py index c52bd1a315..42d22eda4f 100644 --- a/olive/model/handler/composite.py +++ b/olive/model/handler/composite.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging +from pathlib import Path from typing import Any, Optional, Union from olive.common.config_utils import serialize_to_json, validate_config @@ -32,8 +33,8 @@ class CompositeModelHandler(OliveModelHandler): def __init__( self, - model_components: list[Union[OliveModelHandler, dict[str, Any]]], - model_component_names: list[str], + model_components: Optional[list[Union[OliveModelHandler, dict[str, Any]]]] = None, + model_component_names: Optional[list[str]] = None, model_path: OLIVE_RESOURCE_ANNOTATIONS = None, model_attributes: Optional[dict[str, Any]] = None, ): @@ -43,16 +44,57 @@ def __init__( model_file_format=ModelFileFormat.COMPOSITE_MODEL, model_attributes=model_attributes, ) + + # When components are not provided but model_path is a directory of per-component ONNX + # subfolders, discover them using the subfolder names as component names. + if model_components is None: + discovered = self._discover_components(model_path) + if not discovered: + raise ValueError( + "CompositeModelHandler requires model_components, or a model_path directory containing " + "per-component ONNX subfolders." + ) + model_components, model_component_names = discovered + + if model_component_names is None: + raise ValueError("CompositeModelHandler requires model_component_names when model_components is provided.") + self._model_components = [ validate_config(m, ModelConfig).create_model() if isinstance(m, dict) else m for m in model_components ] - assert all(isinstance(m, OliveModelHandler) for m in self._model_components), ( - "All components must be OliveModelHandler or dict" - ) + if not all(isinstance(m, OliveModelHandler) for m in self._model_components): + raise ValueError("All components must be OliveModelHandler or dict.") - assert len(self._model_components) == len(model_component_names), "Number of components and names must match" + if len(self._model_components) != len(model_component_names): + raise ValueError( + f"Number of components ({len(self._model_components)}) and names " + f"({len(model_component_names)}) must match." + ) self.model_component_names = model_component_names + @staticmethod + def _discover_components( + model_path: OLIVE_RESOURCE_ANNOTATIONS, + ) -> Optional[tuple[list[dict[str, Any]], list[str]]]: + """Build component configs from a directory of per-component ONNX subfolders. + + Returns ``(model_components, model_component_names)`` or ``None`` if discovery is not + applicable (model_path is not a local directory of component subfolders). + """ + from olive.model.utils.onnx_utils import discover_onnx_components + + if not model_path or not Path(str(model_path)).is_dir(): + return None + discovered = discover_onnx_components(str(model_path)) + if not discovered: + return None + names = [name for name, _ in discovered] + components = [ + {"type": "ONNXModel", "config": {"model_path": str(model_path), "onnx_file_name": onnx_rel}} + for _, onnx_rel in discovered + ] + return components, names + @property def model_components(self): for m in self._model_components: @@ -77,6 +119,33 @@ def to_json(self, check_object: bool = False): def get_model_components(self) -> list[tuple[str, OliveModelHandler]]: return zip(self.model_component_names, self.model_components) + def select_components(self, names: list[str]) -> "OliveModelHandler": + """Return a handler holding only the named components. + + Returns the unwrapped child handler when exactly one name is given; returns a new + ``CompositeModelHandler`` containing the subset (in the requested order) otherwise. + Raises ``ValueError`` if any name is missing from ``model_component_names``. + """ + if not names: + raise ValueError("select_components requires a non-empty list of names.") + missing = [n for n in names if n not in self.model_component_names] + if missing: + raise ValueError( + f"Unknown component name(s) {missing}. Available components: {list(self.model_component_names)}." + ) + component_map = dict(zip(self.model_component_names, self._model_components)) + selected = [component_map[n] for n in names] + if len(selected) == 1: + child = selected[0] + child.model_attributes = {**(self.model_attributes or {}), **(child.model_attributes or {})} + return child + return CompositeModelHandler( + model_components=selected, + model_component_names=list(names), + model_path=self.model_path, + model_attributes=self.model_attributes, + ) + def load_model(self, rank: int = None, cache_model: bool = True): raise NotImplementedError diff --git a/olive/model/handler/diffusers.py b/olive/model/handler/diffusers.py index e7f13f4f9b..da119f7ba4 100644 --- a/olive/model/handler/diffusers.py +++ b/olive/model/handler/diffusers.py @@ -34,7 +34,7 @@ class DiffusersModelHandler(OliveModelHandler): """ resource_keys: tuple[str, ...] = ("model_path", "adapter_path") - json_config_keys: tuple[str, ...] = ("model_variant", "load_kwargs") + json_config_keys: tuple[str, ...] = ("model_variant", "load_kwargs", "components") def __init__( self, @@ -42,6 +42,7 @@ def __init__( model_variant: Union[str, DiffusersModelVariant] = DiffusersModelVariant.AUTO, load_kwargs: Optional[dict[str, Any]] = None, adapter_path: OLIVE_RESOURCE_ANNOTATIONS = None, + components: Optional[list[str]] = None, model_attributes: Optional[dict[str, Any]] = None, ): """Initialize DiffusersModelHandler. @@ -51,6 +52,10 @@ def __init__( model_variant: Model variant: 'sd15', 'sdxl', 'flux', or 'auto' for auto-detection. load_kwargs: Additional kwargs for loading the model (e.g., torch_dtype, variant). adapter_path: Path to LoRA adapter weights. + components: Optional subset of exportable component names this handler is scoped to + (e.g. ``["text_encoder"]``). When set, ``get_exportable_components`` returns only + these (in variant order); used by ``builds.components`` to optimize one component at + a time. When ``None``, all of the variant's components are exportable. model_attributes: Additional model attributes. """ @@ -67,6 +72,7 @@ def __init__( self.model_variant = DiffusersModelVariant(model_variant) self.load_kwargs = load_kwargs or {} + self.components = list(components) if components else None self._pipeline = None @property @@ -307,7 +313,18 @@ def get_exportable_components(self) -> list[DC]: } if variant not in variant_components: raise ValueError(f"Unknown model variant: {variant}") - return variant_components[variant] + full = variant_components[variant] + if self.components: + requested = list(self.components) + available = {str(c) for c in full} + unknown = [name for name in requested if name not in available] + if unknown: + raise ValueError( + f"Unknown component(s) {unknown} for variant {variant}. Available: {sorted(available)}." + ) + # preserve the variant's canonical component order + return [c for c in full if str(c) in set(requested)] + return full def get_pipeline_type(self) -> DiffusersModelVariant: """Get the pipeline type for OnnxConfig lookup. diff --git a/olive/model/handler/mixin/hf.py b/olive/model/handler/mixin/hf.py index a5e51c162b..464c6a98b7 100644 --- a/olive/model/handler/mixin/hf.py +++ b/olive/model/handler/mixin/hf.py @@ -10,9 +10,11 @@ from olive.common.hf.utils import ( get_generation_config, get_model_config, + get_processor, get_tokenizer, save_model_config, save_module_files, + save_processor, save_tokenizer, ) @@ -64,6 +66,11 @@ def get_hf_tokenizer(self) -> Union["PreTrainedTokenizer", "PreTrainedTokenizerF # TODO(anyone): only provide relevant kwargs, no use case for now to provide kwargs return get_tokenizer(self.model_path) + def get_hf_processor(self): + """Get processor for the model if one exists.""" + processor_exclude_keys = {"torch_dtype", "dtype", "device_map", "max_memory", "quantization_config"} + return get_processor(self.model_path, **self.get_load_kwargs(list(processor_exclude_keys))) + def save_metadata(self, output_dir: str, exclude_load_keys: Optional[list[str]] = None, **kwargs) -> list[str]: """Save model metadata files to the output directory. @@ -117,6 +124,15 @@ def save_metadata(self, output_dir: str, exclude_load_keys: Optional[list[str]] tokenizer_filepaths = save_tokenizer(self.get_hf_tokenizer(), output_dir, **kwargs) saved_filepaths.extend([fp for fp in tokenizer_filepaths if Path(fp).exists()]) + should_save_processor = (self.task or "").replace("-with-past", "") in {"image-text-to-text"} + if should_save_processor and not any( + (output_dir / name).exists() for name in ("preprocessor_config.json", "processor_config.json") + ): + processor = self.get_hf_processor() + if processor is not None: + processor_filepaths = save_processor(processor, output_dir, **kwargs) + saved_filepaths.extend([fp for fp in processor_filepaths if Path(fp).exists()]) + logger.debug("Save metadata files to %s: %s", output_dir, saved_filepaths) return saved_filepaths diff --git a/olive/model/utils/onnx_utils.py b/olive/model/utils/onnx_utils.py index 73be98a115..3bbed5ff70 100644 --- a/olive/model/utils/onnx_utils.py +++ b/olive/model/utils/onnx_utils.py @@ -63,6 +63,40 @@ def get_onnx_file_path(model_path: str, onnx_file_name: Optional[str] = None) -> raise ValueError(f"No .onnx file found in the model folder {model_path}.") +def discover_onnx_components(model_dir: str) -> list[tuple[str, str]]: + """Discover per-component ONNX subfolders in a directory. + + A multi-component ONNX package lays out each component in its own subfolder, with a + ``model.onnx`` inside: + + model_dir/decoder/model.onnx + model_dir/vision_encoder/model.onnx + model_dir/embedding/model.onnx + + Args: + model_dir: Directory that contains one subfolder per component. + + Returns: + A list of ``(component_name, onnx_file_relpath)`` tuples sorted by component name, where + ``component_name`` is the subfolder name and ``onnx_file_relpath`` is the path to the + component's ``.onnx`` file relative to ``model_dir``. Empty if no component subfolders are + found. + + """ + model_dir_path = Path(model_dir) + if not model_dir_path.is_dir(): + return [] + + components: list[tuple[str, str]] = [] + for sub_dir in sorted(p for p in model_dir_path.iterdir() if p.is_dir()): + onnx_files = list(sub_dir.glob("*.onnx")) + if len(onnx_files) == 1: + components.append((sub_dir.name, f"{sub_dir.name}/{onnx_files[0].name}")) + elif (sub_dir / "model.onnx").exists(): + components.append((sub_dir.name, f"{sub_dir.name}/model.onnx")) + return components + + def get_additional_file_path(model_dir: str, file_name: str) -> Optional[str]: """Get the full path to the additional file. diff --git a/olive/passes/openvino/optimum_intel.py b/olive/passes/openvino/optimum_intel.py index d898e665eb..8e11cf2db5 100644 --- a/olive/passes/openvino/optimum_intel.py +++ b/olive/passes/openvino/optimum_intel.py @@ -4,8 +4,11 @@ # -------------------------------------------------------------------------- import logging import os +import tempfile +from contextlib import contextmanager from copy import deepcopy from pathlib import Path +from threading import Lock from typing import Any, Optional, Union from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE @@ -18,6 +21,24 @@ from olive.passes.pass_config import BasePassConfig, PassConfigParam, get_user_script_data_config logger = logging.getLogger(__name__) +_TEMP_DIR_LOCK = Lock() + + +@contextmanager +def _use_output_tempdir(output_model_path: str): + with _TEMP_DIR_LOCK: + original_tmpdir = os.environ.get("TMPDIR") + original_tempdir = tempfile.tempdir + os.environ["TMPDIR"] = output_model_path + tempfile.tempdir = output_model_path + try: + yield + finally: + tempfile.tempdir = original_tempdir + if original_tmpdir is None: + os.environ.pop("TMPDIR", None) + else: + os.environ["TMPDIR"] = original_tmpdir def infer_task( @@ -501,40 +522,28 @@ def _run_for_config( # Workaround for optimum-intel using Path.rename() which fails across filesystems. # Set tempdir to output path so temp files are on the same filesystem as the cache. - import tempfile - Path(output_model_path).mkdir(parents=True, exist_ok=True) - original_tmpdir = os.environ.get("TMPDIR") - original_tempdir = tempfile.tempdir - os.environ["TMPDIR"] = output_model_path - tempfile.tempdir = output_model_path - - export_optimum_intel( - model.model_name_or_path, - output_model_path, - **extra_args, - ) - if apply_main_quantize: - _main_quantize( - model_name_or_path=model.model_name_or_path, - task=extra_args.get("task", "auto"), - library_name=lib_name, - quantization_config=quant_config, - output=Path(output_model_path), - cache_dir=config.ov_quant_config.get("cache_dir", None) if config.ov_quant_config else None, - trust_remote_code=config.ov_quant_config.get("trust_remote_code", False) - if config.ov_quant_config - else False, - model_kwargs=model.load_kwargs.__dict__ if model.load_kwargs else None, + with _use_output_tempdir(output_model_path): + export_optimum_intel( + model.model_name_or_path, + output_model_path, + **extra_args, ) + if apply_main_quantize: + _main_quantize( + model_name_or_path=model.model_name_or_path, + task=extra_args.get("task", "auto"), + library_name=lib_name, + quantization_config=quant_config, + output=Path(output_model_path), + cache_dir=config.ov_quant_config.get("cache_dir", None) if config.ov_quant_config else None, + trust_remote_code=config.ov_quant_config.get("trust_remote_code", False) + if config.ov_quant_config + else False, + model_kwargs=model.load_kwargs.__dict__ if model.load_kwargs else None, + ) except Exception as e: raise RuntimeError(f"OpenVINO optimum export failed: {e}") from e - finally: - tempfile.tempdir = original_tempdir - if original_tmpdir is None: - os.environ.pop("TMPDIR", None) - else: - os.environ["TMPDIR"] = original_tmpdir # check the exported components exported_models = [name.stem for name in Path(output_model_path).iterdir() if name.suffix == ".xml"] diff --git a/olive/passes/pytorch/quant_utils.py b/olive/passes/pytorch/quant_utils.py index 5ee0ce00d8..5646b29799 100644 --- a/olive/passes/pytorch/quant_utils.py +++ b/olive/passes/pytorch/quant_utils.py @@ -21,7 +21,7 @@ ) from olive.common.quant.nn import QuantEmbedding, QuantLinear from olive.common.quant.utils import WeightQuantizer -from olive.common.utils import tensor_data_to_device +from olive.common.utils import get_attr, set_attr, tensor_data_to_device from olive.constants import PrecisionBits from olive.passes.pass_config import PassConfigParam from olive.passes.pytorch.common import inherit_hf_from_hf @@ -80,7 +80,245 @@ def get_quantizer_config(allow_embeds: bool = False) -> dict[str, PassConfigPara } -def get_qkv_quantization_groups(wrapper: ModelWrapper, module_names: set[str] | None = None) -> list[tuple[str, ...]]: +def _root_module_name(name: str, name_prefix: str = "") -> str: + """Return the module name relative to the saved model root.""" + return f"{name_prefix}{name}" if name else name_prefix.rstrip(".") + + +def _is_in_component(name: str, source_paths: list[str]) -> bool: + """Return True if ``name`` (root-relative) belongs to any of the component's sub-trees.""" + if not source_paths: + return True + return any(name == path or name.startswith(f"{path}.") for path in source_paths) + + +def _get_component_source_paths(model: HfModelHandler) -> list[str]: + """Return the component's dotted sub-module paths from model attributes. + + Reads ``component_source_paths`` from model attributes. Falls back to the legacy singular + ``component_source_path`` string for backward compatibility. + """ + attributes = model.model_attributes or {} + source_paths = attributes.get("component_source_paths") + if source_paths: + return list(source_paths) + legacy = attributes.get("component_source_path") + return [legacy] if legacy else [] + + +def _component_slice_path(source_paths: list[str]) -> str | None: + """Return the sub-module to slice and quantize for the given component paths. + + A single-path component slices to that path directly. A component spanning several + disjoint sub-trees (e.g. ``model.layers`` + ``model.norm`` + ``lm_head``) slices to their + greatest common dotted ancestor (the root when they share none), and ``_is_in_component`` + then restricts quantization to the declared sub-trees within that slice. + """ + if not source_paths: + return None + if len(source_paths) == 1: + return source_paths[0] + split = [path.split(".") for path in source_paths] + common: list[str] = [] + for segments in zip(*split): + if len(set(segments)) == 1: + common.append(segments[0]) + else: + break + return ".".join(common) or None + + +def _path_with_leaf(source_paths: list[str], leaves: set[str]) -> str | None: + for path in source_paths: + if path.rsplit(".", 1)[-1] in leaves: + return path + return None + + +def _module_path(root_model: torch.nn.Module, target: torch.nn.Module | None) -> str | None: + if target is None: + return None + return next((name for name, module in root_model.named_modules() if module is target), None) + + +def _root_component_model_wrapper( + root_model: torch.nn.Module, + backbone_path: str, + layers_path: str, + source_paths: list[str], + embedding_path: str | None = None, +) -> ModelWrapper: + """Create a text wrapper whose structural paths are relative to the full HF model.""" + backbone = get_attr(root_model, backbone_path) if backbone_path else root_model + wrapper = ModelWrapper(backbone.config) + wrapper.LAYERS = {"default": layers_path} + + norm_path = _path_with_leaf( + source_paths, + {"norm", "final_layer_norm", "layer_norm"}, + ) + if norm_path is None: + norm_path = _module_path(root_model, getattr(backbone, "norm", None)) + if norm_path is not None: + wrapper.PRE_HEAD_LAYERNORM = {"default": norm_path} + + head_path = _path_with_leaf( + source_paths, + {"lm_head", "proj_out", "output_projection", "codec_head"}, + ) + if head_path is None: + get_output_embeddings = getattr(root_model, "get_output_embeddings", None) + output_embeddings = get_output_embeddings() if callable(get_output_embeddings) else None + head_path = _module_path(root_model, output_embeddings) + if head_path is not None: + wrapper.LM_HEAD = {"default": head_path} + + if embedding_path is None: + get_input_embeddings = getattr(backbone, "get_input_embeddings", None) + input_embeddings = get_input_embeddings() if callable(get_input_embeddings) else None + if input_embeddings is None: + input_embeddings = getattr(backbone, "embed_tokens", None) + embedding_path = _module_path(root_model, input_embeddings) + if embedding_path is not None: + wrapper.EMBEDDINGS = {"default": [embedding_path]} + + rotary_path = _module_path(root_model, getattr(backbone, "rotary_emb", None)) + if rotary_path is not None: + wrapper.ROTARY_EMBEDDING = {"default": rotary_path} + + wrapper.set_model(root_model) + return wrapper + + +class _GenericComponentWrapper(ModelWrapper): + """Minimal wrapper for non-decoder components that do not expose transformer layers.""" + + def __init__(self, model: torch.nn.Module, config) -> None: + super().__init__(config) + super().set_model(model, initialize_layer_wrappers=False) + + def maybe_untie_word_embeddings(self): + if not getattr(self.config, "tie_word_embeddings", False): + return + + root_model = self.olive_root_model if self.olive_root_model is not None else self.model + + # T5 reuses one Embedding module at several paths, while BART uses + # distinct modules that share one Parameter. Split both forms before + # replacing only the selected component. + embedding_aliases: dict[int, list[tuple[str, torch.nn.Embedding]]] = {} + for name, module in root_model.named_modules(remove_duplicate=False): + if name and isinstance(module, torch.nn.Embedding): + embedding_aliases.setdefault(id(module), []).append((name, module)) + for aliases in embedding_aliases.values(): + for name, module in aliases[1:]: + set_attr(root_model, name, deepcopy(module)) + + tied_weights: dict[int, list[torch.nn.Module]] = {} + for module in root_model.modules(): + if isinstance(module, (torch.nn.Embedding, torch.nn.Linear)): + tied_weights.setdefault(id(module.weight), []).append(module) + for modules in tied_weights.values(): + if len(modules) < 2 or not any(isinstance(module, torch.nn.Embedding) for module in modules): + continue + for module in modules[1:]: + weight = module.weight + module.weight = torch.nn.Parameter( + weight.detach().clone(), + requires_grad=weight.requires_grad, + ) + + for module in root_model.modules(): + module_config = getattr(module, "config", None) + if module_config is not None and hasattr(module_config, "tie_word_embeddings"): + module_config.tie_word_embeddings = False + self.config.tie_word_embeddings = False + + def get_embeds(self, return_name: bool = True): + raise AttributeError("The selected component has no language-model embeddings.") + + def get_lm_head(self, return_name: bool = True): + raise AttributeError("The selected component has no language-model head.") + + +def _component_model_wrapper( + root_model: torch.nn.Module, + source_paths: list[str], + component_role: str | None, +) -> tuple[ModelWrapper, str]: + """Wrap a selected component while retaining paths relative to the saved HF model. + + Decoder components can span disjoint runtime sub-trees (for example a nested + language backbone plus a top-level ``lm_head``). In that case the wrapper must + operate on the full model so replacement and saving preserve every component, + while its structural lookups point at the selected decoder paths. + """ + if not source_paths: + if component_role not in {None, "decoder"}: + return _GenericComponentWrapper(root_model, root_model.config), "" + return ModelWrapper.from_model(root_model), "" + + if component_role not in {None, "decoder", "embedding"}: + slice_path = _component_slice_path(source_paths) + component_model = get_attr(root_model, slice_path) if slice_path else root_model + config = getattr(component_model, "config", root_model.config) + return _GenericComponentWrapper(component_model, config), (f"{slice_path}." if slice_path else "") + + slice_path = _component_slice_path(source_paths) + embedding_path = _path_with_leaf(source_paths, {"embed_tokens", "shared"}) + if embedding_path is not None: + backbone_path = embedding_path.rpartition(".")[0] + backbone = get_attr(root_model, backbone_path) if backbone_path else root_model + if hasattr(backbone, "layers"): + layers_path = f"{backbone_path}.layers" if backbone_path else "layers" + wrapper = _root_component_model_wrapper( + root_model, + backbone_path, + layers_path, + source_paths, + embedding_path, + ) + return wrapper, "" + + layers_path = _path_with_leaf(source_paths, {"layers"}) + if layers_path is not None and (component_role is not None or not slice_path): + backbone_path = layers_path.rpartition(".")[0] + return _root_component_model_wrapper( + root_model, + backbone_path, + layers_path, + source_paths, + ), "" + + if component_role in {"decoder", "embedding"}: + component_model = get_attr(root_model, slice_path) if slice_path else root_model + config = getattr(component_model, "config", root_model.config) + return _GenericComponentWrapper(component_model, config), (f"{slice_path}." if slice_path else "") + + if slice_path: + quant_model = get_attr(root_model, slice_path) + return ModelWrapper.from_model(quant_model), f"{slice_path}." + + return ModelWrapper.from_model(root_model), "" + + +def _validate_component_source_paths( + root_model: torch.nn.Module, + source_paths: list[str], +) -> None: + missing = [path for path in source_paths if get_attr(root_model, path) is None] + if missing: + raise ValueError( + "Component source path(s) do not exist in the loaded Hugging Face model: " + f"{missing}. The paths must match runtime model.named_modules() names." + ) + + +def get_qkv_quantization_groups( + wrapper: ModelWrapper, + module_names: set[str] | None = None, + name_prefix: str = "", +) -> list[tuple[str, ...]]: """Get attention input projection groups that must share quantization settings. Names are resolved from ``wrapper.model.named_modules()`` to stay correct for any layer @@ -92,10 +330,11 @@ def get_qkv_quantization_groups(wrapper: ModelWrapper, module_names: set[str] | qkv_groups = [] for layer_wrapper in wrapper.get_layer_wrappers(): attn_inputs, _ = layer_wrapper.get_attention_inputs() + attn_input_names = (module_to_name.get(id(module)) for module in attn_inputs) group = tuple( - name - for name in (module_to_name.get(id(module)) for module in attn_inputs) - if name is not None and (module_names is None or name in module_names) + root_name + for root_name in (_root_module_name(name, name_prefix) for name in attn_input_names) + if root_name and (module_names is None or root_name in module_names) ) if len(group) > 1: qkv_groups.append(group) @@ -125,6 +364,8 @@ def normalize_qkv_quant_config( wrapper: ModelWrapper, qcfg: OliveHfQuantizationConfig, locked_modules: set[str] | None = None, + module_names: set[str] | None = None, + name_prefix: str = "", ) -> OliveHfQuantizationConfig: """Promote split QKV projection overrides to one shared quantization config. @@ -138,7 +379,7 @@ def normalize_qkv_quant_config( locked members of one group disagree, the group is left untouched. """ locked_modules = locked_modules or set() - for group in get_qkv_quantization_groups(wrapper): + for group in get_qkv_quantization_groups(wrapper, module_names=module_names, name_prefix=name_prefix): group_qargs = {name: qcfg.get_qlinear_init_args(name) for name in group} if len({tuple(qargs.items()) for qargs in group_qargs.values()}) == 1: continue @@ -207,29 +448,97 @@ def prepare_model( if existing_qcfg.get("quant_method", None) != OliveHfQuantizationMethod.OLIVE: raise ValueError("Model has an existing quantization configuration that is not compatible with this pass.") - wrapper = ModelWrapper.from_model(load_hf_base_model(model)) + component_source_paths = _get_component_source_paths(model) + component_attributes = model.model_attributes or {} + component_name = component_attributes.get("component_name") + component_role = component_attributes.get("component_role") + if component_name and component_name != "model" and not component_source_paths: + raise ValueError( + f"Component {component_name!r} has no runtime source paths; refusing to apply " + "a PyTorch quantization pass to the whole model." + ) + root_model = load_hf_base_model(model) + _validate_component_source_paths(root_model, component_source_paths) + wrapper, name_prefix = _component_model_wrapper( + root_model, + component_source_paths, + component_role, + ) + wrapper.olive_root_model = root_model + wrapper.olive_component_path = name_prefix.rstrip(".") or None + wrapper.olive_component_role = component_role wrapper.model.eval() excluded_attn_inputs = _collect_excluded_attn_inputs(wrapper) if exclude_attn_inputs else set() - fresh_qcfg = normalize_qkv_quant_config(wrapper, get_quant_config(model, config)) + selected_module_names = {_root_module_name(name, name_prefix) for name, _ in wrapper.model.named_modules()} + fresh_qcfg = normalize_qkv_quant_config( + wrapper, + get_quant_config(model, config), + module_names=selected_module_names, + name_prefix=name_prefix, + ) - originally_tied_embeddings = wrapper.config.tie_word_embeddings + originally_tied_embeddings = getattr(wrapper.config, "tie_word_embeddings", False) if fresh_qcfg.lm_head or fresh_qcfg.embeds: wrapper.maybe_untie_word_embeddings() - lm_head_name = wrapper.get_lm_head()[1] - embeds_name = wrapper.get_embeds()[1][0] + declared_head_name = _path_with_leaf( + component_source_paths, + {"lm_head", "proj_out", "output_projection", "codec_head", "output"}, + ) + try: + lm_head_name = _root_module_name(wrapper.get_lm_head()[1], name_prefix) + except AttributeError: + lm_head_name = declared_head_name + if fresh_qcfg.lm_head and lm_head_name is None: + raise + declared_embeds_name = _path_with_leaf( + component_source_paths, + {"embed_tokens", "shared", "tok_embeddings", "text_embedding", "codec_embedding"}, + ) + component_embedding_names = [ + name + for name, module in root_model.named_modules() + if isinstance(module, torch.nn.Embedding) and _is_in_component(name, component_source_paths) + ] + try: + embeds_name = _root_module_name(wrapper.get_embeds()[1][0], name_prefix) + except AttributeError: + embeds_name = declared_embeds_name or next( + ( + name + for name in component_embedding_names + if name.rsplit(".", 1)[-1] + in {"embed_tokens", "shared", "tok_embeddings", "text_embedding", "codec_embedding"} + ), + component_embedding_names[0] if len(component_embedding_names) == 1 else None, + ) + if fresh_qcfg.embeds and not component_embedding_names: + raise ValueError("The selected component has no torch.nn.Embedding modules to quantize.") from None def should_quantize(module: torch.nn.Module, name: str) -> bool: + root_name = _root_module_name(name, name_prefix) if module in excluded_attn_inputs: return False + # When the slice spans more than the component (multi-path components slice to a + # common ancestor), restrict quantization to modules inside the declared sub-trees. + if not _is_in_component(root_name, component_source_paths): + return False if isinstance(module, torch.nn.Linear): - return name != lm_head_name or fresh_qcfg.lm_head + return lm_head_name is None or root_name != lm_head_name or fresh_qcfg.lm_head if fresh_qcfg.embeds and isinstance(module, torch.nn.Embedding): - return name == embeds_name + if component_source_paths or isinstance(wrapper, _GenericComponentWrapper): + return root_name in component_embedding_names + return root_name == embeds_name return False + fresh_names = { + _root_module_name(name, name_prefix) + for name, module in wrapper.model.named_modules() + if should_quantize(module, name) + } + # Pre-existing quantized weights are immutable. If we're merging with an existing # checkpoint, build the final qcfg first (merge fresh into existing, then renormalize # QKV with already-quantized modules locked) so that the quant_info we attach below @@ -241,9 +550,10 @@ def should_quantize(module: torch.nn.Module, name: str) -> bool: if existing_qcfg: on_disk_overrides = set((existing_qcfg.get("overrides") or {}).keys()) already_quantized = { - name for name, module in wrapper.model.named_modules() if isinstance(module, (QuantLinear, QuantEmbedding)) + _root_module_name(name, name_prefix) + for name, module in wrapper.model.named_modules() + if isinstance(module, (QuantLinear, QuantEmbedding)) } - fresh_names = {name for name, module in wrapper.model.named_modules() if should_quantize(module, name)} merged = existing_qcfg merged["overrides"] = existing_qcfg.get("overrides") or {} for name in fresh_names: @@ -254,17 +564,36 @@ def should_quantize(module: torch.nn.Module, name: str) -> bool: merged["lm_head"] |= fresh_qcfg.lm_head merged["embeds"] |= fresh_qcfg.embeds qcfg = OliveHfQuantizationConfig(**merged) - qcfg = normalize_qkv_quant_config(wrapper, qcfg, locked_modules=already_quantized) + qcfg = normalize_qkv_quant_config( + wrapper, + qcfg, + locked_modules=already_quantized, + module_names=selected_module_names, + name_prefix=name_prefix, + ) else: qcfg = fresh_qcfg + existing_modules_to_not_convert = { + excluded + for excluded in qcfg.modules_to_not_convert or [] + if not any(excluded in fresh_name for fresh_name in fresh_names) + } + unquantized_modules = { + name + for name, module in root_model.named_modules() + if isinstance(module, (torch.nn.Linear, torch.nn.Embedding)) and name not in fresh_names + } + qcfg.modules_to_not_convert = sorted(existing_modules_to_not_convert | unquantized_modules) or None + new_qargs: dict[str, dict[str, int | bool]] = {} def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module: # TODO(jambayk): validate that the module and config are compatible - qargs = qcfg.get_qlinear_init_args(name) + root_name = _root_module_name(name, name_prefix) + qargs = qcfg.get_qlinear_init_args(root_name) module.quant_info = QuantInfo(quantizer=WeightQuantizer(**qargs)) - new_qargs[name] = qargs + new_qargs[root_name] = qargs return module replace_matching_submodules(wrapper.model, should_quantize, add_quant_info, description="Preparing model") @@ -280,6 +609,8 @@ def add_quant_info(module: torch.nn.Module, name: str) -> torch.nn.Module: word_embeddings_eligible_for_tieing = ( originally_tied_embeddings + and not isinstance(wrapper, _GenericComponentWrapper) + and lm_head_name is not None and embeds_name in new_qargs and lm_head_name in new_qargs and new_qargs[embeds_name] == new_qargs[lm_head_name] @@ -459,6 +790,13 @@ def run_layerwise_quantization( Device string used for calibration. """ + component_role = wrapper.olive_component_role + if component_role not in {None, "decoder"} or isinstance(wrapper, _GenericComponentWrapper): + raise ValueError( + "Layerwise calibration requires a decoder component with identifiable transformer layers. " + "Use RTN or KQuant for generic encoder/vision/embedding components." + ) + from tqdm.auto import tqdm if device is None: @@ -567,22 +905,29 @@ def quantize_and_pack(module: torch.nn.Module, _: str) -> QuantLinear | QuantEmb zero_points=module.quant_info.zero_points, ).to("cpu") # move the original module to CPU - replace_matching_submodules( + root_model = wrapper.olive_root_model if wrapper.olive_root_model is not None else wrapper.model + packed_model = replace_matching_submodules( wrapper.model, should_quantize, quantize_and_pack, description="Quantizing and packing linear layers", ) + if packed_model is not wrapper.model: + component_path = wrapper.olive_component_path + if component_path: + set_attr(root_model, component_path, packed_model) + else: + root_model = packed_model if retie_word_embeddings: - tie_quant_word_embeddings(wrapper.model) + tie_quant_word_embeddings(packed_model) quant_config.tie_word_embeddings = True - wrapper.model.quantization_method = quant_config.quant_method - wrapper.model.config.quantization_config = quant_config + root_model.quantization_method = quant_config.quant_method + root_model.config.quantization_config = quant_config # save the quantized model - wrapper.model.save_pretrained(output_model_path) + root_model.save_pretrained(output_model_path) model.save_metadata(output_model_path) return inherit_hf_from_hf(model, output_model_path, adapter_path=model.adapter_path) diff --git a/olive/systems/utils/misc.py b/olive/systems/utils/misc.py index f5ae6c3ead..a4d54c64c3 100644 --- a/olive/systems/utils/misc.py +++ b/olive/systems/utils/misc.py @@ -28,7 +28,12 @@ def create_new_environ( python_environment_path: Optional[Union[Path, str]] = None, ): """Create a copy of the current environment with the given environment variables and paths prepended.""" + from olive.cache import get_cache_dir_from_env + environ = deepcopy(os.environ) + cache_dir = get_cache_dir_from_env() + if cache_dir: + environ["OLIVE_CACHE_DIR"] = cache_dir if environment_variables: environ.update(environment_variables) if prepend_to_path: diff --git a/olive/workflows/run/builds.py b/olive/workflows/run/builds.py new file mode 100644 index 0000000000..5ef2d9467d --- /dev/null +++ b/olive/workflows/run/builds.py @@ -0,0 +1,145 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import re +from collections import OrderedDict +from copy import deepcopy +from pathlib import Path +from typing import Annotated, Optional, Union + +from pydantic import ConfigDict, Field, StringConstraints + +from olive.common.config_utils import ConfigBase, load_config_file +from olive.common.constants import DEFAULT_WORKFLOW_ID +from olive.evaluator.olive_evaluator import OliveEvaluatorConfig +from olive.model import ModelConfig +from olive.search.search_strategy import SearchStrategyConfig +from olive.systems.system_config import SystemConfig +from olive.workflows.run.config import RunConfig + +BUILD_DEFAULT_KEY = "_default" +BUILD_NAME_PATTERN = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*") +NonEmptyString = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] + + +class BuildConfigPartial(ConfigBase): + """Partial build configuration used by the ``_default`` entry.""" + + model_config = ConfigDict(extra="forbid") + + components: Optional[list[NonEmptyString]] = Field(None, min_length=1) + pipeline: Optional[list[NonEmptyString]] = Field(None, min_length=1) + output_dir: Optional[NonEmptyString] = None + host: Optional[Union[SystemConfig, str]] = None + target: Optional[Union[SystemConfig, str]] = None + evaluator: Optional[Union[OliveEvaluatorConfig, str]] = None + search_strategy: Optional[Union[SearchStrategyConfig, bool]] = None + + +class BuildConfig(BuildConfigPartial): + """A build that expands into one ordinary Olive run configuration.""" + + pipeline: list[NonEmptyString] = Field(..., min_length=1) + output_dir: NonEmptyString = Field(...) + + +def parse_run_config( + run_config: Union[str, Path, dict], +) -> Union[RunConfig, OrderedDict[str, RunConfig]]: + """Parse one ordinary run config or expand and prevalidate every configured build.""" + raw_run_config = deepcopy(run_config) if isinstance(run_config, dict) else load_config_file(run_config) + if not isinstance(raw_run_config, dict): + raise TypeError("Olive run configuration must be a dictionary.") + if "builds" not in raw_run_config: + return RunConfig.model_validate(raw_run_config) + + parsed_builds = OrderedDict() + for build_name, build_config in expand_builds(raw_run_config).items(): + try: + parsed_builds[build_name] = RunConfig.model_validate(deepcopy(build_config)) + except (TypeError, ValueError) as exc: + raise ValueError(f"Invalid build {build_name!r}: {exc}") from exc + return parsed_builds + + +def expand_builds(run_config: dict) -> OrderedDict[str, dict]: + """Expand ``builds`` into independent, ordinary Olive run configurations.""" + if not isinstance(run_config, dict): + raise TypeError("Multi-build configuration must be a dictionary.") + + source_config = deepcopy(run_config) + if "builds" not in source_config: + return OrderedDict() + raw_builds = source_config.pop("builds") + if not isinstance(raw_builds, dict): + raise ValueError("`builds` must be a dictionary keyed by build name.") + + builds = _parse_builds(raw_builds) + passes = source_config.get("passes") or {} + workflow_id = source_config.get("workflow_id", DEFAULT_WORKFLOW_ID) + expanded = OrderedDict() + + for build_name, build in builds.items(): + missing_passes = [pass_name for pass_name in build.pipeline if pass_name not in passes] + if missing_passes: + raise ValueError( + f"Build {build_name!r} references unknown pass(es) {missing_passes}. Known passes: {sorted(passes)}." + ) + + child_config = deepcopy(source_config) + child_config["workflow_id"] = f"{workflow_id}_{build_name}" + child_config["passes"] = OrderedDict((pass_name, deepcopy(passes[pass_name])) for pass_name in build.pipeline) + _set_engine_value(child_config, "output_dir", build.output_dir) + + for field_name in ("host", "target", "evaluator", "search_strategy"): + value = getattr(build, field_name) + if value is not None: + _set_engine_value(child_config, field_name, value) + + if build.components: + input_model = child_config.get("input_model") + if input_model is None: + raise ValueError(f"Build {build_name!r} selects components but no input_model is configured.") + child_config["input_model"] = ( + ModelConfig.model_validate(deepcopy(input_model)).select_components(build.components).model_dump() + ) + + expanded[build_name] = child_config + + return expanded + + +def _parse_builds(raw_builds: dict) -> OrderedDict[str, BuildConfig]: + default_raw = raw_builds.get(BUILD_DEFAULT_KEY, {}) + if not isinstance(default_raw, dict): + raise ValueError("`builds._default` must be a dictionary.") + default_config = BuildConfigPartial.model_validate(default_raw).model_dump(exclude_none=True) + builds = OrderedDict() + + for build_name, raw_build in raw_builds.items(): + if build_name == BUILD_DEFAULT_KEY: + continue + if not isinstance(build_name, str) or not BUILD_NAME_PATTERN.fullmatch(build_name): + raise ValueError(f"Invalid build name {build_name!r}. Use letters, numbers, dots, underscores, or hyphens.") + if not isinstance(raw_build, dict): + raise ValueError(f"Build {build_name!r} must be a dictionary.") + builds[build_name] = BuildConfig.model_validate({**default_config, **raw_build}) + + if not builds: + raise ValueError("`builds` must contain at least one named build in addition to `_default`.") + return builds + + +def _set_engine_value(run_config: dict, field_name: str, value) -> None: + run_config.pop(field_name, None) + engine_config = run_config.get("engine") or {} + if hasattr(engine_config, "model_dump"): + engine_config = engine_config.model_dump() + elif not isinstance(engine_config, dict): + raise ValueError("`engine` must be a dictionary.") + else: + engine_config = deepcopy(engine_config) + + engine_config[field_name] = value.model_dump() if hasattr(value, "model_dump") else deepcopy(value) + run_config["engine"] = engine_config diff --git a/olive/workflows/run/run.py b/olive/workflows/run/run.py index 89100e1c1c..173df5a6d1 100644 --- a/olive/workflows/run/run.py +++ b/olive/workflows/run/run.py @@ -3,16 +3,20 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import logging +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor, as_completed from copy import deepcopy from pathlib import Path from typing import TYPE_CHECKING, Optional, Union +from olive.cache import CacheConfig, isolated_cache_env from olive.common.utils import set_tempdir from olive.hardware.constants import ExecutionProvider from olive.logging import set_default_logger_severity, set_ort_logger_severity, set_verbosity_info from olive.package_config import OlivePackageConfig from olive.systems.accelerator_creator import create_accelerator from olive.systems.common import SystemType +from olive.workflows.run.builds import parse_run_config from olive.workflows.run.config import RunConfig if TYPE_CHECKING: @@ -116,15 +120,9 @@ def run_engine(package_config: OlivePackageConfig, run_config: RunConfig): # check if target is not used used_passes_configs = get_used_passes_configs(run_config) - target_not_used = ( - # no evaluator given (also implies no search) - engine.evaluator_config is None - # no pass specific evaluator - # no pass needs to run on target - and all( - pass_config.evaluator is None and not get_run_on_target(package_config, pass_config) - for pass_config in used_passes_configs - ) + target_not_used = engine.evaluator_config is None and all( + pass_config.evaluator is None and not get_run_on_target(package_config, pass_config) + for pass_config in used_passes_configs ) is_ep_required = is_execution_provider_required(run_config, package_config) @@ -160,24 +158,188 @@ def run( package_config = OlivePackageConfig.get_default_config_path() package_config = OlivePackageConfig.parse_file_or_obj(package_config) - run_config: RunConfig = RunConfig.parse_file_or_obj(run_config) + parsed_config = parse_run_config(run_config) + if isinstance(parsed_config, dict): + if list_required_packages: + _list_required_packages(package_config, parsed_config.values()) + return None + return _run_builds_in_parallel(package_config, parsed_config) if list_required_packages: - # set the log level to INFO for packages - set_verbosity_info() - required_packages = get_required_packages(package_config, run_config) - generate_files_from_packages(required_packages, "olive_requirements.txt") + _list_required_packages(package_config, [parsed_config]) return None + return _run_single(package_config, parsed_config) + +def _run_builds_in_parallel(package_config: OlivePackageConfig, build_configs: dict[str, RunConfig]) -> OrderedDict: + _validate_parallel_write_dirs(build_configs) + docker_systems, docker_cleanup_groups = _prepare_shared_docker_systems(build_configs) + results = {} + errors = {} + try: + with ThreadPoolExecutor(max_workers=len(build_configs), thread_name_prefix="olive-build") as executor: + future_to_name = { + executor.submit( + _run_named_build, + deepcopy(package_config), + build_name, + build_config, + docker_systems.get(build_name), + ): build_name + for build_name, build_config in build_configs.items() + } + for future in as_completed(future_to_name): + build_name = future_to_name[future] + try: + results[build_name] = future.result() + except Exception as exc: # pylint: disable=broad-exception-caught + errors.setdefault(build_name, []).append(exc) + finally: + for build_names, exc in _remove_shared_docker_images(docker_cleanup_groups): + build_name = build_names[0] + cleanup_error = RuntimeError(f"Failed to clean shared Docker image for builds {list(build_names)}: {exc}") + errors.setdefault(build_name, []).append(cleanup_error) + + if errors: + failed_names = [build_name for build_name in build_configs if build_name in errors] + first_failed = failed_names[0] + details = "; ".join( + f"{build_name}: {type(error).__name__}: {error}" + for build_name in failed_names + for error in errors[build_name] + ) + raise RuntimeError(f"Build(s) {failed_names} failed: {details}") from errors[first_failed][0] + + return OrderedDict((build_name, results[build_name]) for build_name in build_configs) + + +def _validate_parallel_write_dirs(build_configs: dict[str, RunConfig]) -> None: + write_dirs = {} + for build_name, build_config in build_configs.items(): + output_dir = Path(build_config.engine.output_dir) + artifact_dir = output_dir.parent if output_dir.suffix and not output_dir.is_dir() else output_dir + write_dirs[build_name] = { + "artifact": artifact_dir.resolve(), + "cache": _get_build_cache_dir(build_config), + } + + checked_dirs = {} + for build_name, build_dirs in write_dirs.items(): + for other_name, other_dirs in checked_dirs.items(): + for build_dir_type, build_dir in build_dirs.items(): + for other_dir_type, other_dir in other_dirs.items(): + if _paths_overlap(build_dir, other_dir): + raise ValueError( + f"Parallel builds {other_name!r} and {build_name!r} have overlapping writable " + f"directories: {other_dir_type} directory {other_dir} and " + f"{build_dir_type} directory {build_dir}." + ) + checked_dirs[build_name] = build_dirs + + +def _get_build_cache_dir(run_config: RunConfig) -> Path: + cache_config = run_config.engine.cache_config + if cache_config is None: + cache_config = CacheConfig(cache_dir=run_config.engine.cache_dir) + elif isinstance(cache_config, dict): + cache_config = CacheConfig.model_validate(cache_config) + return (Path(cache_config.get_local_cache_dir()) / run_config.workflow_id).resolve() + + +def _paths_overlap(first: Path, second: Path) -> bool: + return first == second or first in second.parents or second in first.parents + + +def _prepare_shared_docker_systems(build_configs: dict[str, RunConfig]): + builds_by_image = {} + image_build_configs = {} + for build_name, build_config in build_configs.items(): + host = build_config.engine.host + if host and host.type == SystemType.Docker: + image_name = _normalize_docker_image_name(host.config.image_name) + image_build_config = host.config.model_dump(include={"build_context_path", "dockerfile", "build_args"}) + if image_name in image_build_configs and image_build_configs[image_name] != image_build_config: + other_name = builds_by_image[image_name][0] + raise ValueError( + f"Parallel Docker builds {other_name!r} and {build_name!r} use image {image_name!r} with " + "conflicting image build settings." + ) + image_build_configs[image_name] = image_build_config + builds_by_image.setdefault(image_name, []).append(build_name) + + docker_systems = {} + cleanup_groups = [] + for image_name, build_names in builds_by_image.items(): + if len(build_names) < 2: + continue + + group_systems = [] + try: + for build_name in build_names: + docker_system = build_configs[build_name].engine.host.create_system() + docker_system.clean_image = False + docker_systems[build_name] = docker_system + group_systems.append(docker_system) + except Exception as exc: + if group_systems and all(build_configs[name].engine.host.config.clean_image for name in build_names): + cleanup_groups.append((tuple(build_names), group_systems[0])) + cleanup_errors = _remove_shared_docker_images(cleanup_groups) + cleanup_details = "".join( + f" Cleanup for builds {list(names)} also failed: {cleanup_error}." + for names, cleanup_error in cleanup_errors + ) + raise RuntimeError( + f"Failed to prepare shared Docker image {image_name!r} for builds {build_names}: {exc}." + f"{cleanup_details}" + ) from exc + + if all(build_configs[name].engine.host.config.clean_image for name in build_names): + cleanup_groups.append((tuple(build_names), group_systems[0])) + + return docker_systems, cleanup_groups + + +def _normalize_docker_image_name(image_name: str) -> str: + if "@" in image_name or ":" in image_name.rsplit("/", 1)[-1]: + return image_name + return f"{image_name}:latest" + + +def _remove_shared_docker_images(cleanup_groups): + errors = [] + for build_names, docker_system in cleanup_groups: + try: + docker_system.remove() + except Exception as exc: # pylint: disable=broad-exception-caught + errors.append((build_names, exc)) + return errors + + +def _run_named_build(package_config: OlivePackageConfig, build_name: str, run_config: RunConfig, docker_system=None): + logger.info("Running build %s", build_name) + with isolated_cache_env(_get_build_cache_dir(run_config)): + if docker_system is None: + return _run_single(package_config, run_config) + return _run_single(package_config, run_config, docker_system) + + +def _run_single(package_config: OlivePackageConfig, run_config: RunConfig, docker_system=None): if run_config.engine.host and run_config.engine.host.type == SystemType.Docker: - docker_system = run_config.engine.host.create_system() + docker_system = docker_system or run_config.engine.host.create_system() return docker_system.run_workflow(run_config) - # set log level for olive set_default_logger_severity(run_config.engine.log_severity_level) return run_engine(package_config, run_config) +def _list_required_packages(package_config: OlivePackageConfig, run_configs) -> None: + set_verbosity_info() + required_packages = set() + for run_config in run_configs: + required_packages.update(get_required_packages(package_config, run_config)) + generate_files_from_packages(required_packages, "olive_requirements.txt") + + def generate_files_from_packages(packages, file_name): file_path = Path(file_name) if file_path.exists(): diff --git a/pyproject.toml b/pyproject.toml index 52a945e7d6..00fcd8b94c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -191,3 +191,4 @@ classmethod-decorators = ["classmethod", "pydantic.field_validator", "pydantic.m "test/**" = ["INP001"] "scripts/**" = ["INP001"] "olive/cli/**" = ["T201"] +"multi_comp_recipe/**" = ["T201"] diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index 999e189132..69761c8044 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -6,7 +6,7 @@ import subprocess import sys from pathlib import Path -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -112,6 +112,67 @@ def test_workflow_run_command(mock_run, tempdir, list_required_packages, tmp_pat ) +@patch("olive.workflows.run") +def test_workflow_run_command_prints_build_outputs(mock_run, tmp_path, capsys): + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "builds": { + "_default": {"output_dir": "out/default"}, + "first": {"pipeline": ["convert"]}, + "second": {"pipeline": ["convert"], "output_dir": "out/second"}, + "missing": {"pipeline": ["convert"], "output_dir": "out/missing"}, + } + } + ) + ) + output = MagicMock() + output.has_output_model.return_value = True + missing_output = MagicMock() + missing_output.has_output_model.return_value = False + mock_run.return_value = {"first": output, "second": output, "missing": missing_output} + + cli_main(["run", "--run-config", str(config_path)]) + + stdout = capsys.readouterr().out + assert "Build 'first': model is saved under out/default" in stdout + assert "Build 'second': model is saved under out/second" in stdout + assert "Build 'missing': no output model produced" in stdout + + +def test_workflow_run_command_rejects_test_with_builds(tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "builds": { + "only": {"pipeline": ["convert"], "output_dir": "out/only"}, + } + } + ) + ) + + with pytest.raises(ValueError, match="not supported with multi-build"): + cli_main(["run", "--run-config", str(config_path), "--test"]) + + +def test_workflow_run_command_rejects_output_path_with_builds(tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps( + { + "builds": { + "only": {"pipeline": ["convert"], "output_dir": "out/only"}, + } + } + ) + ) + + with pytest.raises(ValueError, match="Set output_dir on each build"): + cli_main(["run", "--run-config", str(config_path), "--output_path", str(tmp_path / "output")]) + + @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) def test_workflow_run_command_with_overrides(mock_repo_exists, mock_run, tmp_path): diff --git a/test/common/test_hf.py b/test/common/test_hf.py index 4204606e0b..acc08c089c 100644 --- a/test/common/test_hf.py +++ b/test/common/test_hf.py @@ -7,7 +7,7 @@ import pytest import torch -from transformers import BertConfig, GPT2Config, Qwen3Config +from transformers import AutoModelForSeq2SeqLM, BertConfig, GPT2Config, Qwen3Config from olive.common.hf.model_io import get_model_dummy_input, get_model_io_config from olive.common.hf.utils import ( @@ -28,6 +28,20 @@ def test_load_model_from_task(): assert isinstance(model, torch.nn.Module) +@pytest.mark.parametrize("task", ["text2text-generation", "text2text-generation-with-past"]) +def test_load_model_from_task_uses_seq2seq_model_class(task): + loaded_model = MagicMock(spec=torch.nn.Module) + + with ( + patch("olive.common.hf.utils.get_model_config", return_value=MagicMock(quantization_config=None)), + patch("olive.common.hf.utils.from_pretrained", return_value=loaded_model) as mock_from_pretrained, + ): + model = load_model_from_task(task, "dummy-model") + + assert model is loaded_model + mock_from_pretrained.assert_called_once_with(AutoModelForSeq2SeqLM, "dummy-model", "model") + + @pytest.mark.parametrize( ("model_config", "hidden_layers_attr"), [ diff --git a/test/common/test_hf_wrapper.py b/test/common/test_hf_wrapper.py index 667d645ce0..cd5ebb6c68 100644 --- a/test/common/test_hf_wrapper.py +++ b/test/common/test_hf_wrapper.py @@ -130,10 +130,38 @@ def test_hf_wrapper_lfm2(): assert attn_modules == [] assert attn_names == [] - attn_out_modules, attn_out_names = layer_wrapper.get_attention_outputs() - assert attn_out_modules == [] - assert attn_out_names == [] - # LFM2 must have both layer types assert has_attn_layer, "Expected at least one attention layer" assert has_conv_layer, "Expected at least one conv layer" + + +def test_hf_wrapper_qwen3_vl_text(): + pytest.importorskip("transformers.models.qwen3_vl") + from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLTextConfig + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextModel + + config = Qwen3VLTextConfig( # pylint: disable=unexpected-keyword-arg + vocab_size=128, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=4, + head_dim=4, + ) + model = Qwen3VLTextModel(config) + model_wrapper = ModelWrapper.from_model(model) + + assert model_wrapper.model_type == "qwen3_vl_text" + assert isinstance(model_wrapper.get_embeds(False)[0], nn.Embedding) + assert len(model_wrapper.get_layers(False)) == 1 + assert model_wrapper.get_pre_head_layernorm(False).__class__.__name__.endswith("RMSNorm") + assert model_wrapper.get_rotary_embed(False).__class__.__name__.endswith("RotaryEmbedding") + + layer_wrapper = model_wrapper.get_layer_wrappers()[0] + for key in ["get_attention_inputs", "get_attention_outputs", "get_mlp_inputs", "get_mlp_outputs"]: + modules, names = getattr(layer_wrapper, key)() + assert modules + assert names + for module in modules: + assert isinstance(module, nn.Linear) diff --git a/test/mcp/test_worker.py b/test/mcp/test_worker.py new file mode 100644 index 0000000000..7a588adfa4 --- /dev/null +++ b/test/mcp/test_worker.py @@ -0,0 +1,50 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +import importlib.util +from pathlib import Path + +# pylint: disable=protected-access + + +def _load_worker_module(): + worker_path = Path(__file__).parents[2] / "mcp" / "src" / "olive_mcp" / "worker.py" + spec = importlib.util.spec_from_file_location("olive_mcp_worker_test", worker_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load MCP worker from {worker_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +worker = _load_worker_module() + + +def test_serialize_workflow_output_handles_build_mapping(): + result = worker.serialize_workflow_output({"first": None, "second": None}) + + assert result == { + "status": "success", + "builds": { + "first": {"status": "success", "output_models": []}, + "second": {"status": "success", "output_models": []}, + }, + } + + +def test_validate_config_accepts_multi_build_config(): + config = { + "input_model": {"type": "ONNXModel", "model_path": "model.onnx"}, + "passes": { + "convert": {"type": "OnnxConversion"}, + "tune": {"type": "OrtSessionParamsTuning"}, + }, + "evaluate_input_model": False, + "builds": { + "convert": {"pipeline": ["convert"], "output_dir": "out/convert"}, + "tune": {"pipeline": ["tune"], "output_dir": "out/tune"}, + }, + } + + assert worker._handle_validate_config({"config": config}) == {"valid": True, "message": "Config is valid."} diff --git a/test/model/test_composite_model.py b/test/model/test_composite_model.py index ca3fcd1f40..d38c60cf93 100644 --- a/test/model/test_composite_model.py +++ b/test/model/test_composite_model.py @@ -2,8 +2,11 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- +from unittest.mock import Mock + import pytest +from olive.common import mobius_utils from olive.model.config.model_config import ModelConfig from olive.model.handler.composite import CompositeModelHandler from olive.model.handler.onnx import ONNXModelHandler @@ -42,3 +45,277 @@ def test_composite_model(as_handler): assert composite_json["config"]["model_components"][0]["config"]["model_attributes"] == {"attr0": "value0"} model_config = ModelConfig.from_json(composite_json) assert model_config.type == CompositeModelHandler.model_type + + +def test_composite_model_missing_names_raises_value_error(): + # model_components provided but model_component_names omitted: must raise a clear ValueError + # (not a TypeError from len(None), which would also be silenced under `python -O` if asserted). + with pytest.raises(ValueError, match="requires model_component_names"): + CompositeModelHandler([get_onnx_model(), get_onnx_model()]) + + +def test_composite_model_component_name_count_mismatch_raises_value_error(): + with pytest.raises(ValueError, match="must match"): + CompositeModelHandler([get_onnx_model(), get_onnx_model()], ["only_one_name"]) + + +def _build_composite_handler(): + return CompositeModelHandler( + [get_onnx_model(), get_onnx_model(), get_onnx_model()], + ["text_encoder", "unet", "vae_decoder"], + model_attributes={"shared": "value"}, + ) + + +def test_select_components_single_returns_unwrapped_child(): + composite = _build_composite_handler() + selected = composite.select_components(["unet"]) + assert isinstance(selected, ONNXModelHandler) + # parent attributes should be inherited by the unwrapped child + assert selected.model_attributes == {"shared": "value"} + + +def test_select_components_multiple_returns_subset_composite(): + composite = _build_composite_handler() + selected = composite.select_components(["vae_decoder", "text_encoder"]) + assert isinstance(selected, CompositeModelHandler) + # order from the call is preserved + assert list(selected.model_component_names) == ["vae_decoder", "text_encoder"] + + +def test_select_components_unknown_name_raises(): + composite = _build_composite_handler() + with pytest.raises(ValueError, match="Unknown component"): + composite.select_components(["no_such_component"]) + + +def test_select_components_empty_list_raises(): + composite = _build_composite_handler() + with pytest.raises(ValueError, match="non-empty"): + composite.select_components([]) + + +def test_model_config_select_components_multiple_returns_composite_config(): + composite_config = ModelConfig.model_validate( + { + "type": "CompositeModel", + "config": { + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": "a.onnx"}}, + {"type": "ONNXModel", "config": {"model_path": "b.onnx"}}, + {"type": "ONNXModel", "config": {"model_path": "c.onnx"}}, + ], + "model_component_names": ["text_encoder", "unet", "vae_decoder"], + }, + } + ) + selected = composite_config.select_components(["vae_decoder", "text_encoder"]) + assert isinstance(selected, ModelConfig) + assert selected.type == "compositemodel" + assert list(selected.config["model_component_names"]) == ["vae_decoder", "text_encoder"] + assert [c["config"]["model_path"] for c in selected.config["model_components"]] == ["c.onnx", "a.onnx"] + + +def test_model_config_select_components_on_non_composite_raises(): + onnx_config = ModelConfig.model_validate({"type": "ONNXModel", "config": {"model_path": "a.onnx"}}) + with pytest.raises(ValueError, match="only supported on CompositeModel"): + onnx_config.select_components(["any"]) + + +def test_model_config_select_components_single_inherits_parent_attributes(): + composite_config = ModelConfig.model_validate( + { + "type": "CompositeModel", + "config": { + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": "a.onnx", "model_attributes": {"child": "c"}}}, + {"type": "ONNXModel", "config": {"model_path": "b.onnx"}}, + ], + "model_component_names": ["text_encoder", "unet"], + "model_attributes": {"shared": "s", "child": "parent"}, + }, + } + ) + selected = composite_config.select_components(["text_encoder"]) + assert isinstance(selected, ModelConfig) + assert selected.type == "onnxmodel" + # parent-only keys are inherited; child keys win on conflict + assert selected.config["model_attributes"] == {"shared": "s", "child": "c"} + + +def test_model_config_get_components_returns_none_for_non_composite(): + onnx_config = ModelConfig.model_validate({"type": "ONNXModel", "config": {"model_path": "a.onnx"}}) + assert onnx_config.get_components() is None + + +def test_model_config_get_components_returns_names_for_composite(): + composite_config = ModelConfig.model_validate( + { + "type": "CompositeModel", + "config": { + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": "a.onnx"}}, + {"type": "ONNXModel", "config": {"model_path": "b.onnx"}}, + ], + "model_component_names": ["text_encoder", "unet"], + }, + } + ) + assert composite_config.get_components() == ["text_encoder", "unet"] + + +def _make_export_package(root): + """Create an export package with one model.onnx subfolder per component.""" + for name in ["decoder", "vision_encoder", "embedding"]: + comp_dir = root / name + comp_dir.mkdir(parents=True) + (comp_dir / "model.onnx").write_bytes(b"onnx") + return root + + +def test_discover_onnx_components_empty_for_flat_dir(tmp_path): + from olive.model.utils.onnx_utils import discover_onnx_components + + (tmp_path / "model.onnx").write_bytes(b"onnx") + assert not discover_onnx_components(str(tmp_path)) + + +def test_composite_handler_discovers_components_from_directory(tmp_path): + _make_export_package(tmp_path) + handler = CompositeModelHandler(model_path=str(tmp_path)) + assert list(handler.model_component_names) == ["decoder", "embedding", "vision_encoder"] + for _, component in handler.get_model_components(): + assert isinstance(component, ONNXModelHandler) + + +def test_model_config_get_components_discovers_directory_composite(tmp_path): + _make_export_package(tmp_path) + config = ModelConfig.model_validate({"type": "CompositeModel", "config": {"model_path": str(tmp_path)}}) + assert config.get_components() == ["decoder", "embedding", "vision_encoder"] + + +def test_model_config_select_components_discovers_directory_composite(tmp_path): + _make_export_package(tmp_path) + config = ModelConfig.model_validate({"type": "CompositeModel", "config": {"model_path": str(tmp_path)}}) + selected = config.select_components(["decoder"]) + assert isinstance(selected, ModelConfig) + assert selected.type == "onnxmodel" + assert selected.config["onnx_file_name"] == "decoder/model.onnx" + + +def test_model_config_get_components_hfmodel_uses_mobius(monkeypatch): + inspect_components = Mock( + return_value=[mobius_utils.ComponentInfo(name="decoder", role="decoder", source_paths=["model.language_model"])] + ) + monkeypatch.setattr(mobius_utils, "inspect_components", inspect_components) + config = ModelConfig.model_validate( + { + "type": "HfModel", + "config": { + "model_path": "some/vlm", + "load_kwargs": {"trust_remote_code": True}, + "model_attributes": {"mobius_task": "qwen-vl"}, + }, + } + ) + + assert config.get_components() == ["decoder"] + inspect_components.assert_called_once_with("some/vlm", task="qwen-vl", trust_remote_code=True) + + +def test_model_config_select_components_hfmodel_tags_component(monkeypatch): + monkeypatch.setattr( + mobius_utils, + "inspect_components", + lambda *args, **kwargs: [ + mobius_utils.ComponentInfo(name="decoder", role="decoder", source_paths=["model.language_model"]) + ], + ) + config = ModelConfig.model_validate({"type": "HfModel", "config": {"model_path": "some/vlm"}}) + + selected = config.select_components(["decoder"]) + + assert selected.type == "hfmodel" + assert selected.config["model_path"] == "some/vlm" + assert selected.config["model_attributes"] == { + "component_name": "decoder", + "component_role": "decoder", + "component_source_paths": ["model.language_model"], + } + + +def test_model_config_select_components_hfmodel_multiple_names_raises(): + config = ModelConfig.model_validate({"type": "HfModel", "config": {"model_path": "some/vlm"}}) + + with pytest.raises(ValueError, match="one at a time"): + config.select_components(["decoder", "vision_encoder"]) + + +def test_model_config_select_components_hfmodel_missing_paths_raises(monkeypatch): + monkeypatch.setattr( + mobius_utils, + "inspect_components", + lambda *args, **kwargs: [ + mobius_utils.ComponentInfo(name="decoder", role="decoder"), + mobius_utils.ComponentInfo(name="vision_encoder", role="encoder", source_paths=["model.visual"]), + ], + ) + config = ModelConfig.model_validate({"type": "HfModel", "config": {"model_path": "some/vlm"}}) + + with pytest.raises(ValueError, match="no runtime source paths"): + config.select_components(["decoder"]) + + +def test_model_config_select_components_hfmodel_whole_model_allows_empty_paths(monkeypatch): + monkeypatch.setattr( + mobius_utils, + "inspect_components", + lambda *args, **kwargs: [mobius_utils.ComponentInfo(name="model", role="decoder")], + ) + config = ModelConfig.model_validate({"type": "HfModel", "config": {"model_path": "some/llm"}}) + + selected = config.select_components(["model"]) + + assert selected.config["model_attributes"] == {"component_name": "model", "component_role": "decoder"} + + +def _make_diffusers_dir(tmp_path): + """Create a minimal local diffusers dir so is_valid_diffusers_model passes offline.""" + (tmp_path / "model_index.json").write_text("{}") + return tmp_path + + +def test_model_config_get_components_diffusersmodel(tmp_path): + model_dir = _make_diffusers_dir(tmp_path) + config = ModelConfig.model_validate( + {"type": "DiffusersModel", "config": {"model_path": str(model_dir), "model_variant": "sdxl"}} + ) + assert config.get_components() == [ + "text_encoder", + "text_encoder_2", + "unet", + "vae_encoder", + "vae_decoder", + ] + + +def test_model_config_select_components_diffusersmodel_scopes_subset(tmp_path): + model_dir = _make_diffusers_dir(tmp_path) + config = ModelConfig.model_validate( + {"type": "DiffusersModel", "config": {"model_path": str(model_dir), "model_variant": "sdxl"}} + ) + selected = config.select_components(["unet", "text_encoder"]) + assert selected.type == "diffusersmodel" + # preserved in the variant's canonical order, not the requested order + assert selected.config["components"] == ["text_encoder", "unet"] + # the scoped config now exposes only the selected components + assert selected.get_components() == ["text_encoder", "unet"] + + +def test_model_config_select_components_diffusersmodel_unknown_raises(tmp_path): + model_dir = _make_diffusers_dir(tmp_path) + config = ModelConfig.model_validate( + {"type": "DiffusersModel", "config": {"model_path": str(model_dir), "model_variant": "sd"}} + ) + with pytest.raises(ValueError, match="Unknown component name"): + config.select_components(["text_encoder_2"]) # SDXL-only; not in SD diff --git a/test/model/test_diffusers_model.py b/test/model/test_diffusers_model.py index acbe6113a4..44279f6d47 100644 --- a/test/model/test_diffusers_model.py +++ b/test/model/test_diffusers_model.py @@ -207,3 +207,22 @@ def test_adapter_path_property(self): def test_adapter_path_property_none(self): model = DiffusersModelHandler(model_path=self.model_path, model_variant=DiffusersModelVariant.SD) assert model.adapter_path is None + + @patch("olive.model.handler.diffusers.is_valid_diffusers_model", return_value=True) + def test_get_exportable_components_raises_for_unknown_component(self, mock_is_valid): + model = DiffusersModelHandler( + model_path=self.model_path, + model_variant=DiffusersModelVariant.SD, + components=["text_encoder_2"], # SDXL-only; not in SD + ) + with pytest.raises(ValueError, match="Unknown component"): + model.get_exportable_components() + + @patch("olive.model.handler.diffusers.is_valid_diffusers_model", return_value=True) + def test_to_json_round_trips_components(self, mock_is_valid): + model = DiffusersModelHandler( + model_path=self.model_path, + model_variant=DiffusersModelVariant.SDXL, + components=["text_encoder"], + ) + assert model.to_json()["config"]["components"] == ["text_encoder"] diff --git a/test/model/test_hf_model.py b/test/model/test_hf_model.py index 019ddefe2d..fb196b245d 100644 --- a/test/model/test_hf_model.py +++ b/test/model/test_hf_model.py @@ -66,11 +66,41 @@ def test_save_metadata(self, local, trust_remote_code, tokenizer_exists, tmp_pat if tokenizer_exists: olive_model.get_hf_tokenizer().save_pretrained(tmp_path) saved_filepaths = olive_model.save_metadata(tmp_path) - # transformers>=5.0.0 - assert len(saved_filepaths) == (4 if tokenizer_exists else 7) assert all(Path(fp).exists() for fp in saved_filepaths) assert isinstance(transformers.AutoConfig.from_pretrained(tmp_path), transformers.Phi3Config) assert isinstance(transformers.AutoTokenizer.from_pretrained(tmp_path), transformers.PreTrainedTokenizerBase) + assert (tmp_path / "generation_config.json").exists() + + def test_save_metadata_saves_processor_when_available(self, tmp_path): + olive_model = HfModelHandler( + model_path=self.local_path, task="image-text-to-text", load_kwargs={"revision": self.revision} + ) + + class MockProcessor: + @staticmethod + def save_pretrained(output_dir, **kwargs): + processor_path = Path(output_dir) / "preprocessor_config.json" + processor_path.write_text('{"processor": true}') + return (str(processor_path),) + + with patch.object(olive_model, "get_hf_processor", return_value=MockProcessor()): + saved_filepaths = olive_model.save_metadata(tmp_path) + + assert str(tmp_path / "preprocessor_config.json") in saved_filepaths + assert (tmp_path / "preprocessor_config.json").exists() + + def test_save_metadata_does_not_save_processor_when_task_is_none(self, tmp_path): + # task can be None when not provided in the config; save_metadata must not crash on + # `self.task.replace(...)` and must skip processor saving in that case. + olive_model = HfModelHandler(model_path=self.local_path, task=None, load_kwargs={"revision": self.revision}) + + with patch.object(olive_model, "get_hf_processor") as mock_get_processor: + saved_filepaths = olive_model.save_metadata(tmp_path) + + mock_get_processor.assert_not_called() + assert all(Path(fp).exists() for fp in saved_filepaths) + assert not (tmp_path / "preprocessor_config.json").exists() + assert not (tmp_path / "processor_config.json").exists() @pytest.mark.parametrize("local", [True, False]) def test_save_pretrained_metadata(self, local, tmp_path): diff --git a/test/passes/onnx/test_mnb_to_qdq.py b/test/passes/onnx/test_mnb_to_qdq.py index 391130e2af..ee35e5cc05 100644 --- a/test/passes/onnx/test_mnb_to_qdq.py +++ b/test/passes/onnx/test_mnb_to_qdq.py @@ -68,13 +68,16 @@ def __init__(self): def forward(self, x): return self.f3(self.f2(self.f1(x))) - model = TestModel() + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + model = TestModel() + example_input = torch.randn(1, 1, in_dim) # base model base_path = tmp_path / "base.onnx" torch.onnx.export( model, - torch.randn(1, 1, in_dim), + example_input, base_path, input_names=["input"], output_names=["output"], @@ -158,10 +161,10 @@ def test_mnb_to_qdq(create_mnb_model, nodes_to_exclude, add_zero_point, use_sign qdq_session = onnxruntime.InferenceSession(str(qdq_model.model_path), disabled_optimizers=disabled_optimizers) qdq_session.disable_fallback() - input_data = {"input": np.random.randn(1, 1, in_dim).astype(np.float32)} + input_data = {"input": np.random.default_rng(0).standard_normal((1, 1, in_dim)).astype(np.float32)} original_output = original_session.run(None, input_data)[0] qdq_output = qdq_session.run(None, input_data)[0] assert original_output.shape == qdq_output.shape assert original_output.dtype == qdq_output.dtype # acc level 4 is used for 8 bit, so the tolerance is higher - np.testing.assert_allclose(original_output, qdq_output, atol=1e-2 if bits == 8 else 1e-4) + np.testing.assert_allclose(original_output, qdq_output, atol=2e-2 if bits == 8 else 1e-4) diff --git a/test/passes/pytorch/test_quant_utils.py b/test/passes/pytorch/test_quant_utils.py index 6c3f8a79bf..ef9f7ff071 100644 --- a/test/passes/pytorch/test_quant_utils.py +++ b/test/passes/pytorch/test_quant_utils.py @@ -8,17 +8,29 @@ import pytest import torch -from transformers import LlamaConfig, LlamaForCausalLM +from transformers import ( + BartConfig, + BartForConditionalGeneration, + BertConfig, + BertForSequenceClassification, + LlamaConfig, + LlamaForCausalLM, + T5Config, + T5ForConditionalGeneration, +) from olive.common.hf.wrapper import ModelWrapper from olive.common.quant.hf_utils import OliveHfQuantizationConfig +from olive.common.quant.nn import QuantEmbedding, QuantLinear from olive.constants import PrecisionBits from olive.model import HfModelHandler from olive.passes.pytorch import quant_utils as quant_utils_module from olive.passes.pytorch.quant_utils import ( _quant_config_rank, + finalize, normalize_qkv_quant_config, prepare_model, + run_layerwise_quantization, ) from test.utils import get_tiny_phi3 @@ -44,12 +56,13 @@ def input_model_fixture(tmp_path_factory): return HfModelHandler(save_path) -def _baseline_pass_config(overrides=None): +def _baseline_pass_config(overrides=None, *, embeds=False): return SimpleNamespace( bits=PrecisionBits.BITS4, sym=False, group_size=16, lm_head=False, + embeds=embeds, overrides=overrides, ) @@ -66,6 +79,40 @@ def fake(self, exclude_load_keys=None): monkeypatch.setattr(HfModelHandler, "get_hf_model_config", fake) +class _NestedDecoderRoot(torch.nn.Module): + def __init__(self, decoder): + super().__init__() + self.decoder = decoder + self.vision = torch.nn.Linear(2, 2) + self.config = decoder.config + self.saved_state_keys = set() + + def save_pretrained(self, output_dir): + self.saved_state_keys = set(self.state_dict()) + self.config.save_pretrained(output_dir) + + +def _make_nested_decoder_root(input_model): + return _NestedDecoderRoot(LlamaForCausalLM.from_pretrained(input_model.model_path)) + + +class _NestedBackboneRoot(torch.nn.Module): + """VLM-like root where the text backbone and LM head are disjoint.""" + + def __init__(self, causal_lm): + super().__init__() + self.model = torch.nn.Module() + self.model.language_model = causal_lm.model + self.lm_head = causal_lm.lm_head + self.vision = torch.nn.Linear(2, 2) + self.config = causal_lm.config + self.saved_state_keys = set() + + def save_pretrained(self, output_dir): + self.saved_state_keys = set(self.state_dict()) + self.config.save_pretrained(output_dir) + + # --------------------------------------------------------------------------- # _quant_config_rank # --------------------------------------------------------------------------- @@ -204,6 +251,384 @@ def test_prepare_model_no_existing_quant_config_no_overrides_quantizes_all_linea assert eligible is False +def test_prepare_model_component_source_path_quantizes_only_selected_component(input_model, monkeypatch): + """A selected HfModel component should not attach quant_info outside that submodule.""" + root_model = _make_nested_decoder_root(input_model) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={"component_source_paths": ["decoder"]}, + ) + + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + assert wrapper.model is root_model.decoder + assert hasattr(root_model.decoder.model.layers[0].self_attn.q_proj, "quant_info") + assert not hasattr(root_model.vision, "quant_info") + assert "vision" in qcfg.modules_to_not_convert + assert "decoder.model.embed_tokens" in qcfg.modules_to_not_convert + assert "decoder.lm_head" in qcfg.modules_to_not_convert + assert not any(name.startswith("decoder.model.layers") for name in qcfg.modules_to_not_convert) + + +def test_prepare_model_rejects_component_source_path_missing_at_runtime(input_model, monkeypatch): + root_model = _make_nested_decoder_root(input_model) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={"component_source_paths": ["model.language_model"]}, + ) + + with pytest.raises(ValueError, match=r"model\.language_model.*named_modules"): + prepare_model(model, _baseline_pass_config()) + + +def test_prepare_model_rejects_selected_component_without_source_paths(input_model): + model = HfModelHandler( + input_model.model_path, + model_attributes={"component_name": "decoder", "component_role": "decoder"}, + ) + + with pytest.raises(ValueError, match="no runtime source paths"): + prepare_model(model, _baseline_pass_config()) + + +def test_prepare_model_whole_encoder_component_uses_generic_wrapper(input_model, monkeypatch): + root_model = BertForSequenceClassification( + BertConfig( # pylint: disable=unexpected-keyword-arg + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + vocab_size=128, + ) + ) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={"component_name": "model", "component_role": "encoder"}, + ) + + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + assert wrapper.hidden_size == root_model.config.hidden_size + assert wrapper.num_attention_heads == root_model.config.num_attention_heads + assert wrapper.num_hidden_layers == root_model.config.num_hidden_layers + assert wrapper.get_layer_wrappers() == [] + assert hasattr(root_model.bert.encoder.layer[0].attention.self.query, "quant_info") + assert hasattr(root_model.classifier, "quant_info") + assert "bert.embeddings.word_embeddings" in qcfg.modules_to_not_convert + assert "bert.embeddings.position_embeddings" in qcfg.modules_to_not_convert + assert "bert.embeddings.token_type_embeddings" in qcfg.modules_to_not_convert + + +def test_finalize_whole_encoder_reloads_all_embeddings_as_float( + input_model, + monkeypatch, + tmp_path, +): + root_model = BertForSequenceClassification( + BertConfig( # pylint: disable=unexpected-keyword-arg + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + vocab_size=128, + ) + ) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + task="text-classification", + model_attributes={"component_name": "model", "component_role": "encoder"}, + ) + model.save_metadata = lambda *_, **__: [] + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + output_model = finalize(model, str(tmp_path), wrapper, qcfg, device="cpu") + reloaded = output_model.load_model() + + assert isinstance(reloaded.bert.encoder.layer[0].attention.self.query, QuantLinear) + assert isinstance(reloaded.classifier, QuantLinear) + assert isinstance(reloaded.bert.embeddings.word_embeddings, torch.nn.Embedding) + assert isinstance(reloaded.bert.embeddings.position_embeddings, torch.nn.Embedding) + assert isinstance(reloaded.bert.embeddings.token_type_embeddings, torch.nn.Embedding) + + +def test_layerwise_quantization_rejects_embedding_role_with_decoder_wrapper(input_model, monkeypatch): + root_model = _NestedBackboneRoot(LlamaForCausalLM.from_pretrained(input_model.model_path)) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={ + "component_role": "embedding", + "component_source_paths": ["model.language_model.embed_tokens"], + }, + ) + wrapper, _, _ = prepare_model(model, _baseline_pass_config(embeds=True)) + + with pytest.raises(ValueError, match="Layerwise calibration requires a decoder"): + run_layerwise_quantization( + model, + wrapper, + data_config=None, + input_hook=lambda *_: None, + process_module=lambda *_: None, + update_before_process=False, + include_lm_head=False, + ) + + +def test_prepare_model_multi_path_component_slices_common_ancestor(input_model, monkeypatch): + """A multi-path component slices to the common ancestor, quantizing only declared sub-trees.""" + root_model = _make_nested_decoder_root(input_model) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + # decoder.model.layers (transformer blocks) + decoder.lm_head, common ancestor = "decoder". + model = HfModelHandler( + input_model.model_path, + model_attributes={"component_source_paths": ["decoder.model.layers", "decoder.lm_head"]}, + ) + + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + # Sliced to the common ancestor submodule, not the whole root. + assert wrapper.model is root_model.decoder + # Linear inside a declared sub-tree is quantized. + assert hasattr(root_model.decoder.model.layers[0].self_attn.q_proj, "quant_info") + # A sibling module inside the slice but outside the declared sub-trees is excluded. + assert "decoder.model.embed_tokens" in qcfg.modules_to_not_convert + # A module outside the slice entirely is excluded. + assert "vision" in qcfg.modules_to_not_convert + + +def test_finalize_multi_path_vlm_decoder_quantizes_and_saves_full_model( + input_model, + monkeypatch, + tmp_path, +): + root_model = _NestedBackboneRoot(LlamaForCausalLM.from_pretrained(input_model.model_path)) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={ + "component_source_paths": [ + "model.language_model.layers", + "model.language_model.norm", + "lm_head", + ] + }, + ) + model.save_metadata = lambda *_, **__: [] + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + finalize(model, str(tmp_path), wrapper, qcfg, device="cpu") + + assert isinstance(root_model.model.language_model.layers[0].self_attn.q_proj, QuantLinear) + assert isinstance(root_model.model.language_model.embed_tokens, torch.nn.Embedding) + assert isinstance(root_model.lm_head, torch.nn.Linear) + assert isinstance(root_model.vision, torch.nn.Linear) + assert any( + key.startswith("model.language_model.layers.0.self_attn.q_proj.qweight") for key in root_model.saved_state_keys + ) + assert "vision.weight" in root_model.saved_state_keys + + +def test_finalize_bart_decoder_reloads_unquantized_modules_as_float( + input_model, + monkeypatch, + tmp_path, +): + root_model = BartForConditionalGeneration( + BartConfig( # pylint: disable=unexpected-keyword-arg + d_model=16, + encoder_layers=1, + decoder_layers=1, + encoder_attention_heads=4, + decoder_attention_heads=4, + encoder_ffn_dim=32, + decoder_ffn_dim=32, + vocab_size=128, + ) + ) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + task="text2text-generation", + model_attributes={ + "component_role": "decoder", + "component_source_paths": ["model.decoder", "lm_head"], + }, + ) + model.save_metadata = lambda *_, **__: [] + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + output_model = finalize(model, str(tmp_path), wrapper, qcfg, device="cpu") + reloaded = output_model.load_model() + + assert isinstance(reloaded.model.decoder.layers[0].self_attn.q_proj, QuantLinear) + assert isinstance(reloaded.model.decoder.embed_tokens, torch.nn.Embedding) + assert isinstance(reloaded.model.decoder.embed_positions, torch.nn.Embedding) + assert isinstance(reloaded.model.encoder.layers[0].self_attn.q_proj, torch.nn.Linear) + assert isinstance(reloaded.lm_head, torch.nn.Linear) + + +def test_finalize_bart_decoder_embeddings_preserves_untied_float_weights( + input_model, + monkeypatch, + tmp_path, +): + root_model = BartForConditionalGeneration( + BartConfig( # pylint: disable=unexpected-keyword-arg + d_model=16, + encoder_layers=1, + decoder_layers=1, + encoder_attention_heads=4, + decoder_attention_heads=4, + encoder_ffn_dim=32, + decoder_ffn_dim=32, + vocab_size=128, + ) + ) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + task="text2text-generation", + model_attributes={ + "component_role": "decoder", + "component_source_paths": ["model.decoder", "lm_head"], + }, + ) + model.save_metadata = lambda *_, **__: [] + wrapper, qcfg, retie = prepare_model(model, _baseline_pass_config(embeds=True)) + expected_shared = root_model.model.shared.weight.detach().clone() + expected_encoder = root_model.model.encoder.embed_tokens.weight.detach().clone() + expected_head = root_model.lm_head.weight.detach().clone() + + output_model = finalize( + model, + str(tmp_path), + wrapper, + qcfg, + device="cpu", + retie_word_embeddings=retie, + ) + reloaded = output_model.load_model() + + assert not retie + assert isinstance(reloaded.model.decoder.embed_tokens, QuantEmbedding) + assert isinstance(reloaded.model.decoder.embed_positions, QuantEmbedding) + assert isinstance(reloaded.model.shared, torch.nn.Embedding) + assert isinstance(reloaded.model.encoder.embed_tokens, torch.nn.Embedding) + assert isinstance(reloaded.lm_head, torch.nn.Linear) + torch.testing.assert_close(reloaded.model.shared.weight, expected_shared) + torch.testing.assert_close(reloaded.model.encoder.embed_tokens.weight, expected_encoder) + torch.testing.assert_close(reloaded.lm_head.weight, expected_head) + + +def test_finalize_t5_shared_embedding_preserves_float_aliases( + input_model, + monkeypatch, + tmp_path, +): + root_model = T5ForConditionalGeneration( + T5Config( # pylint: disable=unexpected-keyword-arg + d_model=16, + d_ff=32, + num_layers=1, + num_decoder_layers=1, + num_heads=4, + vocab_size=128, + ) + ) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + task="text2text-generation", + model_attributes={ + "component_role": "embedding", + "component_source_paths": ["shared"], + }, + ) + model.save_metadata = lambda *_, **__: [] + wrapper, qcfg, retie = prepare_model(model, _baseline_pass_config(embeds=True)) + expected_encoder = root_model.encoder.embed_tokens.weight.detach().clone() + expected_decoder = root_model.decoder.embed_tokens.weight.detach().clone() + expected_head = root_model.lm_head.weight.detach().clone() + + output_model = finalize( + model, + str(tmp_path), + wrapper, + qcfg, + device="cpu", + retie_word_embeddings=retie, + ) + reloaded = output_model.load_model() + + assert not retie + assert isinstance(reloaded.shared, QuantEmbedding) + assert isinstance(reloaded.encoder.embed_tokens, torch.nn.Embedding) + assert isinstance(reloaded.decoder.embed_tokens, torch.nn.Embedding) + assert isinstance(reloaded.lm_head, torch.nn.Linear) + torch.testing.assert_close(reloaded.encoder.embed_tokens.weight, expected_encoder) + torch.testing.assert_close(reloaded.decoder.embed_tokens.weight, expected_decoder) + torch.testing.assert_close(reloaded.lm_head.weight, expected_head) + + +def test_prepare_model_removes_current_component_from_existing_exclusions( + input_model, + monkeypatch, +): + existing = OliveHfQuantizationConfig( + bits=PrecisionBits.BITS4, + symmetric=False, + group_size=16, + modules_to_not_convert=["model.language_model.embed_tokens", "vision"], + ) + _with_existing_quantization_config(monkeypatch, existing) + root_model = _NestedBackboneRoot(LlamaForCausalLM.from_pretrained(input_model.model_path)) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={"component_source_paths": ["model.language_model.embed_tokens"]}, + ) + + _, qcfg, _ = prepare_model( + model, + _baseline_pass_config(embeds=True), + allow_quantized=True, + ) + + assert "model.language_model.embed_tokens" not in qcfg.modules_to_not_convert + assert "vision" in qcfg.modules_to_not_convert + + +def test_finalize_vlm_encoder_component_only_quantizes_encoder( + input_model, + monkeypatch, + tmp_path, +): + root_model = _NestedBackboneRoot(LlamaForCausalLM.from_pretrained(input_model.model_path)) + root_model.vision = torch.nn.Linear(16, 16) + monkeypatch.setattr(quant_utils_module, "load_hf_base_model", lambda _: root_model) + model = HfModelHandler( + input_model.model_path, + model_attributes={ + "component_role": "encoder", + "component_source_paths": ["vision"], + }, + ) + model.save_metadata = lambda *_, **__: [] + wrapper, qcfg, _ = prepare_model(model, _baseline_pass_config()) + + finalize(model, str(tmp_path), wrapper, qcfg, device="cpu") + + assert isinstance(root_model.vision, QuantLinear) + assert isinstance(root_model.model.language_model.layers[0].self_attn.q_proj, torch.nn.Linear) + assert isinstance(root_model.model.language_model.embed_tokens, torch.nn.Embedding) + assert isinstance(root_model.lm_head, torch.nn.Linear) + + def test_prepare_model_promotes_user_override_conflicts_for_qkv(input_model): """User-supplied overrides on K/V promote Q to the most-precise shared config.""" model = HfModelHandler( @@ -494,8 +919,6 @@ def test_prepare_model_locks_default_quantized_qkv_member_without_override(input for V -- that would disagree with V's on-disk weights. Instead, Q/K should be demoted to V's existing default config. """ - from olive.common.quant.nn import QuantLinear - qu = quant_utils_module existing = { diff --git a/test/test_cache.py b/test/test_cache.py index 4063ff1554..23fa96c459 100644 --- a/test/test_cache.py +++ b/test/test_cache.py @@ -3,19 +3,56 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- import json +import os import shutil +from concurrent.futures import ThreadPoolExecutor from pathlib import Path +from threading import Barrier from unittest.mock import ANY, mock_open, patch import pytest -from olive.cache import CacheConfig, OliveCache, SharedCache +from olive.cache import CacheConfig, OliveCache, SharedCache, isolated_cache_env from olive.common.constants import DEFAULT_WORKFLOW_ID +from olive.systems.utils import create_new_environ # pylint: disable=W0201 class TestCache: + def test_cache_env_uses_process_environment_outside_isolated_context(self, tmp_path, monkeypatch): + first_cache_dir = tmp_path / "first" + second_cache_dir = tmp_path / "second" + OliveCache({"cache_dir": first_cache_dir}).set_cache_env() + monkeypatch.setenv("OLIVE_CACHE_DIR", str(second_cache_dir)) + + assert OliveCache.from_cache_env().get_cache_dir() == second_cache_dir + + def test_cache_env_is_scoped_to_parallel_build_thread(self, tmp_path, monkeypatch): + barrier = Barrier(2) + process_cache_dir = tmp_path / "process" + monkeypatch.setenv("OLIVE_CACHE_DIR", str(process_cache_dir)) + + def set_and_read_cache(cache_dir): + with isolated_cache_env(cache_dir): + cache = OliveCache({"cache_dir": cache_dir}) + cache.set_cache_env() + barrier.wait(timeout=2) + return ( + OliveCache.from_cache_env().get_cache_dir(), + Path(create_new_environ()["OLIVE_CACHE_DIR"]), + Path(os.environ["OLIVE_CACHE_DIR"]), + ) + + cache_dirs = [tmp_path / "first", tmp_path / "second"] + with ThreadPoolExecutor(max_workers=2) as executor: + resolved = list(executor.map(set_and_read_cache, cache_dirs)) + + assert [item[0] for item in resolved] == [path.resolve() for path in cache_dirs] + assert [item[1] for item in resolved] == cache_dirs + assert all(item[2] == process_cache_dir for item in resolved) + assert Path(os.environ["OLIVE_CACHE_DIR"]) == process_cache_dir + @pytest.mark.parametrize("clean_cache", [True, False]) def test_cache_init(self, clean_cache, tmp_path): # setup diff --git a/test/workflows/test_run_builds.py b/test/workflows/test_run_builds.py new file mode 100644 index 0000000000..8a369ec9f1 --- /dev/null +++ b/test/workflows/test_run_builds.py @@ -0,0 +1,420 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +import sys +import time +from copy import deepcopy +from pathlib import Path +from threading import Barrier +from unittest.mock import MagicMock, patch + +import pytest + +from olive.workflows import run as olive_run +from test.utils import get_pytorch_model_io_config, pytorch_model_loader + +# pylint: disable=attribute-defined-outside-init + +PT_MODEL = { + "type": "PyTorchModel", + "config": { + "model_loader": pytorch_model_loader, + "io_config": get_pytorch_model_io_config(), + }, +} + + +class TestRunBuilds: + @pytest.fixture(autouse=True) + def setup(self, tmp_path): + self.cache_dir = tmp_path / "cache" + self.template = { + "input_model": PT_MODEL, + "systems": { + "cpu_system": {"type": "LocalSystem", "accelerators": [{"device": "cpu"}]}, + "gpu_system": {"type": "LocalSystem", "accelerators": [{"device": "gpu"}]}, + }, + "passes": { + "convert": {"type": "OnnxConversion"}, + "tune": {"type": "OrtSessionParamsTuning"}, + }, + "engine": { + "evaluate_input_model": False, + "cache_dir": str(self.cache_dir), + }, + } + + def _patch_engine_and_acc(self): + run_mock = MagicMock(return_value=MagicMock(name="WorkflowOutput")) + acc_mock = MagicMock(name="accelerator_spec") + engine_run_patch = patch("olive.engine.engine.Engine.run", run_mock) + accelerator_patch = patch.object(sys.modules[olive_run.__module__], "create_accelerator", return_value=acc_mock) + return run_mock, acc_mock, engine_run_patch, accelerator_patch + + def test_builds_passes_per_build_pipeline_subset_in_declared_order(self): + # `tune` declared first in passes but second in pipeline; engine should receive [convert, tune]. + run_mock, _, engine_run_patch, acc_patch = self._patch_engine_and_acc() + captured: list = [] + + def capture_input_passes(self, pass_configs): + captured.append(list(pass_configs)) + + config = deepcopy(self.template) + config["passes"] = { + "tune": {"type": "OrtSessionParamsTuning"}, + "convert": {"type": "OnnxConversion"}, + } + config["builds"] = { + "only": {"pipeline": ["convert", "tune"], "output_dir": "out/only"}, + } + with ( + engine_run_patch, + acc_patch, + patch("olive.engine.engine.Engine.set_input_passes_configs", capture_input_passes), + ): + olive_run(config) + assert run_mock.call_count == 1 + assert captured == [["convert", "tune"]] + + def test_builds_uses_per_build_output_dir(self): + run_mock, _, engine_run_patch, acc_patch = self._patch_engine_and_acc() + config = deepcopy(self.template) + config["builds"] = { + "first": {"pipeline": ["convert"], "output_dir": "out/first"}, + "second": {"pipeline": ["convert"], "output_dir": "out/second"}, + } + with engine_run_patch, acc_patch: + olive_run(config) + output_dirs = [call.args[3] for call in run_mock.call_args_list] + assert set(output_dirs) == {Path("out/first").resolve(), Path("out/second").resolve()} + + def test_builds_host_target_override_applied_per_build(self): + # Captures the SystemConfig passed to create_accelerator for each build. + run_mock, acc_mock, engine_run_patch, _ = self._patch_engine_and_acc() + seen_targets: list = [] + + def fake_create_accelerator(system_config, **kwargs): + seen_targets.append(system_config.config.accelerators[0].device.lower()) + return acc_mock + + config = deepcopy(self.template) + config["builds"] = { + "cpu_build": { + "pipeline": ["convert"], + "output_dir": "out/cpu", + "host": "cpu_system", + "target": "cpu_system", + }, + "gpu_build": { + "pipeline": ["convert"], + "output_dir": "out/gpu", + "host": "gpu_system", + "target": "gpu_system", + }, + } + with ( + engine_run_patch, + patch.object(sys.modules[olive_run.__module__], "create_accelerator", side_effect=fake_create_accelerator), + ): + olive_run(config) + assert run_mock.call_count == 2 + assert set(seen_targets) == {"cpu", "gpu"} + + def test_builds_components_on_non_composite_input_raises(self): + config = deepcopy(self.template) + config["builds"] = { + "broken": { + "pipeline": ["convert"], + "output_dir": "out/broken", + "components": ["text_encoder"], + }, + } + with pytest.raises(ValueError, match="select_components is only supported"): + olive_run(config) + + def test_builds_components_unknown_name_raises(self): + composite_input = { + "type": "CompositeModel", + "config": { + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": "a.onnx"}}, + {"type": "ONNXModel", "config": {"model_path": "b.onnx"}}, + ], + "model_component_names": ["text_encoder", "unet"], + }, + } + config = deepcopy(self.template) + config["input_model"] = composite_input + config["builds"] = { + "bad": { + "pipeline": ["convert"], + "output_dir": "out/bad", + "components": ["no_such_component"], + }, + } + with pytest.raises(ValueError, match="Unknown component"): + olive_run(config) + + def test_builds_directory_composite_input_runs_per_component(self, tmp_path): + # An export directory loads as a CompositeModel; subfolder names become + # component names, and sibling builds optimize each component. + for name in ["decoder", "vision_encoder"]: + comp_dir = tmp_path / "exported_pkg" / name + comp_dir.mkdir(parents=True) + (comp_dir / "model.onnx").write_bytes(b"onnx") + + run_mock, _, engine_run_patch, acc_patch = self._patch_engine_and_acc() + config = deepcopy(self.template) + config["input_model"] = {"type": "CompositeModel", "config": {"model_path": str(tmp_path / "exported_pkg")}} + config["builds"] = { + "decoder": {"components": ["decoder"], "pipeline": ["convert"], "output_dir": "out/decoder"}, + "vision_encoder": {"components": ["vision_encoder"], "pipeline": ["convert"], "output_dir": "out/vision"}, + } + with engine_run_patch, acc_patch: + result = olive_run(config) + assert set(result) == {"decoder", "vision_encoder"} + assert run_mock.call_count == 2 + + def test_builds_prevalidate_all_configs_before_running(self): + config = deepcopy(self.template) + config["builds"] = { + "valid": {"pipeline": ["convert"], "output_dir": "out/valid"}, + "broken": { + "pipeline": ["convert"], + "output_dir": "out/broken", + "search_strategy": True, + }, + } + + with ( + patch("olive.engine.engine.Engine.run") as run_mock, + pytest.raises(ValueError, match="Invalid build 'broken'"), + ): + olive_run(config) + run_mock.assert_not_called() + + def test_builds_union_required_packages_across_expanded_configs(self): + config = deepcopy(self.template) + config["builds"] = { + "first": {"pipeline": ["convert"], "output_dir": "out/first"}, + "second": {"pipeline": ["tune"], "output_dir": "out/second"}, + } + run_module = sys.modules[olive_run.__module__] + + with ( + patch.object(run_module, "get_required_packages", side_effect=[{"package-a"}, {"package-b"}]), + patch.object(run_module, "generate_files_from_packages") as generate_mock, + ): + olive_run(config, list_required_packages=True) + + generate_mock.assert_called_once_with({"package-a", "package-b"}, "olive_requirements.txt") + + def test_builds_use_standard_docker_workflow_dispatch(self): + config = deepcopy(self.template) + config["systems"]["docker_system"] = { + "type": "Docker", + "dockerfile": "Dockerfile", + "build_context_path": ".", + } + config["builds"] = { + "docker": { + "pipeline": ["convert"], + "output_dir": "out/docker", + "host": "docker_system", + }, + } + docker_system = MagicMock() + docker_output = MagicMock(name="docker_output") + docker_system.run_workflow.return_value = docker_output + + with patch("olive.systems.system_config.SystemConfig.create_system", return_value=docker_system): + result = olive_run(config) + + assert result == {"docker": docker_output} + docker_run_config = docker_system.run_workflow.call_args.args[0] + assert docker_run_config.workflow_id == "default_workflow_docker" + + def test_builds_share_docker_image_while_containers_run_in_parallel(self): + config = deepcopy(self.template) + config["systems"]["docker_system"] = { + "type": "Docker", + "dockerfile": "Dockerfile", + "build_context_path": ".", + } + config["builds"] = { + "first": {"pipeline": ["convert"], "output_dir": "out/first", "host": "docker_system"}, + "second": {"pipeline": ["tune"], "output_dir": "out/second", "host": "docker_system"}, + } + barrier = Barrier(2) + docker_systems = [MagicMock(), MagicMock()] + for docker_system in docker_systems: + docker_system.run_workflow.side_effect = lambda run_config: ( + barrier.wait(timeout=2), + run_config.workflow_id, + )[1] + + with patch( + "olive.systems.system_config.SystemConfig.create_system", + side_effect=docker_systems, + ) as create_system_mock: + result = olive_run(config) + + assert result == { + "first": "default_workflow_first", + "second": "default_workflow_second", + } + assert create_system_mock.call_count == 2 + assert all(docker_system.clean_image is False for docker_system in docker_systems) + docker_systems[0].remove.assert_called_once_with() + docker_systems[1].remove.assert_not_called() + + def test_builds_reject_conflicting_shared_docker_image_settings(self): + config = deepcopy(self.template) + config["systems"].update( + { + "first_docker": { + "type": "Docker", + "dockerfile": "Dockerfile.first", + "build_context_path": ".", + "image_name": "olive-docker", + }, + "second_docker": { + "type": "Docker", + "dockerfile": "Dockerfile.second", + "build_context_path": ".", + "image_name": "olive-docker:latest", + }, + } + ) + config["builds"] = { + "first": {"pipeline": ["convert"], "output_dir": "out/first", "host": "first_docker"}, + "second": {"pipeline": ["tune"], "output_dir": "out/second", "host": "second_docker"}, + } + + with ( + patch("olive.systems.system_config.SystemConfig.create_system") as create_system_mock, + pytest.raises(ValueError, match="conflicting image build settings"), + ): + olive_run(config) + create_system_mock.assert_not_called() + + def test_builds_clean_prepared_shared_docker_image_when_later_preparation_fails(self): + config = deepcopy(self.template) + config["systems"]["docker_system"] = { + "type": "Docker", + "dockerfile": "Dockerfile", + "build_context_path": ".", + } + config["builds"] = { + "first": {"pipeline": ["convert"], "output_dir": "out/first", "host": "docker_system"}, + "second": {"pipeline": ["tune"], "output_dir": "out/second", "host": "docker_system"}, + } + first_system = MagicMock() + + with ( + patch( + "olive.systems.system_config.SystemConfig.create_system", + side_effect=[first_system, RuntimeError("preparation failed")], + ), + pytest.raises(RuntimeError, match="preparation failed"), + ): + olive_run(config) + + first_system.remove.assert_called_once_with() + + def test_builds_run_in_parallel_and_preserve_configured_result_order(self): + config = deepcopy(self.template) + config["builds"] = { + "slow": {"pipeline": ["convert"], "output_dir": "out/slow"}, + "fast": {"pipeline": ["tune"], "output_dir": "out/fast"}, + } + barrier = Barrier(2) + + def run_build(_, run_config): + barrier.wait(timeout=2) + if run_config.workflow_id.endswith("_slow"): + time.sleep(0.05) + return run_config.workflow_id + + with patch.object(sys.modules[olive_run.__module__], "_run_single", side_effect=run_build): + result = olive_run(config) + + assert result == { + "slow": "default_workflow_slow", + "fast": "default_workflow_fast", + } + + def test_builds_report_build_name_when_parallel_execution_fails(self): + config = deepcopy(self.template) + config["builds"] = { + "good": {"pipeline": ["convert"], "output_dir": "out/good"}, + "bad": {"pipeline": ["tune"], "output_dir": "out/bad"}, + } + barrier = Barrier(2) + + def run_build(_, run_config): + barrier.wait(timeout=2) + if run_config.workflow_id.endswith("_bad"): + raise ValueError("bad build") + return run_config.workflow_id + + with ( + patch.object(sys.modules[olive_run.__module__], "_run_single", side_effect=run_build), + pytest.raises(RuntimeError, match=r"Build\(s\) \['bad'\] failed.*bad build"), + ): + olive_run(config) + + def test_builds_report_all_parallel_execution_failures(self): + config = deepcopy(self.template) + config["builds"] = { + "first": {"pipeline": ["convert"], "output_dir": "out/first"}, + "second": {"pipeline": ["tune"], "output_dir": "out/second"}, + } + + def run_build(_, run_config): + raise ValueError(f"{run_config.workflow_id} error") + + with ( + patch.object(sys.modules[olive_run.__module__], "_run_single", side_effect=run_build), + pytest.raises( + RuntimeError, + match=( + r"Build\(s\) \['first', 'second'\] failed:.*default_workflow_first error" + r".*default_workflow_second error" + ), + ), + ): + olive_run(config) + + def test_builds_reject_overlapping_parallel_artifact_directories(self): + config = deepcopy(self.template) + config["builds"] = { + "parent": {"pipeline": ["convert"], "output_dir": "out/shared"}, + "child": {"pipeline": ["tune"], "output_dir": "out/shared/child"}, + } + + with ( + patch.object(sys.modules[olive_run.__module__], "_run_single") as run_mock, + pytest.raises(ValueError, match="overlapping writable directories"), + ): + olive_run(config) + run_mock.assert_not_called() + + def test_builds_reject_artifact_directory_overlapping_another_build_cache(self): + config = deepcopy(self.template) + config["builds"] = { + "artifact": { + "pipeline": ["convert"], + "output_dir": str(self.cache_dir / "default_workflow_cached"), + }, + "cached": {"pipeline": ["tune"], "output_dir": "out/cached"}, + } + + with ( + patch.object(sys.modules[olive_run.__module__], "_run_single") as run_mock, + pytest.raises(ValueError, match=r"artifact directory .* and cache directory"), + ): + olive_run(config) + run_mock.assert_not_called() diff --git a/test/workflows/test_run_config_builds.py b/test/workflows/test_run_config_builds.py new file mode 100644 index 0000000000..802015495e --- /dev/null +++ b/test/workflows/test_run_config_builds.py @@ -0,0 +1,180 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +from copy import deepcopy + +import pytest +from pydantic import ValidationError + +from olive.workflows.run.builds import expand_builds +from olive.workflows.run.config import RunConfig + +# pylint: disable=attribute-defined-outside-init + + +class TestBuildConfigExpansion: + @pytest.fixture(autouse=True) + def setup(self): + self.template = { + "workflow_id": "multi", + "input_model": { + "type": "HfModel", + "model_path": "dummy_model", + "task": "dummy_task", + }, + "systems": { + "local_system": {"type": "LocalSystem", "accelerators": [{"device": "cpu"}]}, + "other_system": {"type": "LocalSystem", "accelerators": [{"device": "gpu"}]}, + }, + "passes": { + "convert": {"type": "OnnxConversion"}, + "tune": {"type": "OrtSessionParamsTuning"}, + }, + "evaluate_input_model": False, + } + + def _expand(self, builds): + config_dict = deepcopy(self.template) + config_dict["builds"] = builds + return expand_builds(config_dict) + + def test_builds_default_values_are_fully_replaced(self): + expanded = self._expand( + { + "_default": { + "pipeline": ["convert", "tune"], + "output_dir": "out/default", + "host": "local_system", + }, + "override": { + "pipeline": ["convert"], + "output_dir": "out/override", + "host": "other_system", + }, + "inherit": {}, + } + ) + + assert list(expanded) == ["override", "inherit"] + assert list(expanded["override"]["passes"]) == ["convert"] + assert expanded["override"]["engine"]["output_dir"] == "out/override" + assert expanded["override"]["engine"]["host"] == "other_system" + assert list(expanded["inherit"]["passes"]) == ["convert", "tune"] + assert expanded["inherit"]["engine"]["output_dir"] == "out/default" + assert expanded["inherit"]["engine"]["host"] == "local_system" + + def test_builds_assign_unique_workflow_ids(self): + expanded = self._expand( + { + "first": {"pipeline": ["convert"], "output_dir": "out/first"}, + "second": {"pipeline": ["tune"], "output_dir": "out/second"}, + } + ) + + assert expanded["first"]["workflow_id"] == "multi_first" + assert expanded["second"]["workflow_id"] == "multi_second" + + def test_builds_do_not_mutate_source_config(self): + config_dict = deepcopy(self.template) + config_dict["builds"] = { + "only": {"pipeline": ["convert"], "output_dir": "out/only"}, + } + original = deepcopy(config_dict) + + expand_builds(config_dict) + + assert config_dict == original + + def test_builds_missing_pipeline_after_merge_errors(self): + with pytest.raises(ValidationError, match="pipeline"): + self._expand( + { + "_default": {"host": "local_system"}, + "broken": {}, + } + ) + + @pytest.mark.parametrize("builds", [None, False, [], {}, {"_default": {}}]) + def test_builds_requires_at_least_one_named_build(self, builds): + with pytest.raises(ValueError, match="builds"): + self._expand(builds) + + @pytest.mark.parametrize("default", [None, False, []]) + def test_builds_default_must_be_a_dictionary(self, default): + with pytest.raises(ValueError, match=r"builds\._default"): + self._expand( + { + "_default": default, + "only": {"pipeline": ["convert"], "output_dir": "out/only"}, + } + ) + + @pytest.mark.parametrize( + ("build", "field_name"), + [ + ({"pipeline": [], "output_dir": "out/only"}, "pipeline"), + ({"pipeline": ["convert"], "output_dir": ""}, "output_dir"), + ({"pipeline": ["convert"], "output_dir": "out/only", "components": []}, "components"), + ], + ) + def test_builds_rejects_empty_required_values(self, build, field_name): + with pytest.raises(ValidationError, match=field_name): + self._expand({"only": build}) + + @pytest.mark.parametrize("build_name", ["", "has space", "../escape"]) + def test_builds_rejects_unsafe_names(self, build_name): + with pytest.raises(ValueError, match="Invalid build name"): + self._expand({build_name: {"pipeline": ["convert"], "output_dir": "out/only"}}) + + def test_builds_invalid_pipeline_reference_errors(self): + with pytest.raises(ValueError, match="unknown pass"): + self._expand( + { + "broken": { + "pipeline": ["convert", "no_such_pass"], + "output_dir": "out/broken", + }, + } + ) + + def test_expanded_build_uses_standard_system_reference_validation(self): + expanded = self._expand( + { + "broken": { + "pipeline": ["convert"], + "output_dir": "out/broken", + "host": "no_such_system", + }, + } + ) + + with pytest.raises(ValidationError, match="not found"): + RunConfig.model_validate(expanded["broken"]) + + def test_empty_default_is_noop(self): + expanded = self._expand( + { + "_default": {}, + "only": {"pipeline": ["convert"], "output_dir": "out/only"}, + } + ) + + assert list(expanded) == ["only"] + assert list(expanded["only"]["passes"]) == ["convert"] + assert "host" not in expanded["only"]["engine"] + + def test_expanded_build_uses_standard_search_validation(self): + expanded = self._expand( + { + "only": { + "pipeline": ["convert"], + "output_dir": "out/only", + "search_strategy": True, + }, + } + ) + + with pytest.raises(ValidationError, match="valid evaluator"): + RunConfig.model_validate(expanded["only"])