diff --git a/docs/source/features/onnx-transformations.md b/docs/source/features/onnx-transformations.md index 7d713615c..558009008 100644 --- a/docs/source/features/onnx-transformations.md +++ b/docs/source/features/onnx-transformations.md @@ -142,6 +142,8 @@ operator is not a portable standard-ONNX optimization. | `SeparateGroupQueryAttentionRoPE` | Move supported GQA-integrated RoPE into separate standard `RotaryEmbedding` nodes. | | `UnpackGroupQueryAttentionQKV` | Split supported packed GQA projections into separate Q/K/V projections. | | `BlockDiagonalAttentionToPackedMHA` | Recognized block-diagonal mask and standard `Attention` to `com.microsoft::PackedMultiHeadAttention`. | +| `ConvertGroupQueryAttentionKVCacheToFp8` | Retype GQA past/present KV-cache I/O to FP8 E4M3 and attach per-layer scale inputs. | +| `InlineModelLocalFunctions` | Inline all model-local function calls and reject any remaining non-standard-domain node without a standard fallback body. | | `ClipToMinMax` | BF16 `Clip` with both bounds, only a lower bound, or only an upper bound to `Max`/`Min`. | | `Rank4RMSNormToRank3` | Rank-4 last-axis RMSNorm with static head dimensions to rank-3 RMSNorm surrounded by reshapes. | | `DecomposeOnnxRotaryEmbedding` | Standard rank-3, full-width, non-interleaved `RotaryEmbedding` to primitive rotate-half operations. | diff --git a/olive/cli/capture_onnx.py b/olive/cli/capture_onnx.py index 3966b798c..cdd02d189 100644 --- a/olive/cli/capture_onnx.py +++ b/olive/cli/capture_onnx.py @@ -19,6 +19,8 @@ update_shared_cache_options, ) from olive.common.utils import set_nested_dict_value +from olive.hardware.constants import DEVICE_TO_EXECUTION_PROVIDERS, ExecutionProvider +from olive.model import ModelConfig from olive.model.utils.diffusers_utils import is_valid_diffusers_model from olive.telemetry import action @@ -30,6 +32,56 @@ class ModelBuilderAccuracyLevel(IntEnum): int8 = 4 +_EP_ALIASES = { + "cuda": ExecutionProvider.CUDAExecutionProvider, + "openvino": ExecutionProvider.OpenVINOExecutionProvider, + "qnn": ExecutionProvider.QNNExecutionProvider, + "trt-rtx": ExecutionProvider.NvTensorRTRTXExecutionProvider, + "vitisai": ExecutionProvider.VitisAIExecutionProvider, +} + + +def _surgery(name: str, **kwargs) -> dict: + return {"surgeon": name, **kwargs} + + +def _resolve_recipe_ep_profile(ep: str, device: str) -> tuple[ExecutionProvider, list[dict]]: + provider = _EP_ALIASES[ep] + if provider not in DEVICE_TO_EXECUTION_PROVIDERS[device]: + raise ValueError(f"Execution provider {ep!r} does not support device {device!r}.") + + if ep == "cuda": + return provider, [ + _surgery( + "AttentionToGroupQueryAttention", + supported_dtypes=["FLOAT16", "BFLOAT16"], + ), + _surgery("PackQKVForGroupQueryAttention"), + _surgery("FuseSkipRMSNormalization"), + _surgery("FuseSkipLayerNormalization"), + ] + if ep == "qnn": + surgeries = [ + _surgery("AttentionToGroupQueryAttention"), + _surgery("PackQKVForGroupQueryAttention"), + _surgery("FuseSkipRMSNormalization"), + _surgery("AttentionMaskToSequenceLengths"), + ] + if device == "npu": + surgeries.append(_surgery("SimplifiedLayerNormToL2Norm")) + return provider, surgeries + if ep == "trt-rtx": + return provider, [ + _surgery( + "AttentionToGroupQueryAttention", + supported_dtypes=["FLOAT16", "BFLOAT16"], + ), + _surgery("PackQKVForGroupQueryAttention"), + ] + + return provider, [] + + def parse_dim_dict(s): try: return {k: int(v) if v.isdigit() else v for k, v in (item.split("=") for item in s.split(","))} @@ -84,10 +136,32 @@ def register_subcommand(parser: ArgumentParser): help=( "Whether to use MobiusBuilder (mobius-onnx) to capture ONNX model. " "Supports multi-component multimodal models (VLMs). " + "Preserves model precision and exports Mobius' canonical graph before applying " + "optional strict-ONNX expansion or execution-provider graph surgeries. " "Requires 'pip install mobius-onnx'." ), ) + mobius_ep_group = sub_parser.add_argument_group("Mobius Builder graph surgery options") + mobius_ep_group.add_argument( + "--execution_provider", + choices=sorted(_EP_ALIASES), + help="Target execution provider used to select post-export graph surgeries.", + ) + mobius_ep_group.add_argument( + "--device", + choices=["cpu", "gpu", "npu"], + help="Target device used with --execution_provider to select post-export graph surgeries.", + ) + mobius_ep_group.add_argument( + "--onnx_standard", + action="store_true", + help=( + "Expand all model-local non-standard functions after Mobius export. " + "May be combined with --execution_provider and --device to retain the target build contract." + ), + ) + # PyTorch Exporter options pte_group = sub_parser.add_argument_group("PyTorch Exporter options") pte_group.add_argument( @@ -130,7 +204,7 @@ def register_subcommand(parser: ArgumentParser): type=str, default="fp16", choices=["fp16", "fp32", "int4", "bf16"], - help="The precision of the ONNX model. Used by Model Builder and Mobius Builder.", + help="The precision of the ONNX model. Used by Model Builder.", ) mb_group.add_argument( "--int4_block_size", @@ -194,7 +268,32 @@ def register_subcommand(parser: ArgumentParser): @action def run(self): - return self._run_workflow() + workflow_output = self._run_workflow() + if workflow_output is None or not self.args.use_mobius_builder or not workflow_output.has_output_model(): + return workflow_output + + if self.args.onnx_standard: + surgeries = [_surgery("InlineModelLocalFunctions")] + decoder_only = False + elif self.args.execution_provider: + _, surgeries = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device) + decoder_only = True + else: + return workflow_output + if not surgeries: + return workflow_output + + from olive.workflows import run as olive_run + + exported_model = workflow_output.get_best_candidate() + post_export_config = self._get_post_export_run_config( + exported_model.olive_model_config, + surgeries, + decoder_only=decoder_only, + ) + if self.args.save_config_file: + self._save_config_file(post_export_config, file_name="graph_surgery_config.json") + return olive_run(post_export_config) def _get_run_config(self, tempdir: str) -> dict: config = deepcopy(TEMPLATE) @@ -217,36 +316,41 @@ def _get_run_config(self, tempdir: str) -> dict: is_diffusers_model = input_model_config["type"].lower() == "diffusersmodel" + if bool(self.args.execution_provider) != bool(self.args.device): + raise ValueError("--execution_provider and --device must be provided together.") + if self.args.execution_provider and not self.args.use_mobius_builder: + raise ValueError("--execution_provider and --device graph-surgery profiles require --use_mobius_builder.") + if self.args.onnx_standard and not self.args.use_mobius_builder: + raise ValueError("--onnx_standard requires --use_mobius_builder.") + # whether model is in fp16 or bf16 (currently not supported by CPU EP) is_fp16_or_bf16 = ( - ( - not self.args.use_model_builder - and not self.args.use_mobius_builder - and self.args.torch_dtype == "float16" + not self.args.use_model_builder and not self.args.use_mobius_builder and self.args.torch_dtype == "float16" + ) or (self.args.use_model_builder and self.args.precision in ("fp16", "bf16")) + + if self.args.use_mobius_builder and self.args.execution_provider: + provider, _ = _resolve_recipe_ep_profile(self.args.execution_provider, self.args.device) + device = self.args.device + elif self.args.use_mobius_builder: + provider = ExecutionProvider.CPUExecutionProvider + device = "cpu" + else: + provider = ( + ExecutionProvider.CUDAExecutionProvider if is_fp16_or_bf16 else ExecutionProvider.CPUExecutionProvider ) - or (self.args.use_model_builder and self.args.precision in ("fp16", "bf16")) - or (self.args.use_mobius_builder and self.args.precision in ("fp16", "bf16")) - ) + device = "gpu" if is_fp16_or_bf16 else "cpu" + to_replace = [ ("input_model", input_model_config), ("output_dir", self.args.output_path), ("log_severity_level", self.args.log_level), - (("systems", "local_system", "accelerators", 0, "device"), "gpu" if is_fp16_or_bf16 else "cpu"), - ( - ("systems", "local_system", "accelerators", 0, "execution_providers"), - [("CUDAExecutionProvider" if is_fp16_or_bf16 else "CPUExecutionProvider")], - ), + (("systems", "local_system", "accelerators", 0, "device"), device), + (("systems", "local_system", "accelerators", 0, "execution_providers"), [provider.value]), ] if self.args.use_mobius_builder: - if self.args.precision not in ("fp32", "fp16", "bf16"): - raise ValueError( - f"MobiusBuilder supports precisions fp32/fp16/bf16; got '{self.args.precision}'. " - "For INT4, capture in fp32/fp16/bf16 first and run a quantization pass afterwards." - ) del config["passes"]["c"] del config["passes"]["m"] - to_replace.append((("passes", "b", "precision"), self.args.precision)) elif is_diffusers_model: del config["passes"]["m"] del config["passes"]["b"] @@ -325,12 +429,43 @@ def _get_run_config(self, tempdir: str) -> dict: return config + def _get_post_export_run_config( + self, + input_model_config: dict, + surgeries: list[dict], + *, + decoder_only: bool, + ) -> dict: + model_config = ModelConfig.model_validate(input_model_config) + components = model_config.get_components() + is_multimodal = components is not None + if decoder_only and is_multimodal and "decoder" not in components: + raise ValueError( + "Execution-provider graph surgeries require a 'decoder' component, " + f"but the exported model contains {components}." + ) + + config = deepcopy(MULTIMODAL_TEMPLATE) + if not decoder_only or not is_multimodal: + config.pop("builds") + + config["input_model"] = model_config.model_dump() + config["output_dir"] = self.args.output_path + config["log_severity_level"] = self.args.log_level + accelerator = config["systems"]["local_system"]["accelerators"][0] + accelerator["device"] = self.args.device or "cpu" + accelerator["execution_providers"] = [ + _EP_ALIASES.get(self.args.execution_provider, ExecutionProvider.CPUExecutionProvider).value + ] + config["passes"]["g"]["surgeries"] = surgeries + update_shared_cache_options(config, self.args) + return config + TEMPLATE = { "systems": { "local_system": { "type": "LocalSystem", - # might need an ep option to set for model builder, it is sensitive to ep "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], } }, @@ -346,3 +481,20 @@ def _get_run_config(self, tempdir: str) -> dict: "target": "local_system", "no_artifacts": True, } + + +MULTIMODAL_TEMPLATE = { + "systems": { + "local_system": { + "type": "LocalSystem", + "accelerators": [{"device": "cpu", "execution_providers": ["CPUExecutionProvider"]}], + } + }, + "passes": {"g": {"type": "GraphSurgeries"}}, + "builds": { + "decoder": {"components": ["decoder"], "pipeline": ["g"]}, + }, + "host": "local_system", + "target": "local_system", + "no_artifacts": True, +} diff --git a/olive/olive_config.json b/olive/olive_config.json index 85ab4233e..dc7ecaeb6 100644 --- a/olive/olive_config.json +++ b/olive/olive_config.json @@ -215,7 +215,7 @@ "module_path": "olive.passes.onnx.mobius_model_builder.MobiusBuilder", "supported_providers": [ "*" ], "supported_accelerators": [ "*" ], - "supported_precisions": [ "fp32", "fp16", "bf16" ], + "supported_precisions": [ "*" ], "supported_algorithms": [ ], "supported_quantization_encodings": [ ] }, diff --git a/olive/passes/onnx/graph_surgery/__init__.py b/olive/passes/onnx/graph_surgery/__init__.py index 9503be8fd..07e5fdce0 100644 --- a/olive/passes/onnx/graph_surgery/__init__.py +++ b/olive/passes/onnx/graph_surgery/__init__.py @@ -8,6 +8,7 @@ from olive.passes.onnx.graph_surgery.attention import ( AttentionToGroupQueryAttention, BlockDiagonalAttentionToPackedMHA, + ConvertGroupQueryAttentionKVCacheToFp8, PackQKVForGroupQueryAttention, SeparateGroupQueryAttentionRoPE, UnpackGroupQueryAttentionQKV, @@ -17,6 +18,7 @@ ClipToMinMax, DecomposeAttention, DecomposeOnnxRotaryEmbedding, + InlineModelLocalFunctions, Rank4RMSNormToRank3, StaticEmptyKV, TensorScatterToScatterND, @@ -32,6 +34,7 @@ "AttentionToGroupQueryAttention", "BlockDiagonalAttentionToPackedMHA", "ClipToMinMax", + "ConvertGroupQueryAttentionKVCacheToFp8", "DecomposeAttention", "DecomposeOnnxRotaryEmbedding", "FuseBiasGelu", @@ -41,6 +44,7 @@ "FuseLayerNormalization", "FuseSkipLayerNormalization", "FuseSkipRMSNormalization", + "InlineModelLocalFunctions", "MoEGraphSurgeryError", "PackQKVForGroupQueryAttention", "ProtoSurgeon", diff --git a/olive/passes/onnx/graph_surgery/attention.py b/olive/passes/onnx/graph_surgery/attention.py index d1cbd1fa1..aa5a1ff42 100644 --- a/olive/passes/onnx/graph_surgery/attention.py +++ b/olive/passes/onnx/graph_surgery/attention.py @@ -6,6 +6,10 @@ from __future__ import annotations +import math +import re +import warnings + import numpy as np import onnx_ir as ir from onnxscript.rewriter import pattern @@ -13,11 +17,103 @@ from onnxscript.rewriter._rewrite_rule import RewriteRuleClassBase from olive.constants import MSFT_DOMAIN -from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon +from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon, Surgeon # ONNXScript binds each rule's named pattern operands to its callbacks. # pylint: disable=arguments-differ +_FP8 = ir.DataType.FLOAT8E4M3FN +_K_SCALE_INDEX = 12 +_V_SCALE_INDEX = 13 +_MIN_GQA_INPUTS_WITH_SCALES = 14 +_LAYER_ID_RE = re.compile(r"\.(\d+)\.") + + +def _validate_fp8_scales(scales: dict[int, tuple[float, float]]) -> dict[int, tuple[float, float]]: + validated = {} + for layer_id, pair in scales.items(): + try: + k_scale, v_scale = pair + k_scale, v_scale = float(k_scale), float(v_scale) + except (TypeError, ValueError) as error: + raise ValueError(f"Layer {layer_id!r} must provide numeric (k_scale, v_scale) values.") from error + if not (math.isfinite(k_scale) and k_scale > 0 and math.isfinite(v_scale) and v_scale > 0): + raise ValueError(f"Layer {layer_id!r} FP8 scales must be finite and greater than zero.") + validated[int(layer_id)] = (k_scale, v_scale) + return validated + + +def _retype_fp8(value: ir.Value | None) -> None: + if value is None: + return + value.type = ir.TensorType(_FP8) + if value.const_value is not None and value.const_value.size == 0: + value.const_value = ir.tensor( + np.zeros(tuple(value.const_value.shape), dtype=_FP8.numpy()), + name=value.name, + ) + + +class ConvertGroupQueryAttentionKVCacheToFp8(Surgeon): + """Convert GroupQueryAttention past/present KV-cache tensors to FP8 E4M3.""" + + def __init__(self, scales: dict[int, tuple[float, float]] | None = None): + self.scales = _validate_fp8_scales(scales or {}) + + @staticmethod + def _scale_initializer(graph: ir.Graph, name: str, value: float) -> ir.Value: + existing = graph.initializers.get(name) + if existing is not None: + return existing + scale = ir.Value( + name=name, + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([1]), + const_value=ir.tensor(np.array([value], dtype=np.float32), name=name), + ) + graph.initializers[name] = scale + return scale + + def call_ir(self, model: ir.Model) -> ir.Model: + converted = 0 + for node in model.graph: + if node.domain != MSFT_DOMAIN or node.op_type != "GroupQueryAttention": + continue + past_key = node.inputs[3] if len(node.inputs) > 4 else None + past_value = node.inputs[4] if len(node.inputs) > 4 else None + if past_key is None or past_value is None: + continue + if any(value.const_value is not None and value.const_value.size > 0 for value in (past_key, past_value)): + warnings.warn( + f"Skipping {node.name!r}: only graph-input or empty-placeholder KV caches can become FP8.", + stacklevel=2, + ) + continue + + layer_match = _LAYER_ID_RE.search(past_key.name or "") + layer_id = int(layer_match.group(1)) if layer_match else -1 + k_value, v_value = self.scales.get(layer_id, (1.0, 1.0)) + + _retype_fp8(past_key) + _retype_fp8(past_value) + _retype_fp8(node.outputs[1] if len(node.outputs) > 1 else None) + _retype_fp8(node.outputs[2] if len(node.outputs) > 2 else None) + + k_scale = self._scale_initializer(model.graph, f"{past_key.name}.fp8_scale", k_value) + v_scale = self._scale_initializer(model.graph, f"{past_value.name}.fp8_scale", v_value) + if len(node.inputs) < _MIN_GQA_INPUTS_WITH_SCALES: + node.resize_inputs(_MIN_GQA_INPUTS_WITH_SCALES) + node.replace_input_with(_K_SCALE_INDEX, k_scale) + node.replace_input_with(_V_SCALE_INDEX, v_scale) + node.attributes.add(ir.AttrString("k_quant_type", "PER_TENSOR")) + node.attributes.add(ir.AttrString("v_quant_type", "PER_TENSOR")) + node.attributes.add(ir.AttrInt64("kv_cache_bit_width", 8)) + converted += 1 + + if converted == 0: + raise ValueError("No retypable GroupQueryAttention KV cache was found.") + return model + def _initializer_dtype(value: ir.Value) -> ir.DataType | None: """Return an initializer's declared dtype, falling back to its tensor dtype.""" @@ -133,9 +229,14 @@ def _static_last_dim(value): return False +def _is_supported_dtype(value: ir.Value, supported_dtypes: frozenset[ir.DataType] | None) -> bool: + return supported_dtypes is None or value.dtype in supported_dtypes + + class _RotaryAttentionToGQA(RewriteRuleClassBase): - def __init__(self): + def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): super().__init__() + self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None self._cos_cache = None @@ -155,8 +256,10 @@ def pattern(self, op, q_pre, k_pre, v, attention_bias, past_key, past_value, cos _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): + def check(self, context, q_pre, attn_out, k_pre, v, attention_bias, cos, sin, past_key, past_value, **_): result = MatchResult() + if not _is_supported_dtype(q_pre, self._supported_dtypes): + return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -258,8 +361,9 @@ def rewrite( class _AttentionToGQA(RewriteRuleClassBase): - def __init__(self): + def __init__(self, supported_dtypes: frozenset[ir.DataType] | None = None): super().__init__() + self._supported_dtypes = supported_dtypes self._seqlens_k = None self._total_seq_len = None @@ -275,8 +379,10 @@ def pattern(self, op, q, k, v, attention_bias, past_key, past_value): _outputs=["attn_out", "present_key", "present_value"], ) - def check(self, context, attn_out, k, v, attention_bias, past_key, past_value, **_): + def check(self, context, q, attn_out, k, v, attention_bias, past_key, past_value, **_): result = MatchResult() + if not _is_supported_dtype(q, self._supported_dtypes): + return result.fail("Attention dtype is unsupported by the target execution provider") if not _local_window_from_attention_bias(attention_bias).recognized: return result.fail("Attention bias cannot be represented by GroupQueryAttention") @@ -360,8 +466,21 @@ def rewrite( class AttentionToGroupQueryAttention(RewriteRuleSurgeon): """Fuse decoder Attention, and standard RoPE when possible, into GQA.""" + def __init__(self, supported_dtypes: list[str] | None = None): + try: + self.supported_dtypes = ( + frozenset(ir.DataType[dtype.upper()] for dtype in supported_dtypes) if supported_dtypes else None + ) + except KeyError as exc: + raise ValueError(f"Unsupported ONNX dtype {exc.args[0]!r}") from exc + def rules(self) -> pattern.RewriteRuleSet: - return pattern.RewriteRuleSet([_RotaryAttentionToGQA.rule(), _AttentionToGQA.rule()]) + return pattern.RewriteRuleSet( + [ + _RotaryAttentionToGQA.rule(self.supported_dtypes), + _AttentionToGQA.rule(self.supported_dtypes), + ] + ) class _PackQKVForGQA(RewriteRuleClassBase): diff --git a/olive/passes/onnx/graph_surgery/lowering.py b/olive/passes/onnx/graph_surgery/lowering.py index 35addbbeb..65652fa14 100644 --- a/olive/passes/onnx/graph_surgery/lowering.py +++ b/olive/passes/onnx/graph_surgery/lowering.py @@ -9,6 +9,7 @@ import numpy as np import onnx_ir as ir from onnx_ir import tape +from onnx_ir.passes.common import InlinePass, RemoveUnusedOpsetsPass from onnxscript.rewriter import pattern from olive.passes.onnx.graph_surgery.base import RewriteRuleSurgeon, Surgeon @@ -17,6 +18,29 @@ # pylint: disable=arguments-differ _MASK_NEGATIVE_INFINITY = float("-inf") +_STANDARD_ONNX_DOMAINS = frozenset({"", "ai.onnx"}) + + +class InlineModelLocalFunctions(Surgeon): + """Inline model-local functions and require the resulting graph to use standard ONNX domains.""" + + def call_ir(self, model: ir.Model) -> ir.Model: + InlinePass()(model) + model.functions.clear() + RemoveUnusedOpsetsPass()(model) + non_standard_ops = sorted( + { + f"{node.domain}::{node.op_type}" + for node in model.graph.all_nodes() + if node.domain not in _STANDARD_ONNX_DOMAINS + } + ) + if non_standard_ops: + raise ValueError( + "Cannot produce strict ONNX because these operators have no model-local " + f"standard function body: {non_standard_ops}." + ) + return model class _BFloat16ClipRule(pattern.RewriteRuleClassBase): diff --git a/olive/passes/onnx/mobius_model_builder.py b/olive/passes/onnx/mobius_model_builder.py index d52b6d26f..87f595df0 100644 --- a/olive/passes/onnx/mobius_model_builder.py +++ b/olive/passes/onnx/mobius_model_builder.py @@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Any, ClassVar from olive.common.utils import StrEnumBase -from olive.constants import Precision from olive.hardware.constants import EXECUTION_PROVIDER_TO_MOBIUS_EP, ExecutionProvider from olive.model import HfModelHandler, ONNXModelHandler from olive.model.handler.composite import CompositeModelHandler @@ -24,18 +23,6 @@ logger = logging.getLogger(__name__) -# Maps Olive Precision values to mobius dtype strings. -# "f32" = 32-bit float (torch.float32), standard full precision. -# "f16" = 16-bit float (torch.float16), half precision — good for GPU inference. -# "bf16" = bfloat16 (torch.bfloat16), brain float — preferred over f16 on newer hardware. -# For INT4/INT8 quantization, use a downstream Olive quantization pass (e.g. OnnxMatMulNBits) -# after this pass rather than setting precision here. -_PRECISION_TO_DTYPE: dict[str, str] = { - Precision.FP32: "f32", - Precision.FP16: "f16", - Precision.BF16: "bf16", -} - class MobiusBuilder(Pass): """Olive pass that uses mobius to build ONNX models from HuggingFace model IDs. @@ -46,6 +33,12 @@ class MobiusBuilder(Pass): pass returns a :class:`~olive.model.handler.composite.CompositeModelHandler` whose components are individual :class:`~olive.model.ONNXModelHandler` objects. Single-component models return a plain :class:`~olive.model.ONNXModelHandler`. + Mobius preserves the source model precision and exports its canonical graph + without automatic rewrites; strict-ONNX expansion and target-specific graph + transformations run in downstream Olive passes. + The target accelerator is forwarded only as a structural build contract, so + requirements such as OpenVINO's rank-4 Gemma4 component interface are preserved + without enabling Mobius graph rewrites. Use ``components_to_export`` to export only a subset of components. This is useful when some components (e.g. a text decoder) are already exported and @@ -89,33 +82,25 @@ class MobiusEP(StrEnumBase): TRT_RTX = "trt-rtx" ONNX_STANDARD = "onnx-standard" - # Maps Olive ExecutionProvider enum values to mobius EP names. EP_MAP: ClassVar[dict[ExecutionProvider, str]] = { ExecutionProvider.CPUExecutionProvider: "cpu", ExecutionProvider.CUDAExecutionProvider: "cuda", ExecutionProvider.DmlExecutionProvider: "dml", + ExecutionProvider.NvTensorRTRTXExecutionProvider: "trt-rtx", + ExecutionProvider.OpenVINOExecutionProvider: "openvino", + ExecutionProvider.QNNExecutionProvider: "qnn", ExecutionProvider.WebGpuExecutionProvider: "webgpu", } @classmethod def is_accelerator_agnostic(cls, accelerator_spec: AcceleratorSpec) -> bool: - # EP selection determines which fused ops are emitted, so this pass is - # EP-specific. + # The graph is accelerator-agnostic, but ORT GenAI runtime packaging + # records the target EP. return False @classmethod def _default_config(cls, accelerator_spec: AcceleratorSpec) -> dict[str, PassConfigParam]: return { - "precision": PassConfigParam( - type_=Precision, - required=False, - default_value=Precision.FP32, - description=( - "Model weight / compute precision. One of: fp32, fp16, bf16. " - "Defaults to fp32. For INT4 quantization, run an Olive " - "quantization pass (e.g. OnnxMatMulNBits) after this pass." - ), - ), "text_only": PassConfigParam( type_=bool, required=False, @@ -159,18 +144,21 @@ def _run_for_config( if not isinstance(model, HfModelHandler): raise ValueError(f"MobiusBuilder requires an HfModelHandler input, got {type(model).__name__}.") - # Map Olive EP to mobius EP. If unsupported/unknown, fall back to mobius default EP. + # The requested EP/device provide structural build requirements and + # ORT GenAI runtime packaging without enabling Mobius graph rewrites. requested_ep = self.accelerator_spec.execution_provider - ep_str: str = EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT) - if ep_str == self.MobiusEP.DEFAULT: + runtime_ep: str = self.EP_MAP.get( + requested_ep, + EXECUTION_PROVIDER_TO_MOBIUS_EP.get(requested_ep, self.MobiusEP.DEFAULT), + ) + if runtime_ep == self.MobiusEP.DEFAULT: logger.warning( "MobiusBuilder: execution provider '%s' on accelerator '%s' is not explicitly supported; " - "falling back to mobius default EP.", + "using the default Mobius runtime configuration.", requested_ep, self.accelerator_spec.accelerator_type, ) - dtype_str: str = _PRECISION_TO_DTYPE.get(config.precision, "f32") model_id: str = model.model_name_or_path load_kwargs = model.get_load_kwargs() @@ -178,10 +166,9 @@ def _run_for_config( trust_remote_code: bool = load_kwargs.get("trust_remote_code", False) logger.info( - "MobiusBuilder: building '%s' (ep=%s, dtype=%s)", + "MobiusBuilder: building '%s' as standard ONNX (runtime ep=%s)", model_id, - ep_str, - dtype_str, + runtime_ep, ) if trust_remote_code: @@ -201,8 +188,8 @@ def _run_for_config( pkg = build( model_id, revision=revision, - dtype=dtype_str, - execution_provider=ep_str, + execution_provider=runtime_ep, + device=str(self.accelerator_spec.accelerator_type), load_weights=True, trust_remote_code=trust_remote_code, **text_only_kwargs, @@ -258,7 +245,7 @@ def components_filter(name: str) -> bool: pkg, str(output_dir), model_id, - ep_str, + runtime_ep, revision=revision, trust_remote_code=trust_remote_code, ) diff --git a/test/cli/test_cli.py b/test/cli/test_cli.py index f12ab8359..d29b57d74 100644 --- a/test/cli/test_cli.py +++ b/test/cli/test_cli.py @@ -12,6 +12,7 @@ from olive.cli.base import TEST_OUTPUT_MARKER_FILE from olive.cli.launcher import main as cli_main +from olive.workflows.run.builds import parse_run_config @pytest.mark.parametrize("console_script", [True, False]) @@ -539,15 +540,8 @@ def test_capture_onnx_command_fix_shape(_, mock_run, use_model_builder, tmp_path @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) -@pytest.mark.parametrize( - ("precision", "use_ort_genai"), - [ - ("fp16", True), - ("fp32", False), - ("bf16", True), - ], -) -def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort_genai, tmp_path): +@pytest.mark.parametrize("use_ort_genai", [True, False]) +def test_capture_onnx_command_use_mobius_builder(_, mock_run, use_ort_genai, tmp_path): # setup output_dir = tmp_path / "output_dir" model_id = "dummy-model-id" @@ -558,8 +552,6 @@ def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort "-o", str(output_dir), "--use_mobius_builder", - "--precision", - precision, ] if use_ort_genai: command_args.append("--use_ort_genai") @@ -573,14 +565,18 @@ def test_capture_onnx_command_use_mobius_builder(_, mock_run, precision, use_ort assert "b" in config["passes"] assert "c" not in config["passes"] assert "m" not in config["passes"] + assert "g" not in config["passes"] assert config["passes"]["b"]["type"] == "MobiusBuilder" - assert config["passes"]["b"]["precision"] == precision + assert "precision" not in config["passes"]["b"] + accelerator = config["systems"]["local_system"]["accelerators"][0] + assert accelerator["device"] == "cpu" + assert accelerator["execution_providers"] == ["CPUExecutionProvider"] assert mock_run.call_count == 1 @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) -def test_capture_onnx_command_use_mobius_builder_rejects_int4(_, __, tmp_path): +def test_capture_onnx_command_use_mobius_builder_ignores_model_builder_precision(_, mock_run, tmp_path): # setup output_dir = tmp_path / "output_dir" command_args = [ @@ -594,11 +590,346 @@ def test_capture_onnx_command_use_mobius_builder_rejects_int4(_, __, tmp_path): "int4", ] - # execute / verify - with pytest.raises(ValueError, match="MobiusBuilder supports precisions fp32/fp16/bf16"): + cli_main(command_args) + + config = mock_run.call_args[0][0] + assert "precision" not in config["passes"]["b"] + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +@pytest.mark.parametrize( + ("ep", "device", "provider", "surgeons"), + [ + ( + "cuda", + "gpu", + "CUDAExecutionProvider", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + ], + ), + ( + "trt-rtx", + "gpu", + "NvTensorRTRTXExecutionProvider", + ["AttentionToGroupQueryAttention", "PackQKVForGroupQueryAttention"], + ), + ( + "qnn", + "gpu", + "QNNExecutionProvider", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "AttentionMaskToSequenceLengths", + ], + ), + ( + "qnn", + "npu", + "QNNExecutionProvider", + [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "AttentionMaskToSequenceLengths", + "SimplifiedLayerNormToL2Norm", + ], + ), + ], +) +def test_capture_onnx_command_adds_recipe_ep_surgeries(_, mock_run, ep, device, provider, surgeons, tmp_path): + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "ONNXModel", + "config": {"model_path": str(tmp_path / "output"), "onnx_file_name": "model.onnx"}, + } + mock_run.side_effect = [exported_output, MagicMock()] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + "--execution_provider", + ep, + "--device", + device, + ] + ) + + export_config, surgery_config = [call.args[0] for call in mock_run.call_args_list] + assert list(export_config["passes"]) == ["b"] + assert "builds" not in export_config + accelerator = surgery_config["systems"]["local_system"]["accelerators"][0] + assert accelerator == {"device": device, "execution_providers": [provider]} + assert list(surgery_config["passes"]) == ["g"] + assert "builds" not in surgery_config + assert [surgery["surgeon"] for surgery in surgery_config["passes"]["g"]["surgeries"]] == surgeons + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_limits_recipe_ep_surgeries_to_multimodal_decoder(_, mock_run, tmp_path): + output_dir = tmp_path / "output" + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "CompositeModel", + "config": { + "model_path": str(output_dir), + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "decoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "vision_encoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "embedding")}}, + ], + "model_component_names": ["decoder", "vision_encoder", "embedding"], + }, + } + mock_run.side_effect = [exported_output, {"decoder": MagicMock()}] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(output_dir), + "--use_mobius_builder", + "--execution_provider", + "cuda", + "--device", + "gpu", + ] + ) + + export_config, surgery_config = [call.args[0] for call in mock_run.call_args_list] + assert list(export_config["passes"]) == ["b"] + assert "builds" not in export_config + assert surgery_config["builds"] == { + "decoder": {"components": ["decoder"], "pipeline": ["g"]}, + } + assert "b" not in surgery_config["passes"] + assert [surgery["surgeon"] for surgery in surgery_config["passes"]["g"]["surgeries"]] == [ + "AttentionToGroupQueryAttention", + "PackQKVForGroupQueryAttention", + "FuseSkipRMSNormalization", + "FuseSkipLayerNormalization", + ] + parsed = parse_run_config(surgery_config) + assert list(parsed) == ["decoder"] + assert list(parsed["decoder"].passes) == ["g"] + assert parsed["decoder"].engine.output_dir == (output_dir / "decoder").resolve() + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_onnx_standard_inlines_all_multimodal_components(_, mock_run, tmp_path): + output_dir = tmp_path / "output" + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "CompositeModel", + "config": { + "model_path": str(output_dir), + "model_components": [ + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "decoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "vision_encoder")}}, + {"type": "ONNXModel", "config": {"model_path": str(output_dir / "embedding")}}, + ], + "model_component_names": ["decoder", "vision_encoder", "embedding"], + }, + } + mock_run.side_effect = [exported_output, MagicMock()] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(output_dir), + "--use_mobius_builder", + "--execution_provider", + "openvino", + "--device", + "npu", + "--onnx_standard", + ] + ) + + _, standard_config = [call.args[0] for call in mock_run.call_args_list] + assert "builds" not in standard_config + assert standard_config["passes"]["g"]["surgeries"] == [ + {"surgeon": "InlineModelLocalFunctions"}, + ] + assert standard_config["systems"]["local_system"]["accelerators"][0] == { + "device": "npu", + "execution_providers": ["OpenVINOExecutionProvider"], + } + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_onnx_standard_without_target_ep(_, mock_run, tmp_path): + output_dir = tmp_path / "output" + exported_output = MagicMock() + exported_output.has_output_model.return_value = True + exported_output.get_best_candidate.return_value.olive_model_config = { + "type": "ONNXModel", + "config": {"model_path": str(output_dir), "onnx_file_name": "model.onnx"}, + } + mock_run.side_effect = [exported_output, MagicMock()] + + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(output_dir), + "--use_mobius_builder", + "--onnx_standard", + ] + ) + + _, standard_config = [call.args[0] for call in mock_run.call_args_list] + assert standard_config["passes"]["g"]["surgeries"] == [ + {"surgeon": "InlineModelLocalFunctions"}, + ] + assert standard_config["systems"]["local_system"]["accelerators"][0] == { + "device": "cpu", + "execution_providers": ["CPUExecutionProvider"], + } + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +@pytest.mark.parametrize( + ("ep", "device", "provider"), + [ + ("openvino", "cpu", "OpenVINOExecutionProvider"), + ("openvino", "gpu", "OpenVINOExecutionProvider"), + ("openvino", "npu", "OpenVINOExecutionProvider"), + ("vitisai", "npu", "VitisAIExecutionProvider"), + ], +) +def test_capture_onnx_command_recipe_ep_without_surgeries(_, mock_run, ep, device, provider, tmp_path): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + "--execution_provider", + ep, + "--device", + device, + ] + ) + + config = mock_run.call_args[0][0] + assert list(config["passes"]) == ["b"] + assert config["systems"]["local_system"]["accelerators"][0] == { + "device": device, + "execution_providers": [provider], + } + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +@pytest.mark.parametrize( + ("extra_args", "message"), + [ + (["--execution_provider", "cuda"], "must be provided together"), + (["--device", "gpu"], "must be provided together"), + (["--execution_provider", "cuda", "--device", "npu"], "does not support device"), + (["--execution_provider", "trt-rtx", "--device", "npu"], "does not support device"), + ], +) +def test_capture_onnx_command_rejects_invalid_recipe_ep_device(_, __, extra_args, message, tmp_path): + command_args = [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + *extra_args, + ] + + with pytest.raises(ValueError, match=message): cli_main(command_args) +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_rejects_execution_provider_without_recipe_profile(_, __, tmp_path): + with pytest.raises(SystemExit) as exc_info: + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--use_mobius_builder", + "--execution_provider", + "dml", + "--device", + "gpu", + ] + ) + + assert exc_info.value.code == 2 + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_rejects_ep_device_without_mobius(_, __, tmp_path): + with pytest.raises(ValueError, match="require --use_mobius_builder"): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--execution_provider", + "cuda", + "--device", + "gpu", + ] + ) + + +@patch("olive.workflows.run") +@patch("huggingface_hub.repo_exists", return_value=True) +def test_capture_onnx_command_rejects_onnx_standard_without_mobius(_, __, tmp_path): + with pytest.raises(ValueError, match="requires --use_mobius_builder"): + cli_main( + [ + "capture-onnx-graph", + "-m", + "dummy-model-id", + "-o", + str(tmp_path / "output"), + "--onnx_standard", + ] + ) + + @patch("olive.workflows.run") @patch("huggingface_hub.repo_exists", return_value=True) @pytest.mark.parametrize("conflicting_flag", ["--use_model_builder", "--use_dynamo_exporter"]) diff --git a/test/passes/onnx/test_graph_surgeries_attention.py b/test/passes/onnx/test_graph_surgeries_attention.py index 862c49e98..466620fe2 100644 --- a/test/passes/onnx/test_graph_surgeries_attention.py +++ b/test/passes/onnx/test_graph_surgeries_attention.py @@ -85,7 +85,7 @@ def _run_surgeries(tmp_path, model, *surgeons): graph_surgeries = create_pass_from_dict( GraphSurgeries, { - "surgeries": [{"surgeon": surgeon} for surgeon in surgeons], + "surgeries": [{"surgeon": surgeon} if isinstance(surgeon, str) else surgeon for surgeon in surgeons], "remove_duplicate_initializers": False, }, disable_search=True, @@ -131,6 +131,7 @@ def _make_attention_model( with_cache=True, k_dim=32, v_dim=32, + dtype=ir.DataType.FLOAT16, ): nodes = [] inputs = [_value("attention_mask", ir.DataType.INT64, [1, "total_sequence"])] @@ -142,19 +143,19 @@ def _make_attention_model( ) inputs.extend(mask_inputs) - cos_cache = _value("cos_cache", ir.DataType.FLOAT16, [128, 8]) - sin_cache = _value("sin_cache", ir.DataType.FLOAT16, [128, 8]) + cos_cache = _value("cos_cache", dtype, [128, 8]) + sin_cache = _value("sin_cache", dtype, [128, 8]) position_ids = _value("position_ids", ir.DataType.INT64, [1, 2]) if rotary: inputs.extend([cos_cache, sin_cache, position_ids]) graph_outputs = [] for layer in range(num_layers): - q = _value(f"q_{layer}", ir.DataType.FLOAT16, [1, 2, 64]) - k = _value(f"k_{layer}", ir.DataType.FLOAT16, [1, 2, k_dim]) - v = _value(f"v_{layer}", ir.DataType.FLOAT16, [1, 2, v_dim]) - past_key = _value(f"past_key_{layer}", ir.DataType.FLOAT16, [1, 2, 4, k_dim]) - past_value = _value(f"past_value_{layer}", ir.DataType.FLOAT16, [1, 2, 4, v_dim]) + q = _value(f"q_{layer}", dtype, [1, 2, 64]) + k = _value(f"k_{layer}", dtype, [1, 2, k_dim]) + v = _value(f"v_{layer}", dtype, [1, 2, v_dim]) + past_key = _value(f"past_key_{layer}", dtype, [1, 2, 4, k_dim]) + past_value = _value(f"past_value_{layer}", dtype, [1, 2, 4, v_dim]) inputs.extend([q, k, v]) if with_cache: inputs.extend([past_key, past_value]) @@ -196,7 +197,7 @@ def _make_attention_model( attention.outputs, ([1, 2, 64], [1, 2, 6, k_dim], [1, 2, 6, v_dim]), ): - output.dtype = ir.DataType.FLOAT16 + output.dtype = dtype output.shape = ir.Shape(shape) graph_outputs.extend(attention.outputs) @@ -422,6 +423,7 @@ def _make_block_diagonal_attention_model(*, false_bias=-10000.0, shared_segments def test_attention_surgeries_register_on_module_import(): expected = { "attentiontogroupqueryattention", + "convertgroupqueryattentionkvcachetofp8", "separategroupqueryattentionrope", "unpackgroupqueryattentionqkv", "blockdiagonalattentiontopackedmha", @@ -430,6 +432,53 @@ def test_attention_surgeries_register_on_module_import(): assert expected <= Surgeon.registry.keys() +def test_convert_gqa_kv_cache_to_fp8_adds_scales_and_retypes_io(tmp_path): + query = _value("query", ir.DataType.FLOAT16, [1, 1, 32]) + key = _value("key", ir.DataType.FLOAT16, [1, 1, 16]) + value = _value("value", ir.DataType.FLOAT16, [1, 1, 16]) + past_key = _value("past_key_values.0.key", ir.DataType.FLOAT16, [1, 1, 0, 16]) + past_value = _value("past_key_values.0.value", ir.DataType.FLOAT16, [1, 1, 0, 16]) + seqlens = _value("seqlens", ir.DataType.INT32, [1]) + total_sequence = _value("total_sequence", ir.DataType.INT32, []) + gqa = _node( + "GroupQueryAttention", + [query, key, value, past_key, past_value, seqlens, total_sequence], + domain=_MS_DOMAIN, + num_outputs=3, + output_names=["output", "present.0.key", "present.0.value"], + ) + model = _model( + [query, key, value, past_key, past_value, seqlens, total_sequence], + list(gqa.outputs), + [gqa], + ) + + rewritten = _run_surgeries( + tmp_path, + model, + { + "surgeon": "ConvertGroupQueryAttentionKVCacheToFp8", + "scales": {0: [0.25, 0.5]}, + }, + ) + rewritten_gqa = next(node for node in rewritten.graph if node.op_type == "GroupQueryAttention") + + assert rewritten_gqa.inputs[3].dtype == ir.DataType.FLOAT8E4M3FN + assert rewritten_gqa.inputs[4].dtype == ir.DataType.FLOAT8E4M3FN + assert rewritten_gqa.outputs[1].dtype == ir.DataType.FLOAT8E4M3FN + assert rewritten_gqa.outputs[2].dtype == ir.DataType.FLOAT8E4M3FN + assert len(rewritten_gqa.inputs) == 14 + np.testing.assert_array_equal( + rewritten_gqa.inputs[12].const_value.numpy(), + np.array([0.25], dtype=np.float32), + ) + np.testing.assert_array_equal( + rewritten_gqa.inputs[13].const_value.numpy(), + np.array([0.5], dtype=np.float32), + ) + assert rewritten_gqa.attributes.get_int("kv_cache_bit_width") == 8 + + def test_attention_to_gqa_fuses_rotary_preserves_attributes_outputs_and_shared_inputs(tmp_path): model = _make_attention_model(rotary=True, num_layers=2, sliding_window=64, k_dim=32, v_dim=32) rewritten = _run_surgeries(tmp_path, model, "AttentionToGroupQueryAttention") @@ -470,6 +519,26 @@ def test_attention_to_gqa_fallback_uses_external_rope_and_preserves_three_output assert len(rewritten.graph.outputs) == 3 +@pytest.mark.parametrize( + ("dtype", "expected_gqa"), + [(ir.DataType.FLOAT16, 1), (ir.DataType.FLOAT, 0)], +) +@pytest.mark.parametrize("rotary", [False, True]) +def test_attention_to_gqa_filters_target_unsupported_dtypes(tmp_path, dtype, expected_gqa, rotary): + model = _make_attention_model(dtype=dtype, rotary=rotary) + rewritten = _run_surgeries( + tmp_path, + model, + { + "surgeon": "AttentionToGroupQueryAttention", + "supported_dtypes": ["FLOAT16", "BFLOAT16"], + }, + ) + + assert _count_ops(rewritten)["GroupQueryAttention"] == expected_gqa + assert _count_ops(rewritten)["Attention"] == 1 - expected_gqa + + @pytest.mark.parametrize( "model", [ diff --git a/test/passes/onnx/test_graph_surgeries_lowering.py b/test/passes/onnx/test_graph_surgeries_lowering.py index 4c3d3a84b..8a8dd0cd7 100644 --- a/test/passes/onnx/test_graph_surgeries_lowering.py +++ b/test/passes/onnx/test_graph_surgeries_lowering.py @@ -11,6 +11,7 @@ import onnx_ir as ir import onnxruntime as ort import pytest +from onnx import TensorProto, helper from onnx.reference import ReferenceEvaluator from olive.model import ONNXModelHandler @@ -59,6 +60,7 @@ def _metadata(value: ir.Value): def test_lowering_module_registers_all_surgeons(): expected = { lowering_surgeries.ClipToMinMax, + lowering_surgeries.InlineModelLocalFunctions, lowering_surgeries.Rank4RMSNormToRank3, lowering_surgeries.DecomposeOnnxRotaryEmbedding, lowering_surgeries.TensorScatterToScatterND, @@ -68,6 +70,53 @@ def test_lowering_module_registers_all_surgeons(): assert {Surgeon.registry[surgeon.__name__.lower()] for surgeon in expected} == expected +def _model_with_local_identity_function() -> ir.Model: + function = helper.make_function( + "com.microsoft", + "CustomIdentity", + ["x"], + ["y"], + [helper.make_node("Identity", ["x"], ["y"])], + [helper.make_opsetid("", 24)], + ) + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1]) + graph = helper.make_graph( + [helper.make_node("CustomIdentity", ["x"], ["y"], domain="com.microsoft")], + "custom_identity", + [x], + [y], + ) + model = helper.make_model( + graph, + functions=[function], + opset_imports=[ + helper.make_opsetid("", 24), + helper.make_opsetid("com.microsoft", 1), + ], + ) + return ir.from_proto(model) + + +def test_inline_model_local_functions_produces_standard_onnx(tmp_path): + output = _apply_surgery( + tmp_path, + _model_with_local_identity_function(), + "InlineModelLocalFunctions", + ) + + assert _counts(output) == Counter({"Identity": 1}) + assert not output.functions + + +def test_inline_model_local_functions_rejects_custom_op_without_function(tmp_path): + model = _model_with_local_identity_function() + model.functions.clear() + + with pytest.raises(ValueError, match="no model-local standard function body"): + _apply_surgery(tmp_path, model, "InlineModelLocalFunctions") + + def _clip_model( dtype: ir.DataType, *, diff --git a/test/passes/onnx/test_mobius_model_builder.py b/test/passes/onnx/test_mobius_model_builder.py index fe43864ca..0ffee8199 100644 --- a/test/passes/onnx/test_mobius_model_builder.py +++ b/test/passes/onnx/test_mobius_model_builder.py @@ -79,7 +79,7 @@ def _make_hf_model(model_path: str, load_kwargs: dict | None = None, task: str | def _make_pass(ep: str = ExecutionProvider.CPUExecutionProvider, text_only: bool | None = None) -> MobiusBuilder: accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=ep) - pass_config = {"precision": "fp32"} + pass_config = {} if text_only is not None: pass_config["text_only"] = text_only return create_pass_from_dict( @@ -198,12 +198,12 @@ def __exit__(self, *args): def test_default_config_params(): - """MobiusBuilder must declare precision, and must not declare execution_provider or trust_remote_code.""" + """MobiusBuilder preserves model precision and does not expose build environment options.""" accelerator_spec = AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider ) config = MobiusBuilder._default_config(accelerator_spec) # pylint: disable=protected-access - assert "precision" in config + assert "precision" not in config assert config["text_only"].default_value is False assert config["text_only"].required is False assert "execution_provider" not in config @@ -211,7 +211,7 @@ def test_default_config_params(): def test_is_not_accelerator_agnostic(): - """Pass must be EP-specific because it chooses fused ops based on the EP.""" + """Pass remains EP-specific because runtime packaging records the target EP.""" accelerator_spec = AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider ) @@ -222,10 +222,16 @@ def test_ep_map_covers_common_providers(): assert ExecutionProvider.CPUExecutionProvider in MobiusBuilder.EP_MAP assert ExecutionProvider.CUDAExecutionProvider in MobiusBuilder.EP_MAP assert ExecutionProvider.DmlExecutionProvider in MobiusBuilder.EP_MAP + assert ExecutionProvider.NvTensorRTRTXExecutionProvider in MobiusBuilder.EP_MAP + assert ExecutionProvider.OpenVINOExecutionProvider in MobiusBuilder.EP_MAP + assert ExecutionProvider.QNNExecutionProvider in MobiusBuilder.EP_MAP assert ExecutionProvider.WebGpuExecutionProvider in MobiusBuilder.EP_MAP assert MobiusBuilder.EP_MAP[ExecutionProvider.CPUExecutionProvider] == "cpu" assert MobiusBuilder.EP_MAP[ExecutionProvider.CUDAExecutionProvider] == "cuda" assert MobiusBuilder.EP_MAP[ExecutionProvider.DmlExecutionProvider] == "dml" + assert MobiusBuilder.EP_MAP[ExecutionProvider.NvTensorRTRTXExecutionProvider] == "trt-rtx" + assert MobiusBuilder.EP_MAP[ExecutionProvider.OpenVINOExecutionProvider] == "openvino" + assert MobiusBuilder.EP_MAP[ExecutionProvider.QNNExecutionProvider] == "qnn" assert MobiusBuilder.EP_MAP[ExecutionProvider.WebGpuExecutionProvider] == "webgpu" @@ -257,7 +263,8 @@ def test_single_component_returns_onnx_handler(tmp_path): mock_build.assert_called_once() call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == "cpu" - assert call_kwargs["dtype"] == "f32" + assert call_kwargs["device"] == "cpu" + assert "dtype" not in call_kwargs def test_text_only_default_omits_mobius_build_kwarg(tmp_path): @@ -406,8 +413,8 @@ def test_multi_component_returns_composite_handler(tmp_path): # --------------------------------------------------------------------------- -def test_ep_auto_detected_from_accelerator(tmp_path): - """Execution provider is determined by the Olive accelerator spec.""" +def test_target_accelerator_is_forwarded_without_enabling_mobius_rewrites(tmp_path): + """The target accelerator controls the build contract and runtime packaging.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -416,17 +423,52 @@ def test_ep_auto_detected_from_accelerator(tmp_path): ) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp16"}, + {}, disable_search=True, accelerator_spec=accelerator_spec, ) - with _patch_build(pkg) as mock_build: + with ( + patch("mobius.build", return_value=pkg) as mock_build, + patch.object(MobiusBuilder, "_write_genai_config", return_value={}) as mock_write, + ): p.run(_make_hf_model("org/model"), out) call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == "cuda" - assert call_kwargs["dtype"] == "f16" + assert call_kwargs["device"] == "gpu" + assert "dtype" not in call_kwargs + mock_write.assert_called_once_with( + pkg, + str(out), + "org/model", + "cuda", + revision=None, + trust_remote_code=False, + ) + + +def test_openvino_target_contract_keeps_standard_export(tmp_path): + """OpenVINO structural requirements are independent of graph rewrites.""" + out = tmp_path / "out" + pkg = _fake_pkg(["model"], out) + accelerator_spec = AcceleratorSpec( + accelerator_type=Device.NPU, + execution_provider=ExecutionProvider.OpenVINOExecutionProvider, + ) + p = create_pass_from_dict( + MobiusBuilder, + {}, + disable_search=True, + accelerator_spec=accelerator_spec, + ) + + with _patch_build(pkg) as mock_build: + p.run(_make_hf_model("org/model"), out) + + call_kwargs = mock_build.call_args.kwargs + assert call_kwargs["execution_provider"] == "openvino" + assert call_kwargs["device"] == "npu" def test_hf_load_options_forwarded_to_build_and_genai_config(tmp_path): @@ -509,8 +551,8 @@ def test_write_genai_config_defaults_hf_load_options(tmp_path): ) -def test_unsupported_ep_falls_back_to_default(tmp_path): - """If accelerator EP is unsupported, pass should fall back to mobius default EP.""" +def test_unsupported_ep_uses_standard_export(tmp_path): + """An unsupported runtime EP does not change the standard ONNX export path.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -521,7 +563,7 @@ def test_unsupported_ep_falls_back_to_default(tmp_path): ) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp32"}, + {}, disable_search=True, accelerator_spec=accelerator_spec, ) @@ -531,10 +573,11 @@ def test_unsupported_ep_falls_back_to_default(tmp_path): call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["device"] == "npu" -def test_none_execution_provider_falls_back_to_default(tmp_path): - """If execution_provider is None, pass should fall back to mobius default EP.""" +def test_none_execution_provider_uses_standard_export(tmp_path): + """An unspecified runtime EP does not change the standard ONNX export path.""" out = tmp_path / "out" pkg = _fake_pkg(["model"], out) @@ -542,7 +585,7 @@ def test_none_execution_provider_falls_back_to_default(tmp_path): accelerator_spec = AcceleratorSpec(accelerator_type=Device.CPU, execution_provider=None) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp32"}, + {}, disable_search=True, accelerator_spec=accelerator_spec, ) @@ -552,6 +595,7 @@ def test_none_execution_provider_falls_back_to_default(tmp_path): call_kwargs = mock_build.call_args.kwargs assert call_kwargs["execution_provider"] == MobiusBuilder.MobiusEP.DEFAULT + assert call_kwargs["device"] == "cpu" @pytest.mark.skipif(not _HAS_REAL_MOBIUS, reason="mobius-onnx is not publicly available in CI yet") @@ -637,7 +681,7 @@ def test_trust_remote_code_warning_logged(tmp_path): pkg = _fake_pkg(["model"], out) p = create_pass_from_dict( MobiusBuilder, - {"precision": "fp32"}, + {}, disable_search=True, accelerator_spec=AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider @@ -672,13 +716,13 @@ def test_no_warning_when_trust_remote_code_false(tmp_path): # --------------------------------------------------------------------------- -def _make_filtered_pass(components_to_export, precision: str = "fp16") -> MobiusBuilder: +def _make_filtered_pass(components_to_export) -> MobiusBuilder: accelerator_spec = AcceleratorSpec( accelerator_type=Device.CPU, execution_provider=ExecutionProvider.CPUExecutionProvider ) return create_pass_from_dict( MobiusBuilder, - {"precision": precision, "components_to_export": components_to_export}, + {"components_to_export": components_to_export}, disable_search=True, accelerator_spec=accelerator_spec, )