diff --git a/optimum/exporters/openvino/__init__.py b/optimum/exporters/openvino/__init__.py index 94ea4f103b..320a15e38a 100644 --- a/optimum/exporters/openvino/__init__.py +++ b/optimum/exporters/openvino/__init__.py @@ -15,7 +15,7 @@ import optimum.exporters.openvino.model_configs from .__main__ import main_export -from .convert import export, export_from_model, export_models, export_pytorch_via_onnx +from .convert import export, export_from_model, export_models from .stateful import ensure_stateful_is_available, patch_stateful diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index 1fb3f5d43c..f3d8fed9de 100644 --- a/optimum/exporters/openvino/__main__.py +++ b/optimum/exporters/openvino/__main__.py @@ -28,7 +28,7 @@ from transformers.utils import is_torch_available from openvino import Core, Type, save_model -from optimum.exporters.onnx.base import OnnxConfig +from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import ( DIFFUSERS_IMPORT_ERROR, @@ -227,7 +227,7 @@ def main_export( local_files_only: bool = False, token: Optional[Union[bool, str]] = None, model_kwargs: Optional[Dict[str, Any]] = None, - custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None, + custom_export_configs: Optional[Dict[str, "OpenVINOConfig"]] = None, fn_get_submodels: Optional[Callable] = None, ov_config: "OVConfig" = None, stateful: bool = True, @@ -283,7 +283,7 @@ def main_export( the export. This argument should be used along the `custom_export_configs` argument in case, for example, the model inputs/outputs are changed (for example, if `model_kwargs={"output_attentions": True}` is passed). - custom_export_configs (`Optional[Dict[str, OnnxConfig]]`, defaults to `None`): + custom_export_configs (`Optional[Dict[str, OpenVINOConfig]]`, defaults to `None`): Experimental usage: override the default export config used for the given model. This argument may be useful for advanced users that desire a finer-grained control on the export. An example is available [here](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model). fn_get_submodels (`Optional[Callable]`, defaults to `None`): Experimental usage: Override the default submodels that are used at the export. This is diff --git a/optimum/exporters/openvino/_traceable_cache.py b/optimum/exporters/openvino/_traceable_cache.py new file mode 100644 index 0000000000..cc2486713e --- /dev/null +++ b/optimum/exporters/openvino/_traceable_cache.py @@ -0,0 +1,107 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import logging +from typing import Any + +import torch + + +logger = logging.getLogger(__name__) + + +# Simply removing the nn.Module, same as in https://github.com/huggingface/transformers/pull/35873 +class TraceableCache: + """Base, abstract class for all caches. The actual data structure is specific to each subclass.""" + + def __init__(self): + super().__init__() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. These are specific to each subclass and allow new types of + cache to be created. + + Return: + A tuple containing the updated key and value states. + """ + raise NotImplementedError("Make sure to implement `update` in a subclass.") + + def get_seq_length(self, layer_idx: int | None = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + raise NotImplementedError("Make sure to implement `get_seq_length` in a subclass.") + + # Deprecate in favor of max-cache-shape because we want to be specific by what we mean with "max_length" + # Prev some cache objects didn't have "max_length" (SlidingWindowCache or SinkCache) because the cache object technically handles + # infinite amount of tokens. In the codebase what we really need to check is the max capacity of certain cache instances, so + # we change naming to be more explicit + def get_max_length(self) -> int | None: + logger.warning_once( + "`get_max_cache()` is deprecated for all Cache classes. Use `get_max_cache_shape()` instead. " + "Calling `get_max_cache()` will raise error from v4.48" + ) + return self.get_max_cache_shape() + + def get_max_cache_shape(self) -> int | None: + """Returns the maximum sequence length (i.e. max capacity) of the cache object.""" + raise NotImplementedError("Make sure to implement `get_max_cache_shape` in a subclass.") + + def get_usable_length(self, new_seq_length: int, layer_idx: int | None = 0) -> int: + """Given the sequence length of the new inputs, returns the usable length of the cache.""" + # Cache without size limit -> all cache is usable + # Cache with size limit -> if the length cache plus the length of the new inputs is larger the maximum cache + # length, we will need to evict part of the cache (and thus not all cache is usable) + max_length = self.get_max_cache_shape() + previous_seq_length = self.get_seq_length(layer_idx) + if max_length is not None and previous_seq_length + new_seq_length > max_length: + return max_length - new_seq_length + return previous_seq_length + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx] != []: + device = self.key_cache[layer_idx].device + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) + if self.value_cache[layer_idx] != []: + device = self.value_cache[layer_idx].device + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) + + @property + def seen_tokens(self): + logger.warning_once( + "The `seen_tokens` attribute is deprecated and will be removed in v4.41. Use the `cache_position` " + "model input instead." + ) + if hasattr(self, "_seen_tokens"): + return self._seen_tokens + else: + return None diff --git a/optimum/exporters/openvino/_traceable_decorator.py b/optimum/exporters/openvino/_traceable_decorator.py new file mode 100644 index 0000000000..e12f6f2fad --- /dev/null +++ b/optimum/exporters/openvino/_traceable_decorator.py @@ -0,0 +1,175 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import warnings +from collections import defaultdict +from functools import wraps + +from transformers.utils.generic import logger + + +try: + # transformers>=5.2 + from transformers.utils.output_capturing import _CAN_RECORD_REGISTRY, OutputRecorder +except ImportError: + from transformers.utils.generic import _CAN_RECORD_REGISTRY, OutputRecorder + + +# This is a fixed version of transformers.utils.generic.check_model_inputs +# that fixes issues related to export and tracing +# - adds support for positional args (use_cache), without which use_cache end up being passed twice +# - fixes issue with default capture_flags being None for some models +def traceable_check_model_inputs(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + use_cache = ( + kwargs["use_cache"] if kwargs.get("use_cache") is not None else getattr(self.config, "use_cache", None) + ) + if use_cache is not None: + if getattr(self, "gradient_checkpointing", False) and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + # Prevent passing use_cache twice + if "use_cache" in func.__code__.co_varnames: + use_cache_idx = func.__code__.co_varnames.index("use_cache") - 1 # minus 1 for 'self' + if len(args) > use_cache_idx: + args = list(args) + args[use_cache_idx] = use_cache + args = tuple(args) + else: + kwargs["use_cache"] = use_cache + + return_dict = kwargs.pop("return_dict", None) + if return_dict is None: + return_dict = getattr(self.config, "return_dict", True) + + all_args = kwargs.copy() + if "kwargs" in all_args: + for k, v in all_args["kwargs"].items(): + all_args[k] = v + + capture_flags = _CAN_RECORD_REGISTRY.get(str(self.__class__)) or {} # there is a weak ref for executorch + + recordable_keys = { + f"output_{k}": all_args.get( + f"output_{k}", + getattr( + self.config, + f"output_{k}", + all_args.get("output_attentions", getattr(self.config, "output_attentions", False)), + ), + ) + for k in capture_flags + } + + # We let cross attentions to be saved separately because some models add `cross-attn` layer + # when certain conditions are met. Let's output cross attention if attentions are requested (for BC) + if "output_attentions" in recordable_keys: + recordable_keys["output_cross_attentions"] = recordable_keys["output_attentions"] + + collected_outputs = defaultdict(tuple) + monkey_patched_layers = [] + + # Check attention implementation is properly set for capturing attention outputs + if recordable_keys.get("output_attentions", False): + supported_attn = ["eager", "eager_paged", "flex_attention"] + config_attn = getattr(self.config, "_attn_implementation", None) + sub_configs = [getattr(self.config, key, None) for key in self.config.sub_configs] + sub_configs_attn = [ + getattr(config, "_attn_implementation", None) for config in sub_configs if config is not None + ] + if config_attn not in supported_attn or any(attn not in supported_attn for attn in sub_configs_attn): + warnings.warn( + f"`output_attentions=True` is not supported with `attn_implementation` other than {supported_attn}. " + "Please use `model.set_attn_implementation('eager')` to enable capturing attention outputs.", + UserWarning, + stacklevel=2, + ) + + def make_capture_wrapper(module, orig_forward, key, index): + @wraps(orig_forward) + def wrapped_forward(*args, **kwargs): + if key == "hidden_states" and len(collected_outputs[key]) == 0: + collected_outputs[key] += (args[0],) + output = orig_forward(*args, **kwargs) + if not isinstance(output, tuple): + collected_outputs[key] += (output,) + elif output[index] is not None: + if key not in collected_outputs: + collected_outputs[key] = (output[index],) + else: + collected_outputs[key] += (output[index],) + return output + + return wrapped_forward + + if any(recordable_keys.values()): + capture_tasks = [] + for key, layer_specs in capture_flags.items(): + if not recordable_keys.get(f"output_{key}", False): + continue + if not isinstance(layer_specs, list): + layer_specs = [layer_specs] + for specs in layer_specs: + if not isinstance(specs, OutputRecorder): + index = 0 if "hidden_states" in key else 1 + class_name = None if not isinstance(specs, str) else specs + target_class = specs if not isinstance(specs, str) else None + specs = OutputRecorder(target_class=target_class, index=index, class_name=class_name) + capture_tasks.append((key, specs)) + + for name, module in self.named_modules(): + for key, specs in capture_tasks: + # The second check is for multimodals where only backbone layer suffix is available + if (specs.target_class is not None and isinstance(module, specs.target_class)) or ( + specs.class_name is not None and name.endswith(specs.class_name) + ): + if specs.layer_name is not None and specs.layer_name not in name: + continue + # Monkey patch forward + original_forward = module.forward + module.forward = make_capture_wrapper(module, original_forward, key, specs.index) + monkey_patched_layers.append((module, original_forward)) + + outputs = func(self, *args, **kwargs) + # Restore original forward methods + for module, original_forward in monkey_patched_layers: + module.forward = original_forward + + # Inject collected outputs into model output + for key in collected_outputs: + if key == "hidden_states": + if hasattr(outputs, "vision_hidden_states"): + collected_outputs[key] = collected_outputs[key][:-1] + collected_outputs[key] += (outputs.vision_hidden_states,) + elif hasattr(outputs, "last_hidden_state"): + collected_outputs[key] = collected_outputs[key][:-1] + collected_outputs[key] += (outputs.last_hidden_state,) + + outputs[key] = collected_outputs[key] + elif key == "attentions": + if isinstance(capture_flags[key], list) and len(capture_flags[key]) == 2: + outputs[key] = collected_outputs[key][0::2] + outputs["cross_" + key] = collected_outputs[key][1::2] + else: + outputs[key] = collected_outputs[key] + else: + outputs[key] = collected_outputs[key] + if return_dict is False: + outputs = outputs.to_tuple() + return outputs + + return wrapper diff --git a/optimum/exporters/openvino/base.py b/optimum/exporters/openvino/base.py new file mode 100644 index 0000000000..8c09341f08 --- /dev/null +++ b/optimum/exporters/openvino/base.py @@ -0,0 +1,643 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""OpenVINO configuration base classes.""" + +from __future__ import annotations + +import enum +import inspect +import itertools +import re +from abc import ABC +from collections import OrderedDict +from collections.abc import Iterable +from typing import Any, ClassVar + +from transformers import PretrainedConfig, PreTrainedModel + +from optimum.exporters.base import ExporterConfig +from optimum.exporters.openvino.patching_utils import ModelPatcher, PatchingSpec +from optimum.utils import DEFAULT_DUMMY_SHAPES, DummyInputGenerator, DummySeq2SeqPastKeyValuesGenerator, logging +from optimum.utils.doc import add_dynamic_docstring + + +logger = logging.get_logger(__name__) + + +GENERATE_DUMMY_DOCSTRING = r""" + Generates the dummy inputs necessary for tracing the model. If not explicitly specified, default input shapes are used. + + Args: + framework (`str`, defaults to `"pt"`): + The framework for which to create the dummy inputs. + batch_size (`int`, defaults to {batch_size}): + The batch size to use in the dummy inputs. + sequence_length (`int`, defaults to {sequence_length}): + The sequence length to use in the dummy inputs. + num_choices (`int`, defaults to {num_choices}): + The number of candidate answers provided for multiple choice task. + image_width (`int`, defaults to {width}): + The width to use in the dummy inputs for vision tasks. + image_height (`int`, defaults to {height}): + The height to use in the dummy inputs for vision tasks. + num_channels (`int`, defaults to {num_channels}): + The number of channels to use in the dummpy inputs for vision tasks. + feature_size (`int`, defaults to {feature_size}): + The number of features to use in the dummpy inputs for audio tasks in case it is not raw audio. + This is for example the number of STFT bins or MEL bins. + nb_max_frames (`int`, defaults to {nb_max_frames}): + The number of frames to use in the dummpy inputs for audio tasks in case the input is not raw audio. + audio_sequence_length (`int`, defaults to {audio_sequence_length}): + The number of frames to use in the dummpy inputs for audio tasks in case the input is raw audio. + + Returns: + `Dict`: A dictionary mapping the input names to dummy tensors in the proper framework format. +""" + + +class OpenVINOConfig(ExporterConfig, ABC): + VARIANTS: ClassVar[dict[str, str]] = {"default": "The default OpenVINO variant."} + DEFAULT_VARIANT = "default" + PATCHING_SPECS: list[PatchingSpec] | None = None + _MODEL_PATCHER = ModelPatcher + + _TASK_TO_COMMON_OUTPUTS = { # noqa: RUF012 + "audio-classification": OrderedDict({"logits": {0: "batch_size"}}), + "audio-frame-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "automatic-speech-recognition": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "audio-xvector": OrderedDict({"logits": {0: "batch_size"}, "embeddings": {0: "batch_size"}}), + "depth-estimation": OrderedDict({"predicted_depth": {0: "batch_size", 1: "height", 2: "width"}}), + "document-question-answering": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "feature-extraction": OrderedDict({"last_hidden_state": {0: "batch_size", 1: "sequence_length"}}), + "fill-mask": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "image-classification": OrderedDict({"logits": {0: "batch_size"}}), + "image-segmentation": OrderedDict({"logits": {0: "batch_size", 2: "height", 3: "width"}}), + "image-to-text": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "image-to-image": OrderedDict( + {"reconstruction": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + ), + "keypoint-detection": OrderedDict( + {"heatmaps": {0: "batch_size", 1: "num_keypoints", 2: "height", 3: "width"}} + ), + "mask-generation": OrderedDict({"logits": {0: "batch_size"}}), + "masked-im": OrderedDict({"reconstruction": {0: "batch_size"}}), + "multiple-choice": OrderedDict({"logits": {0: "batch_size", 1: "num_choices"}}), + "object-detection": OrderedDict( + { + "logits": {0: "batch_size", 1: "num_queries"}, + "pred_boxes": {0: "batch_size", 1: "num_queries"}, + } + ), + "question-answering": OrderedDict( + { + "start_logits": {0: "batch_size", 1: "sequence_length"}, + "end_logits": {0: "batch_size", 1: "sequence_length"}, + } + ), + "semantic-segmentation": OrderedDict({"logits": {0: "batch_size", 1: "num_labels", 2: "height", 3: "width"}}), + "text2text-generation": OrderedDict({"logits": {0: "batch_size", 1: "decoder_sequence_length"}}), + "text-classification": OrderedDict({"logits": {0: "batch_size"}}), + "text-generation": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "time-series-forecasting": OrderedDict({"prediction_outputs": {0: "batch_size"}}), + "token-classification": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "visual-question-answering": OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}), + "zero-shot-image-classification": OrderedDict( + { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + ), + "zero-shot-object-detection": OrderedDict( + { + "logits": {0: "batch_size", 1: "num_queries"}, + "pred_boxes": {0: "batch_size", 1: "num_queries"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + ), + } + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + preprocessors: list[Any] | None = None, + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + super().__init__(config=config, task=task, int_dtype=int_dtype, float_dtype=float_dtype) + + self.variant = "default" + self._preprocessors = preprocessors + + @property + def variant(self) -> str: + """For a given OpenVINO config, the variant of the model to export. + + This property allows to define variants of a given model, in case + different users would like to export the model differently (with different inputs/outputs, model split in several OpenVINO or not, etc.). + """ + return self._variant + + @variant.setter + def variant(self, value: str): + if value == "default" and hasattr(self, "DEFAULT_VARIANT"): + value = self.DEFAULT_VARIANT + if value not in self.VARIANTS: + raise ValueError( + f"The variant {value} is not supported for the OpenVINO config {self.__class__.__name__}." + ) + self._variant = value + + @property + def torch_to_ov_input_map(self) -> dict[str, str]: + """Dictionary mapping input names from the PyTorch model to input names from the exported OpenVINO model. + + Override the function when the input names and the exported OpenVINO input names are different. + + Returns: + `Dict[str, str]`: A dictionary mapping the PyTorch model input names to the exported OpenVINO model input names. + """ + return {} + + @property + def torch_to_ov_output_map(self) -> dict[str, str]: + """Dictionary mapping output names from the PyTorch model to output names from the exported OpenVINO model. + + Override the function when the output names and the exported OpenVINO output names are different. + + Returns: + `Dict[str, str]`: A dictionary mapping the PyTorch model output names to the exported OpenVINO model output names. + """ + return {} + + def rename_ambiguous_inputs(self, inputs) -> dict[str, dict[int, str]]: + """Updates the input names of the model to export. + + Override the function when the model input names are ambiguous or too generic. + + Returns: + `Dict[str, Dict[int, str]]`: Updated inputs. + """ + return inputs + + def ordered_inputs(self, model: PreTrainedModel) -> dict[str, dict[int, str]]: + """Re-orders the inputs using the model forward pass signature. + + Args: + model ([`transformers.PreTrainedModel`]): + The model for which we will use the OpenVINOConfig. + + Returns: + `Dict[str, Dict[int, str]]`: The properly ordered inputs. + """ + inputs = self.inputs + inputs = self.rename_ambiguous_inputs(inputs) + + ordered_inputs = {} + if hasattr(model, "forward"): + sig = inspect.signature(model.forward) + else: + sig = inspect.signature(model.call) + + for param in sig.parameters: + param_regex = re.compile(rf"{param}(\..*)?$") + to_insert = [] + for name, dynamic_axes in inputs.items(): + if re.match(param_regex, name): + to_insert.append((name, dynamic_axes)) + # TODO: figure out a smart way of re-ordering potential nested structures. + # to_insert = sorted(to_insert, key=lambda t: t[0]) + for name, dynamic_axes in to_insert: + name = self.torch_to_ov_input_map.get(name, name) + ordered_inputs[name] = dynamic_axes + return ordered_inputs + + @classmethod + def flatten_output_collection_property(cls, name: str, field: Iterable[Any]) -> dict[str, Any]: + """Flattens any potential nested structure expanding the name of the field with the index of the element within the structure. + + Args: + name (`str`): + The name of the nested structure. + field (`Iterable[Any]`): + The structure to potentially flattened. + + Returns: + `Dict[str, Any]`: Outputs with flattened structure and key mapping this new structure. + + """ + if isinstance(field[0], (list, tuple)): + return {f"{name}.{idx}": item for idx, item in enumerate(itertools.chain.from_iterable(field))} + else: + return {f"{name}.{idx}": item for idx, item in enumerate(field)} + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] + ) -> dict[str, Any]: + """Generates inputs for OpenVINO Runtime using the reference model inputs. + + Override this to run inference with seq2seq + models which have the encoder and decoder exported as separate OpenVINO files. + + Args: + reference_model_inputs (`Dict[str, Tensor]`): + Reference inputs for the model. + ov_input_names (`Optional[List[str]]`, defaults to `None`): + Names of the actual inputs to the OpenVINO model. This argument may be required as an unused + input to the model is automatically removed by torch.onnx.export (e.g. encoder_outputs in the decoder with past) + + Returns: + `Dict[str, Tensor]`: The mapping holding the kwargs to provide to the model's forward function + """ + return reference_model_inputs + + def patch_model_for_export( + self, model: PreTrainedModel, model_kwargs: dict[str, Any] | None = None + ) -> ModelPatcher: + return self._MODEL_PATCHER(self, model, model_kwargs=model_kwargs) + + +class OpenVINOConfigWithPast(OpenVINOConfig, ABC): + """Inherits from [`~exporters.onnx.OpenVINOConfig`]. A base class to handle the OpenVINO configuration of decoder-only models.""" + + PAD_ATTENTION_MASK_TO_PAST: bool = False + SUPPORTS_PAST: bool = True + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + preprocessors: list[Any] | None = None, + ): + self.use_past = use_past + self.use_past_in_inputs = use_past_in_inputs + + self.is_merged = False + self.use_cache_branch = None + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if not self.use_past_in_inputs: + common_outputs = super().outputs + # In the other cases, the sequence_length axis is not dynamic, always of length 1 + elif self.task == "feature-extraction": + common_outputs = OrderedDict({"last_hidden_state": {0: "batch_size"}}) + else: + common_outputs = OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}) + if self.use_past: + # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. + self.add_past_key_values(common_outputs, direction="outputs") + return common_outputs + + @property + def values_override(self) -> dict[str, Any] | None: + if hasattr(self._config, "use_cache"): + return {"use_cache": self.use_past} + + @add_dynamic_docstring(text=GENERATE_DUMMY_DOCSTRING, dynamic_elements=DEFAULT_DUMMY_SHAPES) + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs_generators = self._create_dummy_input_generator_classes(**kwargs) + + dummy_inputs = {} + input_names = [key for key in self.inputs if not key.startswith("past_key_values")] + if self.use_past_in_inputs and self.use_cache_branch is not False: + input_names.append("past_key_values") + + for input_name in input_names: + input_was_inserted = False + for dummy_input_gen in dummy_inputs_generators: + if dummy_input_gen.supports_input(input_name): + dummy_inputs[input_name] = self.overwrite_shape_and_generate_input( + dummy_input_gen, + input_name, + framework, + input_shapes=kwargs, + ) + input_was_inserted = True + break + if not input_was_inserted: + raise RuntimeError( + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' + ) + + # refer to https://github.com/huggingface/optimum/pull/764 + if ( + self.use_past_in_inputs + and self.PAD_ATTENTION_MASK_TO_PAST + and self.use_cache_branch is not False + and "attention_mask" in dummy_inputs + and self.task == "text-generation" + ): + seq_len = dummy_inputs["input_ids"].shape[1] + past_seq_len = dummy_inputs["past_key_values"][0][1].shape[-2] + dummy_inputs["attention_mask"] = DummyInputGenerator.pad_input_on_dim( + dummy_inputs["attention_mask"], desired_length=past_seq_len + seq_len, dim=1 + ) + + return dummy_inputs + + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + """The shape passed to the dummy input generator may not always be correct for all of the inputs it manages. + + This method allows + to overwrite some shapes, and generate the dummy input. This should probably be refactored more elegantly. + """ + # models from TextSeq2SeqOpenVINOConfig use decoder_input_ids as input name + # while models from TextDecoderOpenVINOConfig use input_ids, hence the check for both + + # NOTE: The check `self.task != "text-generation" is added following the use of a single OpenVINO for both without/with KV cache, without subgraphs. + # This overwrite may be moved to OpenVINOSeq2SeqConfigWithPast, but I am afraid it would break encoder-decoder models. + if ( + self.use_past + and self.use_past_in_inputs + and self.use_cache_branch is not False + and input_name in ["decoder_input_ids", "input_ids", "position_ids"] + and self.task != "text-generation" + ): + sequence_length = dummy_input_gen.sequence_length + # Use a sequence length of 1 when the KV cache is already populated. + dummy_input_gen.sequence_length = 1 + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + dummy_input_gen.sequence_length = sequence_length + else: + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + + return dummy_input + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + """Fills `input_or_outputs` mapping with past_key_values dynamic axes considering the direction. + + Args: + inputs_or_outputs (`Dict[str, Dict[int, str]]`): + The mapping to fill. + direction (`str`): + either "inputs" or "outputs", it specifies whether `input_or_outputs` is the input mapping or the + output mapping, this is important for axes naming. + """ + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size", 2: decoder_sequence_name} + + def flatten_past_key_values(self, flattened_output, name, idx, t): + flattened_output[f"{name}.{idx}.key"] = t[0] + flattened_output[f"{name}.{idx}.value"] = t[1] + + def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> dict[str, Any]: + flattened_output = {} + if name in ["present", "past_key_values"]: + for idx, t in enumerate(field): + self.flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output = super().flatten_output_collection_property(name, field) + + return flattened_output + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] + ) -> dict[str, Any]: + if self.is_merged is True and self.use_cache_branch is True: + reference_model_inputs["use_cache_branch"] = DummyInputGenerator.constant_tensor(shape=[1], value=True) + elif self.is_merged is True and self.use_cache_branch is False: + reference_model_inputs["use_cache_branch"] = DummyInputGenerator.constant_tensor(shape=[1], value=False) + + # We don't support optional inputs for now, so even though the non-cache branch is used, + # dummy past key values are necessary + batch_size = reference_model_inputs["input_ids"].shape[0] + pkv_generator = self.DUMMY_PKV_GENERATOR_CLASS( + task=self.task, normalized_config=self._normalized_config, sequence_length=1, batch_size=batch_size + ) + reference_model_inputs["past_key_values"] = pkv_generator.generate( + "past_key_values", framework="pt", int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + + return super().generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) + + +class ConfigBehavior(str, enum.Enum): + """Specifies the behavior of the [`~exporters.openvino.base.OpenVINOSeq2SeqConfigWithPast`]. + + - MONOLITH: the config can be used to export the whole seq2seq model as a single file. + - ENCODER: the config can be used to export the encoder part of the seq2seq model. + - DECODER: the config can be used to export the decoder part of the seq2seq model. + """ + + MONOLITH = "monolith" + ENCODER = "encoder" + DECODER = "decoder" + + +class OpenVINOSeq2SeqConfigWithPast(OpenVINOConfigWithPast): + """Inherits from [`~exporters.openvino.OpenVINOConfigWithPast`]. A base class to handle the OpenVINO configuration of encoder-decoder models.""" + + DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + behavior: ConfigBehavior = ConfigBehavior.MONOLITH, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + preprocessors=preprocessors, + ) + self._behavior = behavior + + if self._behavior is ConfigBehavior.ENCODER: + self.task = "feature-extraction" + self.use_past_in_inputs = False + + def with_behavior( + self, + behavior: str | ConfigBehavior, + use_past: bool = False, + use_past_in_inputs: bool = False, + ) -> OpenVINOSeq2SeqConfigWithPast: + """Creates a copy of the current OpenVINOConfig but with a different `ConfigBehavior` and `use_past` value. + + Args: + behavior ([`ConfigBehavior`]): + The behavior to use for the new instance. + use_past (`bool`, defaults to `False`): + Whether or not the OpenVINO config to instantiate is for a model using KV cache. + use_past_in_inputs (`bool`, defaults to `False`): + Whether the KV cache is to be passed as an input to the OpenVINO. + + Returns: + `OpenVINOSeq2SeqConfigWithPast` + """ + if isinstance(behavior, str) and not isinstance(behavior, ConfigBehavior): + behavior = ConfigBehavior(behavior) + + ov_config = self.__class__( + self._config, + task=self.task, + int_dtype=self.int_dtype, + float_dtype=self.float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + behavior=behavior, + preprocessors=self._preprocessors, + ) + ov_config.variant = self.variant + return ov_config + + @property + def torch_to_ov_input_map(self) -> dict[str, str]: + if self._behavior is ConfigBehavior.DECODER: + return { + "decoder_input_ids": "input_ids", + "decoder_attention_mask": "attention_mask", + "encoder_outputs": "encoder_hidden_states", + "attention_mask": "encoder_attention_mask", + } + return {} + + @property + def outputs(self) -> dict[str, dict[int, str]]: + common_outputs = super(OpenVINOConfigWithPast, self).outputs + # Renaming the outputs axes properly. + for name, axes_names in common_outputs.items(): + if self._behavior is ConfigBehavior.ENCODER or "encoder" in name: + sequence_name = "encoder_sequence_length" + else: + sequence_name = "decoder_sequence_length" + + new_axes_names = {} + for axis_idx, axis_name in axes_names.items(): + if "sequence" in axis_name: + if self.use_past_in_inputs is False or self.is_merged is True: + new_axes_names[axis_idx] = sequence_name + else: + # Trick to force it since OpenVINO sometimes infer a dynamic axis where it's not. + new_axes_names[axis_idx] = "1" + else: + new_axes_names[axis_idx] = axis_name + common_outputs[name] = new_axes_names + + if self.use_past: + # When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output. + self.add_past_key_values(common_outputs, direction="outputs") + + return common_outputs + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_decoder_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_decoder_sequence_length + decoder_sequence_length" + name = "present" + + for i in range(self._normalized_config.decoder_num_layers): + inputs_or_outputs[f"{name}.{i}.decoder.key"] = {0: "batch_size", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.decoder.value"] = {0: "batch_size", 2: decoder_sequence_name} + + if ( + self.is_merged is True + or (self._behavior is ConfigBehavior.DECODER and not self.use_past_in_inputs) + or direction == "inputs" + ): + inputs_or_outputs[f"{name}.{i}.encoder.key"] = {0: "batch_size", 2: "encoder_sequence_length"} + inputs_or_outputs[f"{name}.{i}.encoder.value"] = {0: "batch_size", 2: "encoder_sequence_length"} + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if len(t) not in [2, 4]: + raise ValueError( + "past_key_values to flatten should be of length 2 (self-attention only) or 4 (self and cross attention)." + ) + + flattened_output[f"{name}.{idx}.decoder.key"] = t[0] + flattened_output[f"{name}.{idx}.decoder.value"] = t[1] + if len(t) == 4: + flattened_output[f"{name}.{idx}.encoder.key"] = t[2] + flattened_output[f"{name}.{idx}.encoder.value"] = t[3] + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + + if ( + self.use_past_in_inputs + and self.PAD_ATTENTION_MASK_TO_PAST + and self.use_cache_branch is not False + and "decoder_attention_mask" in dummy_inputs + ): + seq_len = dummy_inputs["decoder_input_ids"].shape[1] + past_seq_len = dummy_inputs["past_key_values"][0][1].shape[-2] + dummy_inputs["decoder_attention_mask"] = DummyInputGenerator.pad_input_on_dim( + dummy_inputs["decoder_attention_mask"], desired_length=past_seq_len + seq_len, dim=1 + ) + + return dummy_inputs + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] + ) -> dict[str, Any]: + if self._behavior is ConfigBehavior.DECODER: + if "decoder_input_ids" in reference_model_inputs: + reference_model_inputs["input_ids"] = reference_model_inputs.pop("decoder_input_ids") + if "attention_mask" in reference_model_inputs: + reference_model_inputs["encoder_attention_mask"] = reference_model_inputs.pop("attention_mask") + if "decoder_attention_mask" in reference_model_inputs: + reference_model_inputs["attention_mask"] = reference_model_inputs.pop("decoder_attention_mask") + if "encoder_outputs" in reference_model_inputs: + if "encoder_hidden_states" in ov_input_names: + reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop("encoder_outputs")[0] + else: + reference_model_inputs.pop("encoder_outputs") + + return super().generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) diff --git a/optimum/exporters/openvino/config.py b/optimum/exporters/openvino/config.py new file mode 100644 index 0000000000..f248795b72 --- /dev/null +++ b/optimum/exporters/openvino/config.py @@ -0,0 +1,348 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Common OpenVINO configuration classes that handle most of the features for building model specific configurations.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable +from typing import Any + +from transformers import PretrainedConfig + +from optimum.exporters.openvino.base import ( + ConfigBehavior, + OpenVINOConfig, + OpenVINOConfigWithPast, + OpenVINOSeq2SeqConfigWithPast, +) +from optimum.exporters.tasks import TasksManager +from optimum.utils import ( + DummyAudioInputGenerator, + DummyBboxInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyVisionInputGenerator, + logging, +) + + +logger = logging.get_logger(__name__) + + +class TextEncoderOpenVINOConfig(OpenVINOConfig): + """Handles encoder-based text architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) + + +class TextDecoderOpenVINOConfig(OpenVINOConfigWithPast): + """Handles decoder-based text architectures.""" + + PAD_ATTENTION_MASK_TO_PAST = True + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = DummyPastKeyValuesGenerator + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + preprocessors=preprocessors, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + if self.use_past_in_inputs: + common_inputs = {"input_ids": {0: "batch_size", 1: "sequence_length"}} + common_inputs["attention_mask"] = {0: "batch_size", 1: "past_sequence_length + sequence_length"} + self.add_past_key_values(common_inputs, direction="inputs") + else: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + return common_inputs + + @property + def outputs(self) -> dict[str, dict[int, str]]: + if self.is_merged is False: + common_outputs = super().outputs + else: + # in the merged case, we need to allow the `sequence_length` to be variable, as it is not 1 + # during the first pass without past key values + common_outputs = OrderedDict({"logits": {0: "batch_size", 1: "sequence_length"}}) + self.add_past_key_values(common_outputs, direction="outputs") + return common_outputs + + +class TextDecoderWithPositionIdsOpenVINOConfig(TextDecoderOpenVINOConfig): + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = super().inputs + + # Decoders based on GPT2 require a position_ids input to avoid generating wrong position_ids in the model itself: + # https://github.com/huggingface/transformers/blob/v4.33.1/src/transformers/models/gpt2/modeling_gpt2.py#L802 + if self.task in {"text-generation", "feature-extraction"}: + common_inputs["position_ids"] = {0: "batch_size", 1: "sequence_length"} + + return common_inputs + + +class TextSeq2SeqOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): + """Handles encoder-decoder-based text architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_ids"] = {0: "batch_size", 1: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + def _create_dummy_input_generator_classes(self, **kwargs) -> list[DummyInputGenerator]: + dummy_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[0]( + self.task, self._normalized_config, **kwargs + ) + dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1]( + self.task, + self._normalized_config, + **kwargs, + ) + dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2]( + self.task, + self._normalized_config, + encoder_sequence_length=dummy_text_input_generator.sequence_length, + **kwargs, + ) + dummy_inputs_generators = [ + dummy_text_input_generator, + dummy_decoder_text_input_generator, + dummy_seq2seq_past_key_values_generator, + ] + + return dummy_inputs_generators + + +class VisionOpenVINOConfig(OpenVINOConfig): + """Handles vision architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator,) + + +class TextAndVisionOpenVINOConfig(OpenVINOConfig): + """Handles multi-modal text and vision architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyVisionInputGenerator, DummyBboxInputGenerator) + + +class AudioOpenVINOConfig(OpenVINOConfig): + """Handles audio architectures.""" + + DUMMY_INPUT_GENERATOR_CLASSES = (DummyAudioInputGenerator,) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + return {"input_values": {0: "batch_size", 1: "sequence_length"}} + + +class AudioToTextOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyAudioInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_features"] = {0: "batch_size", 1: "feature_size", 2: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + +class EncoderDecoderBaseOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator,) + + def __init__( + self, + config: PretrainedConfig, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + use_past: bool = False, + use_past_in_inputs: bool = False, + behavior: ConfigBehavior = ConfigBehavior.MONOLITH, + preprocessors: list[Any] | None = None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + use_past=use_past, + use_past_in_inputs=use_past_in_inputs, + behavior=behavior, + preprocessors=preprocessors, + ) + + self.is_decoder_with_past = False + + # Set up the encoder openvino config. + encoder_ov_config_constructor = TasksManager.get_exporter_config_constructor( + exporter="openvino", + task="feature-extraction", + model_type=config.encoder.model_type, + library_name="transformers", + ) + self._encoder_ov_config = encoder_ov_config_constructor( + config.encoder, int_dtype=int_dtype, float_dtype=float_dtype, preprocessors=preprocessors + ) + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = self._encoder_ov_config._normalized_config + + # Set up the decoder openvino config. + decoder_ov_config_constructor = TasksManager.get_exporter_config_constructor( + exporter="openvino", + task="feature-extraction", + model_type=config.decoder.model_type, + library_name="transformers", + ) + kwargs = {} + if issubclass(decoder_ov_config_constructor.func, OpenVINOConfigWithPast): + self.is_decoder_with_past = True + kwargs["use_past"] = use_past + else: + self.use_past = False + + if use_past and not self.is_decoder_with_past: + raise ValueError( + f"The decoder part of the encoder-decoder model is {config.decoder.model_type} which does not need " + "past key values." + ) + + self._decoder_ov_config = decoder_ov_config_constructor( + config.decoder, int_dtype=int_dtype, float_dtype=float_dtype, preprocessors=preprocessors, **kwargs + ) + if issubclass(decoder_ov_config_constructor.func, OpenVINOSeq2SeqConfigWithPast): + self._decoder_ov_config = self._decoder_ov_config.with_behavior( + self._behavior, use_past=kwargs["use_past"], use_past_in_inputs=use_past_in_inputs + ) + + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_ov_config._normalized_config + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS = self._decoder_ov_config._normalized_config + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.encoder_num_attention_heads = ( + self._decoder_ov_config._normalized_config.num_attention_heads + ) + self._normalized_config.DECODER_NORMALIZED_CONFIG_CLASS.decoder_num_attention_heads = ( + self._decoder_ov_config._normalized_config.num_attention_heads + ) + + if isinstance(self._decoder_ov_config, OpenVINOSeq2SeqConfigWithPast): + self._past_key_values_generator = ( + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + ) + else: + self._past_key_values_generator = ( + DummySeq2SeqDecoderTextInputGenerator, + DummyPastKeyValuesGenerator, + ) + + self.DUMMY_INPUT_GENERATOR_CLASSES += self._past_key_values_generator + + @property + def inputs(self) -> dict[str, dict[int, str]]: + common_inputs = {} + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_ids"] = {0: "batch_size", 1: "encoder_sequence_length"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + common_inputs["attention_mask"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + + return common_inputs + + def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], direction: str): + if self.is_decoder_with_past: + return self._decoder_ov_config.add_past_key_values(inputs_or_outputs, direction) + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if self.is_decoder_with_past: + return self._decoder_ov_config.flatten_past_key_values(flattened_output, name, idx, t) + + def flatten_output_collection_property(self, name: str, field: Iterable[Any]) -> dict[str, Any]: + return self._decoder_ov_config.flatten_output_collection_property(name, field) + + def generate_dummy_inputs_for_validation( + self, reference_model_inputs: dict[str, Any], ov_input_names: list[str] + ) -> dict[str, Any]: + if self._behavior is ConfigBehavior.ENCODER: + return self._encoder_ov_config.generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) + else: + if self._behavior is ConfigBehavior.DECODER: + if "decoder_input_ids" in reference_model_inputs: + reference_model_inputs["input_ids"] = reference_model_inputs.pop("decoder_input_ids") + if "attention_mask" in reference_model_inputs: + reference_model_inputs["encoder_attention_mask"] = reference_model_inputs.pop("attention_mask") + if "encoder_outputs" in reference_model_inputs: + if "encoder_hidden_states" in ov_input_names: + reference_model_inputs["encoder_hidden_states"] = reference_model_inputs.pop( + "encoder_outputs" + )[0] + else: + reference_model_inputs.pop("encoder_outputs") + + return self._decoder_ov_config.generate_dummy_inputs_for_validation(reference_model_inputs, ov_input_names) diff --git a/optimum/exporters/openvino/convert.py b/optimum/exporters/openvino/convert.py index a1a928501d..dbe8e403cd 100644 --- a/optimum/exporters/openvino/convert.py +++ b/optimum/exporters/openvino/convert.py @@ -29,6 +29,23 @@ from openvino import Model, save_model from openvino.exceptions import OVTypeError from openvino.tools.ovc import convert_model +from optimum.exporters.openvino.utils import ( + MULTI_MODAL_TEXT_GENERATION_MODELS, + ONNX_SUPPORTED_ARCHITECTURES, + OV_XML_FILE_NAME, + _get_dynamic_shapes_info, + _get_input_info, + _get_kokoro_submodels_fn_and_export_configs, + _get_model_dtype, + _get_open_clip_submodels_fn_and_export_configs, + _normalize_dummy_inputs, + allow_skip_tracing_check, + clear_class_registry, + remove_none_from_dummy_inputs, + save_config, + save_preprocessors, + set_simplified_chat_template, +) from optimum.exporters.tasks import TasksManager from optimum.exporters.utils import ( DECODER_NAME, @@ -50,11 +67,12 @@ _torch_version, _transformers_version, compare_versions, + is_diffusers_available, + is_nncf_available, is_transformers_version, ) -from optimum.utils import DEFAULT_DUMMY_SHAPES, is_diffusers_available +from optimum.utils import DEFAULT_DUMMY_SHAPES -from ...intel.utils.import_utils import is_nncf_available from ...intel.utils.modeling_utils import _infer_library_from_model_or_model_class from .stateful import ( ensure_export_task_support_stateful, @@ -62,23 +80,6 @@ ensure_stateful_is_available, patch_stateful, ) -from .utils import ( - MULTI_MODAL_TEXT_GENERATION_MODELS, - ONNX_SUPPORTED_ARCHITECTURES, - OV_XML_FILE_NAME, - _get_dynamic_shapes_info, - _get_input_info, - _get_kokoro_submodels_fn_and_export_configs, - _get_model_dtype, - _get_open_clip_submodels_fn_and_export_configs, - _normalize_dummy_inputs, - allow_skip_tracing_check, - clear_class_registry, - remove_none_from_dummy_inputs, - save_config, - save_preprocessors, - set_simplified_chat_template, -) logger = logging.getLogger(__name__) @@ -92,14 +93,14 @@ if TYPE_CHECKING: - from optimum.exporters.onnx.base import OnnxConfig + from optimum.exporters.openvino.base import OpenVINOConfig from optimum.intel.openvino.configuration import OVConfig def _set_runtime_options( models_and_export_configs: Dict[ str, - Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"], + Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OpenVINOConfig"], ], task: str, library_name: str, @@ -128,7 +129,7 @@ def _save_model( path: str, ov_config: Optional["OVConfig"] = None, library_name: Optional[str] = None, - config: "OnnxConfig" = None, + config: "OpenVINOConfig" = None, ): compress_to_fp16 = ov_config is not None and ov_config.dtype == "fp16" model = _add_version_info_to_model(model, library_name) @@ -146,9 +147,8 @@ def _save_model( def export( model: Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], - config: "OnnxConfig", + config: "OpenVINOConfig", output: Path, - opset: Optional[int] = None, device: str = "cpu", input_shapes: Optional[Dict] = None, model_kwargs: Optional[Dict[str, Any]] = None, @@ -163,12 +163,10 @@ def export( Args: model ([`PreTrainedModel`]): The model to export. - config ([`~exporters.onnx.config.OnnxConfig`]): - The ONNX configuration associated with the exported model. + config ([`~exporters.openvino.config.OpenVINOConfig`]): + The OpenVINO configuration associated with the exported model. output (`Path`): Directory to store the exported model. - opset (`Optional[int]`, defaults to `None`): - The version of the ONNX operator set to use. device (`str`, *optional*, defaults to `cpu`): The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for export on CUDA devices. @@ -181,7 +179,7 @@ def export( Returns: `Tuple[List[str], List[str]]`: A tuple with an ordered list of the model's inputs, and the named inputs from - the ONNX configuration. + the OpenVINO configuration. """ if not is_torch_available(): raise ImportError( @@ -221,7 +219,6 @@ def export( return export_pytorch( model, config, - opset, output, device=device, input_shapes=input_shapes, @@ -235,77 +232,9 @@ def export( raise RuntimeError("You either provided a non-PyTorch model or the PyTorch library is not installed.") -def export_pytorch_via_onnx( - model: Union["PreTrainedModel", "ModelMixin"], - config: "OnnxConfig", - opset: int, - output: Path, - device: str = "cpu", - input_shapes: Optional[Dict] = None, - model_kwargs: Optional[Dict[str, Any]] = None, - ov_config: Optional["OVConfig"] = None, - library_name: Optional[str] = None, -): - """ - Exports a PyTorch model to an OpenVINO Intermediate Representation via ONNX export. - - Args: - model ([`PreTrainedModel`]): - The model to export. - config ([`~exporters.onnx.config.OnnxConfig`]): - The configuration associated with the exported model. - opset (`int`): - The version of the ONNX operator set to use. - output (`Path`): - Directory to store the exported model. - device (`str`, defaults to `"cpu"`): - The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for - export on CUDA devices. - input_shapes (`optional[Dict]`, defaults to `None`): - If specified, allows to use specific shapes for the example input provided to the exporter. - model_kwargs (optional[Dict[str, Any]], defaults to `None`): - Additional kwargs for model export. - ov_config (`OVConfig`, *optional*): - The configuration containing the parameters related to quantization. - - Returns: - `Tuple[List[str], List[str], bool]`: A tuple with an ordered list of the model's inputs, and the named inputs from - the ONNX configuration and boolean flag - was legacy ONNX path were applied to model or not. - """ - import torch - - from optimum.exporters.onnx.convert import export_pytorch as export_pytorch_to_onnx - - output = Path(output) - orig_torch_onnx_export = torch.onnx.export - torch.onnx.export = functools.partial(orig_torch_onnx_export, do_constant_folding=False) - model.config.torchscript = False - model.config.return_dict = True - onnx_output = output.with_suffix(".onnx") - input_names, output_names = export_pytorch_to_onnx( - model, config, opset, onnx_output, device, input_shapes, model_kwargs - ) - torch.onnx.export = orig_torch_onnx_export - ov_model = convert_model(str(onnx_output)) - - library_name = _infer_library_from_model_or_model_class(model=model, library_name=library_name) - - _save_model( - ov_model, - output.parent / OV_XML_FILE_NAME if output.suffix != ".xml" else output, - ov_config=ov_config, - library_name=library_name, - config=config, - ) - del ov_model - gc.collect() - return input_names, output_names, True - - def export_pytorch( model: Union["PreTrainedModel", "ModelMixin"], - config: "OnnxConfig", - opset: int, + config: "OpenVINOConfig", output: Path, device: str = "cpu", input_shapes: Optional[Dict] = None, @@ -321,10 +250,8 @@ def export_pytorch( Args: model ([`PreTrainedModel`]): The model to export. - config ([`~exporters.onnx.config.OnnxConfig`]): + config ([`~exporters.openvino.config.OpenVINOConfig`]): The configuration associated with the exported model. - opset (`int`): - The version of the ONNX operator set to use. output (`Path`): Directory to store the exported model. device (`str`, defaults to `"cpu"`): @@ -341,7 +268,7 @@ def export_pytorch( Returns: `Tuple[List[str], List[str], bool]`: A tuple with an ordered list of the model's inputs, and the named inputs from - the ONNX configuration and boolean flag - was legacy ONNX path were applied to model or not. + the OpenVINO configuration and boolean flag - was legacy OpenVINO path were applied to model or not. """ import torch from torch.utils._pytree import tree_map @@ -482,10 +409,9 @@ def ts_patched_forward(*args, **kwargs): def export_models( models_and_export_configs: Dict[ - str, Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OnnxConfig"] + str, Tuple[Union["PreTrainedModel", "ModelMixin", "DiffusionPipeline"], "OpenVINOConfig"] ], output_dir: Path, - opset: Optional[int] = None, output_names: Optional[List[str]] = None, device: str = "cpu", input_shapes: Optional[Dict] = None, @@ -499,9 +425,8 @@ def export_models( Export the models to OpenVINO IR format Args: - models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "ModelMixin"], "OnnxConfig"]): + models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "ModelMixin"], "OpenVINOConfig"]): output_dir (Path): output directory for saving models - opset (Optional[int], optional, Default to None): ONNX export opset output_names (Optional[List[str]], optional, Defaults to None): model output names device (str, optional, Defaults to "cpu"): The device on which the model will be exported. Either `cpu` or `cuda`. Only PyTorch is supported for @@ -519,7 +444,7 @@ def export_models( ValueError: if custom names set not equal of number of models Returns: - list of input_names and output_names from ONNX configuration + list of input_names and output_names from OpenVINO configuration """ outputs = [] @@ -539,7 +464,6 @@ def export_models( model=submodel, config=sub_export_config, output=output_path, - opset=opset, device=device, input_shapes=input_shapes, model_kwargs=model_kwargs, @@ -646,9 +570,8 @@ def export_from_model( task: Optional[str] = None, ov_config: Optional["OVConfig"] = None, stateful: bool = True, - opset: Optional[int] = None, model_kwargs: Optional[Dict[str, Any]] = None, - custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None, + custom_export_configs: Optional[Dict[str, "OpenVINOConfig"]] = None, fn_get_submodels: Optional[Callable] = None, preprocessors: List = None, device: str = "cpu", @@ -673,8 +596,8 @@ def export_from_model( model_type = getattr(model.config, "model_type", None) or "" if model_type in ONNX_SUPPORTED_ARCHITECTURES: - logger.warning( - f"The OpenVINO export of {model_type} models is not officially supported by optimum-intel, export at your own risks." + raise ValueError( + f"The OpenVINO export of {model_type} models is not officially supported by optimum-intel and has been removed." ) custom_architecture = library_name == "transformers" and model_type not in TasksManager._SUPPORTED_MODEL_TYPE @@ -724,7 +647,7 @@ def export_from_model( and not getattr(model, "_supports_cache_class", is_transformers_version(">=", "4.54")) ): stateful = False - # TODO: support onnx_config.py in the model repo + # TODO: support ov_config.py in the model repo if custom_architecture and custom_export_configs is None: raise ValueError( f"Trying to export a {model_type} model, that is a custom or unsupported architecture, but no custom export configuration was passed as `custom_export_configs`. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum-intel/issues if you would like the model type {model_type} to be supported natively in the OpenVINO export." @@ -888,7 +811,6 @@ def export_from_model( device=device, ov_config=ov_config, stateful=stateful_submodels, - opset=opset, model_kwargs=model_kwargs, patch_16bit_model=patch_16bit_model, library_name=library_name, diff --git a/optimum/exporters/openvino/input_generators.py b/optimum/exporters/openvino/input_generators.py new file mode 100644 index 0000000000..da2fc7a7f7 --- /dev/null +++ b/optimum/exporters/openvino/input_generators.py @@ -0,0 +1,1737 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Tuple + +import torch + +from optimum.intel.utils.import_utils import is_diffusers_version +from optimum.utils import ( + DEFAULT_DUMMY_SHAPES, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyTimestepInputGenerator, + DummyVisionInputGenerator, + FalconDummyPastKeyValuesGenerator, + MistralDummyPastKeyValuesGenerator, + NormalizedTextConfig, + is_transformers_version, +) +from optimum.utils.input_generators import DTYPE_MAPPER +from optimum.utils.normalized_config import NormalizedConfig, NormalizedVisionConfig + + +class GPTBigCodeDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__(self, task: str, normalized_config: NormalizedTextConfig, **kwargs): + super().__init__(task=task, normalized_config=normalized_config, **kwargs) + self.multi_query = normalized_config.multi_query + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if is_transformers_version("<", "4.54"): + if self.multi_query: + shape = ( + self.batch_size, + self.sequence_length, + self.hidden_size // self.num_attention_heads * 2, + ) + else: + shape = ( + self.batch_size, + self.num_attention_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads * 2, + ) + pkv = [ + self.random_float_tensor(shape, framework=framework, dtype=float_dtype) for _ in range(self.num_layers) + ] + + else: + if self.multi_query: + shape = ( + self.batch_size, + 1, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + else: + shape = ( + self.batch_size, + self.num_attention_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + pkv = [ + ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + return pkv + + +class DummyQwen3VLLMInputGenerator(DummyTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "input_ids", + "attention_mask", + "token_type_ids", + "position_ids", + "visual_pos_masks", + "deepstack_visual_embeds", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + random_num_choices_range: Optional[Tuple[int, int]] = None, + padding_side: str = "right", + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + num_choices=num_choices, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + random_num_choices_range=random_num_choices_range, + padding_side=padding_side, + **kwargs, + ) + self.embed_dim = normalized_config.hidden_size + self.num_layers = len(self.normalized_config.deepstack_visual_indexes) + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + bool_dtype: str = "bool", + ): + if input_name == "deepstack_visual_embeds": + return self.random_float_tensor( + [self.num_layers, 2 * self.sequence_length, self.embed_dim], framework=framework, dtype=float_dtype + ) + if input_name == "visual_pos_masks": + return self.constant_tensor( + shape=[self.batch_size, self.sequence_length], + framework=framework, + value=1, + dtype=DTYPE_MAPPER.pt(bool_dtype), + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class OVMiniCPM3DummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + **kwargs, + ) + self.v_head_dim = getattr(normalized_config, "v_head_dim", self.hidden_size // self.num_attention_heads) + self.k_head_dim = normalized_config.qk_nope_head_dim + normalized_config.qk_rope_head_dim + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + v_shape = ( + self.batch_size, + self.num_key_value_heads, + self.sequence_length, + self.v_head_dim, + ) + k_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.k_head_dim) + return [ + ( + self.random_float_tensor(k_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(v_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class ChatGLM2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.multi_query_group_num = normalized_config.multi_query_group_num + self.head_dim = normalized_config.kv_channels + self.standart_cache_layout = hasattr(normalized_config, "rope_ratio") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if not self.standart_cache_layout: + pkv_shape = ( + self.sequence_length, + self.batch_size, + self.multi_query_group_num, + self.head_dim, + ) + else: + pkv_shape = ( + self.batch_size, + self.multi_query_group_num, + self.sequence_length, + self.head_dim, + ) + return [ + ( + self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class Eagle3DummyGenerator(DummyInputGenerator): + """ + Dummy input generator for Eagle-3 speculative decoding. + + This generator produces synthetic `hidden_states` tensors that mimic the + intermediate hidden-state outputs of a *main (target) model*, which are + required by the Eagle-3 draft model during speculative decoding. + """ + + SUPPORTED_INPUT_NAMES = ("hidden_states",) + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.batch_size = batch_size + self.sequence_length = sequence_length + self.hidden_size = normalized_config.hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + # hidden_states is provided as a concatenation of three hidden-layer outputs from the main model + shape = ( + self.batch_size, + self.sequence_length, + self.hidden_size * 3, + ) + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class Eagle3VLMDummyGenerator(DummyInputGenerator): + """ + Dummy input generator for VLM Eagle-3 speculative decoding. + + Produces `inputs_embeds` (float) and 3D `position_ids` (MRoPE) + required by VLM Eagle-3 draft models targeting Qwen3-VL. + """ + + SUPPORTED_INPUT_NAMES = ("inputs_embeds", "position_ids") + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.batch_size = batch_size + self.sequence_length = sequence_length + self.hidden_size = normalized_config.hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "inputs_embeds": + shape = (self.batch_size, self.sequence_length, self.hidden_size) + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + if input_name == "position_ids": + # The rotary embedding is MRoPE (Multimodal RoPE) + # MRoPE encodes position along three independent axes: temporal, height, and width + # https://github.com/Tencent/AngelSlim/blob/main/angelslim/compressor/speculative/train/models/draft/llama_eagle3.py#L211 + shape = (3, self.batch_size, self.sequence_length) + return self.random_int_tensor(shape, max_value=self.sequence_length, framework=framework, dtype=int_dtype) + + +class QwenDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.kv_channels = normalized_config.kv_channels + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) + past_value_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) + return [ + ( + self.random_float_tensor(past_key_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(past_value_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class OVFalconDummyPastKeyValuesGenerator(FalconDummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + **kwargs, + ) + if normalized_config.new_decoder_architecture: + self.num_kv_heads = normalized_config.num_attention_heads + else: + self.num_kv_heads = normalized_config.num_kv_heads if not normalized_config.multi_query else 1 + + self.head_dim = self.hidden_size // self.num_attention_heads + + +class AquilaDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task, + normalized_config, + batch_size, + sequence_length, + random_batch_size_range, + random_sequence_length_range, + **kwargs, + ) + self.num_key_value_heads = getattr( + normalized_config, "num_key_value_heads", normalized_config.num_attention_heads + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + shape = ( + self.batch_size, + self.num_key_value_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + return [ + ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.num_layers) + ] + + +class Gemma4DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.num_key_value_heads = normalized_config.num_key_value_heads + self.head_dim = normalized_config.head_dim + self.global_head_dim = getattr(normalized_config.config, "global_head_dim", self.head_dim) + self.layer_types = normalized_config.config.layer_types + self.num_kv_shared_layers = normalized_config.config.num_kv_shared_layers + self.sliding_window = normalized_config.config.sliding_window + # Full-attention layers use fewer KV heads than sliding-attention layers (e.g. 2 vs 8 for 26B-A4B) + self.num_global_key_value_heads = ( + getattr(normalized_config.config, "num_global_key_value_heads", None) or self.num_key_value_heads + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + # some layers do not produce their own KV-cache, they use the shared KV-cache + if self.num_kv_shared_layers > 0: + layer_types = self.layer_types[: -self.num_kv_shared_layers] + else: + layer_types = self.layer_types + past_kv_values = [] + for layer_type in layer_types: + if layer_type == "sliding_attention": + shape = ( + self.batch_size, + self.num_key_value_heads, + self.sliding_window, + self.head_dim, + ) + else: + shape = ( + self.batch_size, + self.num_global_key_value_heads, + self.sequence_length, + self.global_head_dim, + ) + past_kv_value = ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + past_kv_values.append(past_kv_value) + + return past_kv_values + + +class DeciDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + ) + self.num_key_value_heads_per_layer = normalized_config.num_key_value_heads_per_layer + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + + for layer_id in range(self.num_layers): + shape = ( + self.batch_size, + self.num_key_value_heads_per_layer[layer_id], + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + past_key_values.append( + ( + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(shape, framework=framework, dtype=float_dtype), + ) + ) + return past_key_values + + +class DummyLLavaMultiModalProjectorInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ["image_features"] + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + self.task = task + + self.batch_size = batch_size + self.hidden_size = normalized_config.hidden_size + self.num_patches = (normalized_config.image_size // normalized_config.patch_size) ** 2 + self.normalized_config = normalized_config + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + shape = [self.batch_size, self.num_patches, self.hidden_size] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class PooledProjectionsDummyInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ["pooled_projections"] + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + self.task = task + self.batch_size = batch_size + self.pooled_projection_dim = normalized_config.config.pooled_projection_dim + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + shape = [self.batch_size, self.pooled_projection_dim] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class DummyTransformerTimestpsInputGenerator(DummyTimestepInputGenerator): + SUPPORTED_INPUT_NAMES = ("timestep", "text_embeds", "time_ids", "timestep_cond", "guidance") + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["timestep", "guidance"]: + shape = [self.batch_size] + return self.random_float_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyUnetVisionInputGenerator(DummyVisionInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name not in ["sample", "latent_sample"]: + return super().generate(input_name, framework, int_dtype, float_dtype) + # add height and width discount for enable any resolution generation + return self.random_float_tensor( + shape=[self.batch_size, self.num_channels, self.height - 1, self.width - 1], + framework=framework, + dtype=float_dtype, + ) + + +class DummyUnetTimestepInputGenerator(DummyTimestepInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name != "timestep": + return super().generate(input_name, framework, int_dtype, float_dtype) + shape = [self.batch_size] + return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=int_dtype) + + +class DummySanaTimestepInputGenerator(DummyTimestepInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name != "timestep": + return super().generate(input_name, framework, int_dtype, float_dtype) + shape = [self.batch_size] + return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) + + +class DummyUnetEncoderInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + random_sequence_length_range: Optional[Tuple[int, int]] = None, + random_num_choices_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + super().__init__( + task, + normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + num_choices=num_choices, + random_batch_size_range=random_batch_size_range, + random_sequence_length_range=random_sequence_length_range, + random_num_choices_range=random_num_choices_range, + **kwargs, + ) + if hasattr(normalized_config.config, "model_max_length"): + self.sequence_length = normalized_config.config.model_max_length + + +class DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "decoder_input_ids", + "decoder_attention_mask", + "encoder_outputs", + "encoder_hidden_states", + "encoder_attention_mask", + ) + + +# DummySanaTransformerVisionInputGenerator inherits from DummyUnetVisionInputGenerator (defined above) +class DummySanaTransformerVisionInputGenerator(DummyUnetVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "pixel_values", + "pixel_mask", + "sample", + "latent_sample", + "guidance", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"] // 8, + height: int = DEFAULT_DUMMY_SHAPES["height"] // 8, + # Reduce img shape by 4 for FLUX to reduce memory usage on conversion + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width=width, height=height, **kwargs) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "guidance": + return self.random_float_tensor([self.batch_size], framework=framework, dtype=float_dtype) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyFluxTransformerInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "pixel_values", + "pixel_mask", + "sample", + "latent_sample", + "hidden_states", + "img_ids", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"] // 4, + height: int = DEFAULT_DUMMY_SHAPES["height"] // 4, + # Reduce img shape by 4 for FLUX to reduce memory usage on conversion + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + if getattr(normalized_config, "in_channels", None): + self.num_channels = normalized_config.in_channels // 4 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["hidden_states", "sample"]: + shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels * 4] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + if input_name == "img_ids": + img_ids_height = self.height // 2 + img_ids_width = self.width // 2 + return self.random_int_tensor( + ( + [self.batch_size, img_ids_height * img_ids_width, 3] + if is_diffusers_version("<", "0.31.0") + else [img_ids_height * img_ids_width, 3] + ), + min_value=0, + max_value=min(img_ids_height, img_ids_width), + framework=framework, + dtype=float_dtype, + ) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "decoder_input_ids", + "decoder_attention_mask", + "encoder_outputs", + "encoder_hidden_states", + "txt_ids", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "txt_ids": + import torch + + shape = ( + [self.batch_size, self.sequence_length, 3] + if is_diffusers_version("<", "0.31.0") + else [self.sequence_length, 3] + ) + dtype = DTYPE_MAPPER.pt(float_dtype) + return torch.full(shape, 0, dtype=dtype) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class LTXVaeDummyInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("pixel_values", "pixel_mask", "sample", "latent_sample", "timestep") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + num_frames: int = 2, + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + self.num_frames = num_frames + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["sample", "latent_sample"]: + return self.random_float_tensor( + [self.batch_size, self.num_channels, self.num_frames, self.height, self.width] + ) + if input_name == "timestep": + return self.random_int_tensor([1], max_value=20, min_value=1, framework=framework, dtype=int_dtype) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class LTXTransformerDummyInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("hidden_states", "width", "height", "num_frames", "rope_interpolation_scale") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = 16, + height: int = 8, + num_frames: int = 2, + frame_rate: int = 10, + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + self.num_frames = num_frames + self.frame_rate = frame_rate + self.vae_spatial_compression_ratio = normalized_config.config.vae_spatial_compression_ratio + self.vae_temporal_compression_ratio = normalized_config.config.vae_temporal_compression_ratio + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + import torch + + if input_name == "hidden_states": + return self.random_float_tensor( + [self.batch_size, self.num_frames * self.height * self.width, self.num_channels] + ) + if input_name == "width": + return torch.tensor(self.width) + if input_name == "height": + return torch.tensor(self.height) + if input_name == "num_frames": + return torch.tensor(self.num_frames) + if input_name == "rope_interpolation_scale": + import torch + + return torch.tensor( + [ + self.vae_temporal_compression_ratio / self.frame_rate, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ] + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyMiniCPMVImageInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("pixel_values", "patch_attention_mask", "position_ids") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height) + self.patch_size = normalized_config.config.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "pixel_values": + return self.random_float_tensor( + shape=[ + self.batch_size, + self.num_channels, + self.patch_size, + (self.height * self.width) // self.patch_size, + ], + framework=framework, + dtype=float_dtype, + ) + + if input_name == "patch_attention_mask": + return self.random_int_tensor( + shape=[self.batch_size, 1, (self.height // self.patch_size) * (self.width // self.patch_size)], + framework=framework, + dtype=float_dtype, + min_value=0, + max_value=2, + ) + + if input_name == "position_ids": + return self.random_int_tensor( + shape=[self.batch_size, (self.height // self.patch_size) * (self.width // self.patch_size)], + max_value=self.patch_size, + ) + + +class DummyMiniCPMVResampleInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("image_feature", "pos_embed", "key_padding_mask") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height) + self.patch_size = normalized_config.config.patch_size + self.hidden_size = normalized_config.config.hidden_size + self.img_hidden_size = normalized_config.config.vision_config.hidden_size + self.feat_size = (normalized_config.config.vision_config.image_size // self.patch_size) * ( + normalized_config.config.vision_config.image_size // self.patch_size + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "image_feature": + return self.random_float_tensor( + shape=[self.batch_size, self.feat_size, self.img_hidden_size], framework=framework, dtype=float_dtype + ) + + if input_name == "key_padding_mask": + return self.constant_tensor( + shape=[self.batch_size, self.feat_size], + framework=framework, + value=1, + dtype=DTYPE_MAPPER.pt(float_dtype), + ) + + if input_name == "pos_embed": + return self.random_float_tensor(shape=[self.feat_size, self.batch_size, self.hidden_size]) + + +class DummyPhi3VisionProjectionInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("input",) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = 336, + height: int = 336, + crop_size=336, + **kwargs, + ): + self.batch_size = batch_size + self._embed_layer_realization = ( + normalized_config.config.embd_layer["embedding_cls"] + if hasattr(normalized_config.config, "embd_layer") + else "image_audio" + ) + if not hasattr(normalized_config.config, "vision_config"): + self.image_dim_out = ( + normalized_config.config.img_processor.get( + "image_dim_out", normalized_config.config.img_processor.get("hidden_size") + ) + if normalized_config.config.img_processor is not None + else 1152 + ) + if "image_embd_layer" in normalized_config.config.embd_layer: + self.crop_size = normalized_config.config.embd_layer["image_embd_layer"].get("crop_size", crop_size) + else: + self.crop_size = normalized_config.config.embd_layer.get("crop_size", crop_size) + else: + self.image_dim_out = normalized_config.config.vision_config.hidden_size + self.crop_size = normalized_config.config.vision_config.crop_size + self.height = height + self.width = width + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + h = self.height // self.crop_size + w = self.width // self.crop_size + feat_size = (h * w + 1) * 144 + 1 + (h + 1) * 12 + if self._embed_layer_realization in ["linear", "image_audio"]: + shape = [self.batch_size, feat_size, self.image_dim_out] + else: + shape = [self.batch_size, feat_size, self.image_dim_out * 4] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class DummyAudioPhi4MMInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ("audio_input", "audio_feature", "audio_mask") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + signal_length=498, + **kwargs, + ): + self.signal_length = signal_length + if hasattr(normalized_config.config, "audio_processor"): + self.audio_chunk_size = ( + signal_length // normalized_config.config.audio_processor["config"]["time_reduction"] + 1 + ) + self.input_size = normalized_config.config.audio_processor["config"]["input_size"] + self.attention_dim = normalized_config.config.audio_processor["config"]["attention_dim"] + else: + self.audio_chunk_size = signal_length // normalized_config.config.audio_config.time_reduction + 1 + self.input_size = normalized_config.config.audio_config.input_size + self.attention_dim = normalized_config.config.audio_config.hidden_size + self.batch_size = batch_size + self.task = task + self.normalized_config = normalized_config + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "audio_input": + return self.random_float_tensor( + [self.batch_size, self.signal_length, self.input_size], framework=framework, dtype=float_dtype + ) + + if input_name == "audio_feature": + return self.random_float_tensor( + [self.batch_size, self.audio_chunk_size, self.attention_dim], framework=framework, dtype=float_dtype + ) + + if input_name == "audio_mask": + return self.random_int_tensor( + [self.batch_size, self.audio_chunk_size, self.audio_chunk_size], + max_value=2, + framework=framework, + dtype="bool", + ) + + +class DummyVisionPositionIdsPhi4InputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("patch_position_ids", "patch_attention_mask") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + if hasattr(normalized_config.config, "vision_conifg"): + self.patch_size = getattr(normalized_config.config.vision_config, "patch_size", 14) + else: + self.patch_size = 14 + self.num_patches_per_side = self.height // self.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "patch_position_ids": + return self.get_vision_position_ids() + if input_name == "patch_attention_mask": + return self.random_int_tensor( + [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], + framework=framework, + dtype="bool", + max_value=2, + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + def get_vision_position_ids(self): + # Adopted from https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py#L494-L512 + import torch + + batch_size = self.batch_size + max_im_h, max_im_w = self.height, self.width + max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size + boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side) + position_ids = torch.full( + size=( + batch_size, + max_nb_patches_h * max_nb_patches_w, + ), + fill_value=0, + ) + patch_attention_mask = torch.ones( + [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], dtype=torch.int64 + ) + patch_attention_mask[0, self.height - 2 :] = 0 + + for batch_idx, p_attn_mask in enumerate(patch_attention_mask): + nb_patches_h = p_attn_mask[:, 0].sum() + nb_patches_w = p_attn_mask[0].sum() + + fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h) + fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w) + + bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) + bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) + + pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten() + position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids + return position_ids + + +class DummyQwen2VLLMInputGenerator(DummyTextInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + generated_input = super().generate(input_name, framework, int_dtype, float_dtype) + if input_name == "position_ids": + return generated_input.unsqueeze(0).expand(3, -1, -1) + return generated_input + + +class DummyQwen3_5LMInputGenerator(DummyTextInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + generated_input = super().generate(input_name, framework, int_dtype, float_dtype) + if input_name == "position_ids": + return generated_input.unsqueeze(0).expand(4, -1, -1) + return generated_input + + +class DummyQwen2VLVisionEmbedInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "hidden_states", + "attention_mask", + "window_attention_mask", + "window_index", + "rotary_pos_emb", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = 1, + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = 420, + height: int = 420, + **kwargs, + ): + self.batch_size = batch_size + self.height = height + self.width = width + self.num_channels = num_channels + self.temporal_patch_size = normalized_config.config.temporal_patch_size + self.patch_size = normalized_config.config.patch_size + if normalized_config.use_embed_dim: + self.embed_dim = ( + normalized_config.config.embed_dim + if hasattr(normalized_config.config, "embed_dim") + else normalized_config.hidden_size + ) + else: + self.embed_dim = self.num_channels * self.temporal_patch_size * self.patch_size * self.patch_size + self.num_heads = normalized_config.config.num_heads + self.spatial_merge_size = None + if hasattr(normalized_config.config, "spatial_merge_size"): + self.spatial_merge_size = normalized_config.config.spatial_merge_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size + grid_t = self.batch_size + + if input_name == "hidden_states": + return self.random_float_tensor( + [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype + ) + + if input_name in ["attention_mask", "window_attention_mask"]: + return self.random_mask_tensor( + [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype + ) + + if input_name == "rotary_pos_emb": + dim = self.embed_dim // self.num_heads // 2 + return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) + + if input_name == "window_index": + if self.spatial_merge_size is None: + raise ValueError( + "`spatial_merge_size` parameter is not found in model config. Can not generate dummy input data for `window_index` input" + ) + spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + hidden_size = (grid_t * grid_h * grid_w) // spatial_merge_unit + return self.random_int_tensor([hidden_size], max_value=hidden_size) + + +# DummyQwen3VLVisionEmbedInputGenerator inherits from DummyQwen2VLVisionEmbedInputGenerator (defined above) +class DummyQwen3VLVisionEmbedInputGenerator(DummyQwen2VLVisionEmbedInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "hidden_states", + "attention_mask", + "rotary_pos_emb", + "input", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size + grid_t = self.batch_size + + if input_name == "hidden_states": + return self.random_float_tensor( + [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype + ) + + if input_name in ["attention_mask"]: + return self.random_mask_tensor( + [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype + ) + + if input_name == "rotary_pos_emb": + dim = self.embed_dim // self.num_heads // 2 + return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) + + if input_name == "input": + return self.constant_tensor( + [4, DEFAULT_DUMMY_SHAPES["sequence_length"]], + framework=framework, + value=0, + dtype=DTYPE_MAPPER.pt(int_dtype), + ) + + +class Qwen3ASRDummySeq2SeqPastKeyValuesGenerator(DummySeq2SeqPastKeyValuesGenerator): + """Custom KV cache generator for Qwen3-ASR with GQA (num_key_value_heads != num_attention_heads). + Qwen3-ASR has no cross-attention, so only self-attention KV cache is generated (2 per layer).""" + + def __init__(self, task, normalized_config, **kwargs): + super().__init__(task, normalized_config, **kwargs) + # Override head count and head_dim for GQA + self.decoder_num_attention_heads = normalized_config.decoder_num_attention_heads + self.decoder_head_dim = getattr(normalized_config, "head_dim", None) + if self.decoder_head_dim is None: + self.decoder_head_dim = self.decoder_hidden_size // normalized_config.num_attention_heads + + def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="fp32"): + if input_name == "past_key_values": + decoder_shape = ( + self.batch_size, + self.decoder_num_attention_heads, + self.sequence_length, + self.decoder_head_dim, + ) + # Qwen3-ASR has no cross-attention, only self-attention KV cache + return [ + ( + self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.decoder_num_layers) + ] + return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) + + +class DummyGemma4VisionInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("pixel_values", "image_position_ids") + + def __init__(self, task, normalized_config, batch_size=DEFAULT_DUMMY_SHAPES["batch_size"], **kwargs): + super().__init__(task, normalized_config, batch_size, **kwargs) + self.patch_size = getattr(normalized_config, "patch_size", 16) + self.pooling_kernel_size = getattr(normalized_config, "pooling_kernel_size", 3) + # Gemma4 processor always pads pixel_values to max_soft_tokens * pooling_kernel_size^2 patches. + # The vision model's pooling uses shape-dependent Python operations that get baked in during tracing, + # so the dummy input must match the actual inference shapes. + max_soft_tokens = getattr(normalized_config, "image_seq_length", None) + if max_soft_tokens is None: + max_soft_tokens = getattr(normalized_config, "max_soft_tokens", 280) + self.num_patches = max_soft_tokens * self.pooling_kernel_size**2 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "pixel_values": + # Gemma4 expects pre-patchified pixel_values: [batch, num_patches, 3 * patch_size^2] + return self.random_float_tensor( + shape=[self.batch_size, self.num_patches, 3 * self.patch_size**2], + framework=framework, + dtype=float_dtype, + ) + if input_name == "image_position_ids": + # Create position ids as a grid. The patch count = h_patches * w_patches + # where both are divisible by pooling_kernel_size for correct pooling. + import math + + k = self.pooling_kernel_size + total_pooled = self.num_patches // (k * k) + # Find roughly square grid for pooled side + pooled_side = int(math.sqrt(total_pooled)) + if pooled_side * pooled_side < total_pooled: + pooled_h = pooled_side + pooled_w = total_pooled // pooled_h + else: + pooled_h = pooled_w = pooled_side + h_patches = pooled_h * k + w_patches = pooled_w * k + pos_ids = torch.stack( + torch.meshgrid(torch.arange(h_patches), torch.arange(w_patches), indexing="ij"), dim=-1 + ).reshape(1, -1, 2) + # Pad to num_patches with -1 (padding position) + if pos_ids.shape[1] < self.num_patches: + pad = torch.full((1, self.num_patches - pos_ids.shape[1], 2), -1, dtype=pos_ids.dtype) + pos_ids = torch.cat([pos_ids, pad], dim=1) + return pos_ids.expand(self.batch_size, -1, -1).clone() + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyVisionPositionIdsInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("patch_attention_mask", "patch_position_ids") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) + self.patch_size = normalized_config.config.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "patch_attention_mask": + shape = [self.batch_size, self.height // self.patch_size, self.width // self.patch_size] + return self.random_int_tensor(shape, max_value=2, framework=framework, dtype="bool") + if input_name == "patch_position_ids": + max_nb_patches_h, max_nb_patches_w = self.height // self.patch_size, self.width // self.patch_size + shape = [self.batch_size, max_nb_patches_h * max_nb_patches_w] + return self.random_int_tensor( + shape, max_value=min(max_nb_patches_h, max_nb_patches_w), framework=framework, dtype=int_dtype + ) + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummySpeechT5InputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "inputs_embeds", + "output_sequence", + "speaker_embeddings", + "spectrogram", + "raw_spectrogram", + "encoder_hidden_states", + ) + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.task = task + self.batch_size = 1 + + self.sequence_length = sequence_length + self.speaker_embedding_dim = normalized_config.speaker_embedding_dim + self.num_mel_bins = normalized_config.num_mel_bins + self.reduction_factor = normalized_config.config.reduction_factor + self.hidden_size = normalized_config.config.hidden_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name in ["output_sequence", "inputs_embeds"]: + shape = [self.batch_size, self.sequence_length, self.num_mel_bins] + elif input_name == "speaker_embeddings": + shape = [self.batch_size, self.speaker_embedding_dim] + elif input_name == "raw_spectrogram": + shape = [self.sequence_length, self.batch_size, self.reduction_factor, self.num_mel_bins] + elif input_name == "encoder_hidden_states": + shape = [self.batch_size, self.sequence_length, self.hidden_size] + elif input_name == "spectrogram": + shape = [self.batch_size, self.sequence_length, self.num_mel_bins] + else: + raise ValueError(f"Unsupported input {input_name} for DummySpeechT5InputGenerator") + + return self.random_float_tensor( + shape=shape, + min_value=0, + max_value=1, + framework=framework, + dtype=float_dtype, + ) + + +class MambaCacheDummyInputGenerator(DummyInputGenerator): + """ + Generates dummy past_ssm_states, past_conv_states and cache_position inputs for Mamba architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params", "cache_position") + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.normalized_config = normalized_config + self.batch_size = batch_size + self.sequence_length = sequence_length + self.intermediate_size = self.normalized_config.config.intermediate_size + self.ssm_state_size = self.normalized_config.config.state_size + self.conv_kernel_size = self.normalized_config.config.conv_kernel + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "cache_params": + ssm_shape = [self.batch_size, self.intermediate_size, self.ssm_state_size] + conv_shape = [self.batch_size, self.intermediate_size, self.conv_kernel_size] + return [ + ( + self.random_float_tensor(ssm_shape, framework=framework, dtype=float_dtype), + self.random_float_tensor(conv_shape, framework=framework, dtype=float_dtype), + ) + for _ in range(self.normalized_config.num_layers) + ] + elif input_name == "cache_position": + return self.random_int_tensor( + shape=[self.conv_kernel_size], + max_value=self.sequence_length, + framework=framework, + dtype=int_dtype, + ) + + raise ValueError(f"Unsupported input name {input_name}") + + +class Zamba2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for Zamba2 architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + + config = normalized_config.config + self.intermediate_size = int(config.mamba_expand * config.hidden_size) + self.conv_kernel_size = config.mamba_d_conv + self.mamba_d_state = config.mamba_d_state + if config.model_type == "zamba2": + self.n_mamba_heads = config.n_mamba_heads + self.mamba_ngroups = config.mamba_ngroups + self.mamba_headdim = config.mamba_headdim + self.head_dim = config.attention_head_dim + # in Zamba2, all layers contain Mamba block + # some of these layers are hybrid so they contain both attention and mamba blocks + self.num_attention_layers = len(config.hybrid_layer_ids) + self.num_mamba_layers = self.num_layers + import logging + + logger = logging.getLogger(__name__) + logger.warning( + "The current support for the 'Zamba2' model type is experimental. " + "Performance is not optimal with high memory consumption. " + "Optimizations and improved support will be available in a future OpenVINO release." + ) + else: + # currently, this else-branch is applied for GraniteMoeHybrid models + self.n_mamba_heads = config.mamba_n_heads + self.mamba_ngroups = config.mamba_n_groups + self.mamba_headdim = config.mamba_d_head + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_attention_layers = config.layer_types.count("attention") + self.num_mamba_layers = config.layer_types.count("mamba") + self.num_attention_heads = config.num_key_value_heads + self.sequence_length = 0 + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + for i in range(self.num_mamba_layers): + conv_state_shape = ( + self.batch_size, + self.intermediate_size + 2 * self.mamba_ngroups * self.mamba_d_state, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(conv_state) + ssm_state_shape = (self.batch_size, self.n_mamba_heads, self.mamba_headdim, self.mamba_d_state) + ssm_state = self.random_float_tensor(ssm_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(ssm_state) + + for i in range(self.num_attention_layers): + kv_shape = (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + past_key_values.append(k) + past_key_values.append(v) + + return past_key_values + + +class Lfm2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy past_key_values inputs for Lfm2 architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + config = normalized_config.config + self.num_conv_layers = config.layer_types.count("conv") + self.num_atten_layers = config.layer_types.count("full_attention") + self.batch_size = batch_size + self.normalized_config = normalized_config + self.hidden_size = self.normalized_config.hidden_size + self.conv_L_cache = self.normalized_config.conv_L_cache + self.num_key_value_heads = self.normalized_config.num_key_value_heads + self.num_hidden_layers = self.normalized_config.num_hidden_layers + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + past_key_values = [] + + for i in range(self.num_conv_layers): + conv_state_shape = (self.batch_size, self.hidden_size, self.conv_L_cache) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + past_key_values.append(conv_state) + + for i in range(self.num_atten_layers): + shape = ( + self.batch_size, + self.num_key_value_heads, + self.sequence_length, + self.hidden_size // self.num_attention_heads, + ) + + kv_shape = shape # (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + past_key_values.append(k) + past_key_values.append(v) + + return past_key_values + + +class DummyVideoChatFlashQwenInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("hidden_states", "rotary_pos_emb") + + def __init__( + self, + task: str, + normalized_config: NormalizedVisionConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], + width: int = DEFAULT_DUMMY_SHAPES["width"], + height: int = DEFAULT_DUMMY_SHAPES["height"], + visual_seq_length: int = DEFAULT_DUMMY_SHAPES["visual_seq_length"], + **kwargs, + ): + super().__init__(task, normalized_config, batch_size, num_channels, width, height, visual_seq_length, **kwargs) + self.num_frames = normalized_config.config.mm_local_num_frames + self.embed_dim = normalized_config.config.mm_hidden_size + self.height = normalized_config.config.image_size + self.width = normalized_config.config.image_size + self.image_size = (self.height, self.width) + self.patch_size = normalized_config.config.patch_size + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "hidden_states": + return self.random_float_tensor( + shape=[ + self.batch_size, + self.num_channels, + self.num_frames, + self.height, + self.width, + ], + framework=framework, + dtype=float_dtype, + ) + elif input_name == "rotary_pos_emb": + grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size + grid_t = self.num_frames + # Source: https://huggingface.co/OpenGVLab/VideoChat-Flash-Qwen2_5-7B_InternVideo2-1B/blob/main/vision_tower_builder.py#L523 + # The first dimension of rotary_pos_emb is fixed to 1 in the original model. + # And the second dimension is the total number of tokens for all frames, which is calculated as grid_h * grid_w * grid_t plus 1 for the cls token. + return self.random_float_tensor( + [1, 1 + grid_h * grid_t * grid_w, self.embed_dim], framework=framework, dtype=float_dtype + ) + + +class DummyVideoChatFlashQwenProjectorInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ["input"] + + def __init__( + self, + task: str, + normalized_config: NormalizedTextConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + random_batch_size_range: Optional[Tuple[int, int]] = None, + **kwargs, + ): + self.task = task + self.batch_size = batch_size + self.hidden_size = normalized_config.config.mm_hidden_size + self.num_patches = normalized_config.config.mm_projector_num_tome_tokens + self.normalized_config = normalized_config + + def generate( + self, + input_name: str, + framework: str = "pt", + int_dtype: str = "int64", + float_dtype: str = "fp32", + ): + shape = [self.batch_size, self.num_patches, self.hidden_size] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + + +class Qwen3NextDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for Qwen3-Next architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + + config = normalized_config.config + self.num_full_attn_layers = config.layer_types.count("full_attention") + self.num_linear_attn_layers = config.layer_types.count("linear_attention") + self.conv_kernel_size = config.linear_conv_kernel_dim + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.num_key_value_heads = config.num_key_value_heads + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + cache_params = [] + + for idx in range(self.num_linear_attn_layers): + d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) + conv_state_shape = ( + self.batch_size, + d_inner, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(conv_state) + num_heads = self.num_v_heads + recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) + recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(recurrent_state) + + for idx in range(self.num_full_attn_layers): + kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + cache_params.append(k) + cache_params.append(v) + + return cache_params + + +class Qwen3_5DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): + """ + Generates dummy cache_params inputs for Qwen3.5 architectures. + """ + + SUPPORTED_INPUT_NAMES = ("cache_params",) + + def __init__( + self, + task: str, + normalized_config, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + super().__init__( + task=task, + normalized_config=normalized_config, + batch_size=batch_size, + sequence_length=sequence_length, + **kwargs, + ) + + config = normalized_config.config + self.num_full_attn_layers = config.layer_types.count("full_attention") + self.num_linear_attn_layers = config.layer_types.count("linear_attention") + self.conv_kernel_size = config.linear_conv_kernel_dim + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.num_key_value_heads = config.num_key_value_heads + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + cache_params = [] + + for idx in range(self.num_linear_attn_layers): + d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) + conv_state_shape = ( + self.batch_size, + d_inner, + self.conv_kernel_size, + ) + conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(conv_state) + num_heads = self.num_v_heads + recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) + recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) + cache_params.append(recurrent_state) + + for idx in range(self.num_full_attn_layers): + kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) + k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) + cache_params.append(k) + cache_params.append(v) + + return cache_params + + +class DummyKokoroInputGenerator(DummyInputGenerator): + """Generates dummy inputs for the Kokoro TTS model.""" + + SUPPORTED_INPUT_NAMES = ("input_ids", "ref_s", "speed") + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], + **kwargs, + ): + self.task = task + self.batch_size = 1 + self.sequence_length = sequence_length + self.style_dim = getattr(normalized_config, "style_dim", 128) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "input_ids": + shape = [self.batch_size, self.sequence_length] + input_ids_value = self.random_int_tensor( + shape=shape, min_value=0, max_value=178, framework=framework, dtype=int_dtype + ) + input_ids_value[:, 0] = 0 + input_ids_value[:, -1] = 0 + return input_ids_value + elif input_name == "ref_s": + shape = [self.batch_size, self.style_dim * 2] + return self.random_float_tensor( + shape=shape, min_value=-1, max_value=1, framework=framework, dtype=float_dtype + ) + elif input_name == "speed": + return self.random_int_tensor(shape=[1], min_value=1, max_value=10, framework=framework, dtype=float_dtype) + else: + raise ValueError(f"Unsupported input {input_name} for DummyKokoroInputGenerator") diff --git a/optimum/exporters/openvino/model_configs.py b/optimum/exporters/openvino/model_configs.py index 19d954b62a..32ee9ed288 100644 --- a/optimum/exporters/openvino/model_configs.py +++ b/optimum/exporters/openvino/model_configs.py @@ -14,136 +14,75 @@ import enum import logging -import math -from copy import deepcopy -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Union -import torch from transformers import AutoConfig, PretrainedConfig, PreTrainedModel -from optimum.exporters.onnx.config import ( - AudioToTextOnnxConfig, - OnnxConfig, - TextDecoderOnnxConfig, - TextDecoderWithPositionIdsOnnxConfig, +from optimum.exporters.openvino.base import ( + ConfigBehavior, + OpenVINOConfig, + OpenVINOConfigWithPast, + OpenVINOSeq2SeqConfigWithPast, ) -from optimum.exporters.onnx.model_configs import ( - AlbertOnnxConfig, - ASTOnnxConfig, - BartOnnxConfig, - BeitOnnxConfig, - BertOnnxConfig, - BlenderbotOnnxConfig, - BlenderbotSmallOnnxConfig, - BloomOnnxConfig, - CamembertOnnxConfig, - CLIPOnnxConfig, - CLIPTextOnnxConfig, - CLIPTextWithProjectionOnnxConfig, - CLIPVisionModelOnnxConfig, - CodeGenOnnxConfig, - ConvBertOnnxConfig, - ConvNextOnnxConfig, - Data2VecAudioOnnxConfig, - Data2VecTextOnnxConfig, - Data2VecVisionOnnxConfig, - DebertaOnnxConfig, - DebertaV2OnnxConfig, - DeiTOnnxConfig, - DistilBertOnnxConfig, - ElectraOnnxConfig, - EsmOnnxConfig, - FalconOnnxConfig, - FlaubertOnnxConfig, - GemmaOnnxConfig, - GPT2OnnxConfig, - GPTBigCodeOnnxConfig, - GPTJOnnxConfig, - GPTNeoOnnxConfig, - GPTNeoXOnnxConfig, - HubertOnnxConfig, - IBertOnnxConfig, - LevitOnnxConfig, - LlamaOnnxConfig, - MarianOnnxConfig, - MistralOnnxConfig, - MobileBertOnnxConfig, - MobileNetV1OnnxConfig, - MobileNetV2OnnxConfig, - MobileViTOnnxConfig, - MPNetOnnxConfig, - MPTOnnxConfig, - NystromformerOnnxConfig, - Olmo2OnnxConfig, - OPTOnnxConfig, - PegasusOnnxConfig, - PerceiverOnnxConfig, - PhiOnnxConfig, - Pix2StructOnnxConfig, - PoolFormerOnnxConfig, - RemBertOnnxConfig, - ResNetOnnxConfig, - RobertaOnnxConfig, - RoFormerOnnxConfig, - SamOnnxConfig, - SegformerOnnxConfig, - SentenceTransformersTransformerOnnxConfig, - SEWDOnnxConfig, - SEWOnnxConfig, - SiglipOnnxConfig, - SiglipTextOnnxConfig, - SiglipTextWithProjectionOnnxConfig, - SpeechT5OnnxConfig, - SqueezeBertOnnxConfig, - SwinOnnxConfig, - T5OnnxConfig, - UNetOnnxConfig, - UniSpeechOnnxConfig, - UniSpeechSATOnnxConfig, - VaeDecoderOnnxConfig, - VaeEncoderOnnxConfig, - VisionEncoderDecoderOnnxConfig, - VisionOnnxConfig, - ViTOnnxConfig, - Wav2Vec2ConformerOnnxConfig, - Wav2Vec2OnnxConfig, - WavLMOnnxConfig, - WhisperOnnxConfig, - XLMOnnxConfig, - XLMRobertaOnnxConfig, +from optimum.exporters.openvino.config import ( + AudioOpenVINOConfig, + AudioToTextOpenVINOConfig, + EncoderDecoderBaseOpenVINOConfig, + TextAndVisionOpenVINOConfig, + TextDecoderOpenVINOConfig, + TextDecoderWithPositionIdsOpenVINOConfig, + TextEncoderOpenVINOConfig, + TextSeq2SeqOpenVINOConfig, + VisionOpenVINOConfig, ) -from optimum.exporters.onnx.model_patcher import ModelPatcher -from optimum.exporters.openvino.utils import ONNX_SUPPORTED_ARCHITECTURES -from optimum.exporters.tasks import TasksManager -from optimum.utils import DEFAULT_DUMMY_SHAPES -from optimum.utils.input_generators import ( - DTYPE_MAPPER, - DummyAudioInputGenerator, - DummyInputGenerator, - DummyPastKeyValuesGenerator, - DummySeq2SeqDecoderTextInputGenerator, - DummySeq2SeqPastKeyValuesGenerator, - DummyTextInputGenerator, - DummyTimestepInputGenerator, - DummyVisionInputGenerator, - FalconDummyPastKeyValuesGenerator, - GemmaDummyPastKeyValuesGenerator, - MistralDummyPastKeyValuesGenerator, -) -from optimum.utils.normalized_config import ( - NormalizedConfig, - NormalizedSeq2SeqConfig, - NormalizedTextConfig, - NormalizedVisionConfig, -) - -from ...intel.utils.import_utils import ( - is_diffusers_available, - is_diffusers_version, - is_openvino_version, - is_transformers_version, +from optimum.exporters.openvino.input_generators import ( + AquilaDummyPastKeyValuesGenerator, + ChatGLM2DummyPastKeyValuesGenerator, + DeciDummyPastKeyValuesGenerator, + DummyAudioPhi4MMInputGenerator, + DummyFluxTextInputGenerator, + DummyFluxTransformerInputGenerator, + DummyGemma4VisionInputGenerator, + DummyKokoroInputGenerator, + DummyLLavaMultiModalProjectorInputGenerator, + DummyMiniCPMVImageInputGenerator, + DummyMiniCPMVResampleInputGenerator, + DummyPhi3VisionProjectionInputGenerator, + DummyQwen2VLLMInputGenerator, + DummyQwen2VLVisionEmbedInputGenerator, + DummyQwen3_5LMInputGenerator, + DummyQwen3VLLMInputGenerator, + DummyQwen3VLVisionEmbedInputGenerator, + DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator, + DummySanaTimestepInputGenerator, + DummySanaTransformerVisionInputGenerator, + DummySpeechT5InputGenerator, + DummyTransformerTimestpsInputGenerator, + DummyUnetEncoderInputGenerator, + DummyUnetTimestepInputGenerator, + DummyUnetVisionInputGenerator, + DummyVideoChatFlashQwenInputGenerator, + DummyVideoChatFlashQwenProjectorInputGenerator, + DummyVisionPositionIdsInputGenerator, + DummyVisionPositionIdsPhi4InputGenerator, + Eagle3DummyGenerator, + Eagle3VLMDummyGenerator, + Gemma4DummyPastKeyValuesGenerator, + GPTBigCodeDummyPastKeyValuesGenerator, + Lfm2DummyPastKeyValuesGenerator, + LTXTransformerDummyInputGenerator, + LTXVaeDummyInputGenerator, + MambaCacheDummyInputGenerator, + OVFalconDummyPastKeyValuesGenerator, + OVMiniCPM3DummyPastKeyValuesGenerator, + PooledProjectionsDummyInputGenerator, + Qwen3_5DummyPastKeyValuesGenerator, + Qwen3ASRDummySeq2SeqPastKeyValuesGenerator, + Qwen3NextDummyPastKeyValuesGenerator, + QwenDummyPastKeyValuesGenerator, + Zamba2DummyPastKeyValuesGenerator, ) -from .model_patcher import ( +from optimum.exporters.openvino.model_patcher import ( AfmoeModelPatcher, AquilaModelPatcher, ArcticModelPatcher, @@ -153,13 +92,14 @@ BlenderbotSmallModelPatcher, BloomModelPatcher, ChatGLMModelPatcher, + CLIPModelPatcher, CodeGenModelPatcher, CommonImageEmbeddingsModelPatcher, DBRXModelPatcher, DeciLMModelPatcher, DeepseekPatcher, FalconModelPatcher, - FluxTransfromerModelPatcher, + FluxTransformerModelPatcher, Gemma2ModelPatcher, Gemma3LMModelPatcher, Gemma4ImageEmbeddingsModelPatcher, @@ -195,10 +135,10 @@ MiniCPMVResamplerModelPatcher, MistralModelPatcher, MixtralModelPatcher, + ModelPatcher, MPTModelPatcher, OVDecoderModelPatcher, OVSeq2SeqModelPatcher, - OVSpeechT5ModelPatcher, PegasusModelPatcher, PersimmonModelPatcher, Phi3ModelPatcher, @@ -221,12 +161,52 @@ Qwen3VLLanguageModelPatcher, Qwen3VLVisionEmbMergerPatcher, QwenModelPatcher, + SAMModelPatcher, SanaTextEncoderModelPatcher, + SentenceTransformersTransformerPatcher, + SpeechT5ModelPatcher, VideoChatFlashQwenVisionEmbeddingModelPatcher, XverseModelPatcher, Zamba2ModelPatcher, _get_model_attribute, ) +from optimum.exporters.tasks import TasksManager +from optimum.intel.utils.import_utils import ( + is_diffusers_available, + is_diffusers_version, + is_openvino_version, + is_transformers_version, +) +from optimum.utils.input_generators import ( + ASTDummyAudioInputGenerator, + BartDummyTextInputGenerator, + BloomDummyPastKeyValuesGenerator, + DummyAudioInputGenerator, + DummyDecoderTextInputGenerator, + DummyInputGenerator, + DummyPastKeyValuesGenerator, + DummyPix2StructInputGenerator, + DummyPointsGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyTextInputGenerator, + DummyVisionEmbeddingsGenerator, + DummyVisionEncoderDecoderPastKeyValuesGenerator, + DummyVisionInputGenerator, + GemmaDummyPastKeyValuesGenerator, + MistralDummyPastKeyValuesGenerator, + PerceiverDummyInputGenerator, + T5DummySeq2SeqPastKeyValuesGenerator, +) +from optimum.utils.normalized_config import ( + NormalizedConfig, + NormalizedConfigManager, + NormalizedEncoderDecoderConfig, + NormalizedSeq2SeqConfig, + NormalizedTextAndVisionConfig, + NormalizedTextConfig, + NormalizedVisionConfig, +) COMMON_TEXT_TASKS = [ @@ -265,10 +245,6 @@ def _warn_potential_accuracy_issue_ov_2026_1(model_type: str, min_transformers_v logger.warning(f"Model type '{model_type}' may have potential accuracy issues with OpenVINO >= 2026.1.0.") -if TYPE_CHECKING: - from transformers.modeling_utils import PreTrainedModel # noqa: F811 - - def init_model_configs(): if "open_clip" not in TasksManager._LIBRARY_TO_SUPPORTED_MODEL_TYPES: TasksManager._LIBRARY_TO_SUPPORTED_MODEL_TYPES["open_clip"] = {} @@ -357,24 +333,6 @@ def init_model_configs(): TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"] = {} TasksManager._DIFFUSERS_TASKS_TO_MODEL_MAPPINGS["text-to-video"]["ltx-video"] = "LTXPipeline" - supported_model_types = [ - "_SUPPORTED_MODEL_TYPE", - "_DIFFUSERS_SUPPORTED_MODEL_TYPE", - "_TIMM_SUPPORTED_MODEL_TYPE", - "_SENTENCE_TRANSFORMERS_SUPPORTED_MODEL_TYPE", - ] - # TODO: remove once models from ONNX_SUPPORTED_ARCHITECTURES are deprecated (optimum-intel v1.29) - for supported_models_config in supported_model_types: - supported_models = getattr(TasksManager, supported_models_config) - for model, export_configs in supported_models.items(): - # adding only the architectures that are already supported via optimum-onnx v0.1.0 - if "onnx" not in export_configs or model not in ONNX_SUPPORTED_ARCHITECTURES: - continue - onnx_config = export_configs["onnx"] - supported_models[model]["openvino"] = deepcopy(onnx_config) - - setattr(TasksManager, supported_models_config, supported_models) - init_model_configs() @@ -383,8 +341,7 @@ def init_model_configs(): @register_in_tasks_manager("baichuan", *["text-generation", "text-generation-with-past"], library_name="transformers") -class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 13 +class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_layers="num_hidden_layers", num_attention_heads="num_attention_heads", hidden_size="hidden_size" ) @@ -404,8 +361,7 @@ class BaichaunOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -413,8 +369,7 @@ class Qwen2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("qwen2_moe", *["text-generation", "text-generation-with-past"], library_name="transformers") -class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -432,7 +387,7 @@ class Qwen2MoEOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class Qwen3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class Qwen3OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.51.0" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator @@ -451,66 +406,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return common_inputs -class DummyQwen3VLLMInputGenerator(DummyTextInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "input_ids", - "attention_mask", - "token_type_ids", - "position_ids", - "visual_pos_masks", - "deepstack_visual_embeds", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - random_num_choices_range: Optional[Tuple[int, int]] = None, - padding_side: str = "right", - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - num_choices=num_choices, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - random_num_choices_range=random_num_choices_range, - padding_side=padding_side, - **kwargs, - ) - self.embed_dim = normalized_config.hidden_size - self.num_layers = len(self.normalized_config.deepstack_visual_indexes) - - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - bool_dtype: str = "bool", - ): - if input_name == "deepstack_visual_embeds": - return self.random_float_tensor( - [self.num_layers, 2 * self.sequence_length, self.embed_dim], framework=framework, dtype=float_dtype - ) - if input_name == "visual_pos_masks": - return self.constant_tensor( - shape=[self.batch_size, self.sequence_length], - framework=framework, - value=1, - dtype=DTYPE_MAPPER.pt(bool_dtype), - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager( "qwen3_vl_text", *[ @@ -519,7 +414,7 @@ def generate( ], library_name="transformers", ) -class Qwen3VLTextOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class Qwen3VLTextOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.57.0" DUMMY_INPUT_GENERATOR_CLASSES = (DummyQwen3VLLMInputGenerator, GemmaDummyPastKeyValuesGenerator) @@ -545,8 +440,7 @@ class Qwen3MoEOpenVINOConfig(Qwen3OpenVINOConfig): @register_in_tasks_manager("minicpm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.53.3" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator @@ -554,49 +448,8 @@ class MiniCPMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): _MODEL_PATCHER = MiniCPMModelPatcher -class OVMiniCPM3DummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - self.v_head_dim = getattr(normalized_config, "v_head_dim", self.hidden_size // self.num_attention_heads) - self.k_head_dim = normalized_config.qk_nope_head_dim + normalized_config.qk_rope_head_dim - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - v_shape = ( - self.batch_size, - self.num_key_value_heads, - self.sequence_length, - self.v_head_dim, - ) - k_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.k_head_dim) - return [ - ( - self.random_float_tensor(k_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(v_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("minicpm3", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.53.3" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, OVMiniCPM3DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = OVMiniCPM3DummyPastKeyValuesGenerator @@ -605,63 +458,15 @@ class MiniCPM3OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("stablelm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class StableLMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = OVDecoderModelPatcher -class ChatGLM2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.multi_query_group_num = normalized_config.multi_query_group_num - self.head_dim = normalized_config.kv_channels - self.standart_cache_layout = hasattr(normalized_config, "rope_ratio") - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if not self.standart_cache_layout: - pkv_shape = ( - self.sequence_length, - self.batch_size, - self.multi_query_group_num, - self.head_dim, - ) - else: - pkv_shape = ( - self.batch_size, - self.multi_query_group_num, - self.sequence_length, - self.head_dim, - ) - return [ - ( - self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(pkv_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("chatglm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): +class ChatGLM2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(vocab_size="padded_vocab_size", num_layers="num_layers") DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, ChatGLM2DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = ChatGLM2DummyPastKeyValuesGenerator @@ -690,7 +495,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) # refer to https://github.com/huggingface/optimum/pull/764 @@ -750,12 +555,10 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire @register_in_tasks_manager("mixtral", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - # The ONNX export of this architecture needs the Trilu operator support, available since opset 14 - DEFAULT_ONNX_OPSET = 14 +class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = ( MistralDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) _MODEL_PATCHER = MixtralModelPatcher @@ -772,12 +575,16 @@ class MixtralOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): ], library_name="transformers", ) -class GemmaOpenVINOConfig(GemmaOnnxConfig): +class GemmaOpenVINOConfig(TextDecoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator + MIN_TRANSFORMERS_VERSION = "4.38.0" _MODEL_PATCHER = OVDecoderModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: - # position_ids was removed from optimum-onnx's gemma config because + # position_ids was removed from optimum-openvino's gemma config because # it's not necessary (it's correctly generated inside the model) # but openvino genai requires it to be present to work properly inputs = super().inputs @@ -786,73 +593,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return inputs -class Eagle3DummyGenerator(DummyInputGenerator): - """ - Dummy input generator for Eagle-3 speculative decoding. - - This generator produces synthetic `hidden_states` tensors that mimic the - intermediate hidden-state outputs of a *main (target) model*, which are - required by the Eagle-3 draft model during speculative decoding. - """ - - SUPPORTED_INPUT_NAMES = ("hidden_states",) - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.batch_size = batch_size - self.sequence_length = sequence_length - self.hidden_size = normalized_config.hidden_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - # hidden_states is provided as a concatenation of three hidden-layer outputs from the main model - shape = ( - self.batch_size, - self.sequence_length, - self.hidden_size * 3, - ) - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - -class Eagle3VLMDummyGenerator(DummyInputGenerator): - """ - Dummy input generator for VLM Eagle-3 speculative decoding. - - Produces `inputs_embeds` (float) and 3D `position_ids` (MRoPE) - required by VLM Eagle-3 draft models targeting Qwen3-VL. - """ - - SUPPORTED_INPUT_NAMES = ("inputs_embeds", "position_ids") - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.batch_size = batch_size - self.sequence_length = sequence_length - self.hidden_size = normalized_config.hidden_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "inputs_embeds": - shape = (self.batch_size, self.sequence_length, self.hidden_size) - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - if input_name == "position_ids": - # The rotary embedding is MRoPE (Multimodal RoPE) - # MRoPE encodes position along three independent axes: temporal, height, and width - # https://github.com/Tencent/AngelSlim/blob/main/angelslim/compressor/speculative/train/models/draft/llama_eagle3.py#L211 - shape = (3, self.batch_size, self.sequence_length) - return self.random_int_tensor(shape, max_value=self.sequence_length, framework=framework, dtype=int_dtype) - - @register_in_tasks_manager( "llama", *[ @@ -865,7 +605,10 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int ], library_name="transformers", ) -class LlamaOpenVINOConfig(LlamaOnnxConfig): +class LlamaOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = OVDecoderModelPatcher def __init__( @@ -962,7 +705,10 @@ class GptOssOpenVINOConfig(LlamaOpenVINOConfig): ], library_name="transformers", ) -class BitnetOpenVINOConfig(LlamaOnnxConfig): +class BitnetOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig MIN_TRANSFORMERS_VERSION = "4.52.1" MAX_TRANSFORMERS_VERSION = "4.57.6" _MODEL_PATCHER = OVDecoderModelPatcher @@ -1027,42 +773,8 @@ class Cohere2OpenVINOConfig(LlamaOpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.48.0" -class QwenDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.kv_channels = normalized_config.kv_channels - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) - past_value_shape = (self.batch_size, self.sequence_length, self.num_attention_heads, self.kv_channels) - return [ - ( - self.random_float_tensor(past_key_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(past_value_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("qwen", *["text-generation", "text-generation-with-past"]) -class QwenOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class QwenOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.55.4" NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_layers="num_hidden_layers", num_attention_heads="num_attention_heads", hidden_size="hidden_size" @@ -1093,7 +805,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) # refer to https://github.com/huggingface/optimum/pull/764 @@ -1143,8 +855,7 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire @register_in_tasks_manager( "starcoder2", *["text-generation", "text-generation-with-past"], library_name="transformers" ) -class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -1152,8 +863,7 @@ class Starcoder2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("internlm2", *["text-generation", "text-generation-with-past"], library_name="transformers") -class InternLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class InternLM2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator @@ -1162,8 +872,7 @@ class InternLM2OpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("orion", *["text-generation", "text-generation-with-past"], library_name="transformers") -class OrionOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class OrionOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator @@ -1173,14 +882,16 @@ class OrionOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("olmo", *["text-generation", "text-generation-with-past"], library_name="transformers") class OlmoOpenVINOConfig(LlamaOpenVINOConfig): - DEFAULT_ONNX_OPSET = 14 NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @register_in_tasks_manager( "mpt", *["text-generation", "text-generation-with-past", "text-classification"], library_name="transformers" ) -class MPTOpenVINOConfig(MPTOnnxConfig): +class MPTOpenVINOConfig(TextDecoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( + num_attention_heads="n_heads", hidden_size="d_model", num_layers="n_layers" + ) _MODEL_PATCHER = MPTModelPatcher @@ -1195,12 +906,13 @@ class MPTOpenVINOConfig(MPTOnnxConfig): ], library_name="transformers", ) -class Phi3OpenVINOConfig(PhiOnnxConfig): +class Phi3OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( MistralDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator - NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) + MIN_TRANSFORMERS_VERSION = "4.36.0" _MODEL_PATCHER = Phi3ModelPatcher @@ -1231,38 +943,12 @@ class PhiMoEOpenVINOConfig(Phi3OpenVINOConfig): ], library_name="transformers", ) -class PhiOpenVINOConfig(PhiOnnxConfig): +class PhiOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = "4.36.0" _MODEL_PATCHER = OVDecoderModelPatcher -class OVFalconDummyPastKeyValuesGenerator(FalconDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - if normalized_config.new_decoder_architecture: - self.num_kv_heads = normalized_config.num_attention_heads - else: - self.num_kv_heads = normalized_config.num_kv_heads if not normalized_config.multi_query else 1 - - self.head_dim = self.hidden_size // self.num_attention_heads - - @register_in_tasks_manager( "falcon", *[ @@ -1275,13 +961,23 @@ def __init__( ], library_name="transformers", ) -class FalconOpenVINOConfig(FalconOnnxConfig): +class FalconOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = ( OVFalconDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES DUMMY_PKV_GENERATOR_CLASS = OVFalconDummyPastKeyValuesGenerator _MODEL_PATCHER = FalconModelPatcher + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = super().inputs + + if self._config.alibi: + common_inputs.pop("position_ids", None) + + return common_inputs + @register_in_tasks_manager( "persimmon", @@ -1294,15 +990,14 @@ class FalconOpenVINOConfig(FalconOnnxConfig): ], library_name="transformers", ) -class PersimmonOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class PersimmonOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = PersimmonModelPatcher @register_in_tasks_manager("biogpt", *["text-generation", "text-generation-with-past"], library_name="transformers") class BioGPTOpenVINOConfig( - TextDecoderWithPositionIdsOnnxConfig if is_transformers_version(">=", "4.52.0") else TextDecoderOnnxConfig + TextDecoderWithPositionIdsOpenVINOConfig if is_transformers_version(">=", "4.52.0") else TextDecoderOpenVINOConfig ): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = OVDecoderModelPatcher @@ -1319,7 +1014,8 @@ class BioGPTOpenVINOConfig( ], library_name="transformers", ) -class GPTNeoOpenVINOConfig(GPTNeoOnnxConfig): +class GPTNeoOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads") _MODEL_PATCHER = GptNeoModelPatcher @@ -1334,7 +1030,8 @@ class GPTNeoOpenVINOConfig(GPTNeoOnnxConfig): ], library_name="transformers", ) -class GPTJOpenVINOConfig(GPTJOnnxConfig): +class GPTJOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = GptJModelPatcher @@ -1350,9 +1047,32 @@ class GPTJOpenVINOConfig(GPTJOnnxConfig): ], library_name="transformers", ) -class BloomOpenVINOConfig(BloomOnnxConfig): +class BloomOpenVINOConfig(TextDecoderOpenVINOConfig): + # Bloom does not require position_ids input. + MIN_TRANSFORMERS_VERSION = "4.36.0" + DUMMY_PKV_GENERATOR_CLASS = BloomDummyPastKeyValuesGenerator + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, BloomDummyPastKeyValuesGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = BloomModelPatcher + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): + if is_transformers_version(">=", "4.44"): + super().add_past_key_values(inputs_or_outputs, direction) + else: + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key"] = {0: "batch_size * num_heads", 2: decoder_sequence_name} + inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size * num_heads", 1: decoder_sequence_name} + @register_in_tasks_manager( "cohere", @@ -1370,8 +1090,7 @@ class CohereOpenVINOConfig(LlamaOpenVINOConfig): @register_in_tasks_manager("xglm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class XGLMConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 13 +class XGLMConfig(TextDecoderWithPositionIdsOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_attention_heads="attention_heads", hidden_size="d_model" ) @@ -1382,49 +1101,8 @@ def __init__(self, *args, **kwargs): _warn_potential_accuracy_issue_ov_2026_1("xglm") -class AquilaDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task, - normalized_config, - batch_size, - sequence_length, - random_batch_size_range, - random_sequence_length_range, - **kwargs, - ) - self.num_key_value_heads = getattr( - normalized_config, "num_key_value_heads", normalized_config.num_attention_heads - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - shape = ( - self.batch_size, - self.num_key_value_heads, - self.sequence_length, - self.hidden_size // self.num_attention_heads, - ) - return [ - ( - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager("aquila", *["text-generation", "text-generation-with-past"], library_name="transformers") -class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, AquilaDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = AquilaDummyPastKeyValuesGenerator @@ -1433,8 +1111,7 @@ class AquilaMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("xverse", *["text-generation", "text-generation-with-past"], library_name="transformers") -class XverseMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class XverseMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = DummyPastKeyValuesGenerator @@ -1443,8 +1120,7 @@ class XverseMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): @register_in_tasks_manager("internlm", *["text-generation", "text-generation-with-past"], library_name="transformers") -class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = DummyPastKeyValuesGenerator @@ -1457,7 +1133,8 @@ class InternLMOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): *["feature-extraction", "feature-extraction-with-past", "text-generation", "text-generation-with-past"], library_name="transformers", ) -class CodeGenOpenVINOConfig(CodeGenOnnxConfig): +class CodeGenOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = CodeGenModelPatcher @@ -1466,8 +1143,7 @@ class CodeGenOpenVINOConfig(CodeGenOnnxConfig): *["text-generation", "text-generation-with-past"], library_name="transformers", ) -class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_attention_heads="n_heads", @@ -1486,8 +1162,7 @@ class DBRXOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): *["text-generation", "text-generation-with-past"], library_name="transformers", ) -class JaisOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class JaisOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DummyPastKeyValuesGenerator) @@ -1501,44 +1176,6 @@ class ArcticOpenVINOConfig(MixtralOpenVINOConfig): _MODEL_PATCHER = ArcticModelPatcher -class OVMistralDummyPastKeyValuesGenerator(MistralDummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - **kwargs, - ) - self.head_dim = getattr(normalized_config, "head_dim", self.hidden_size // self.num_attention_heads) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - shape = ( - self.batch_size, - self.num_key_value_heads, - self.sequence_length, - self.head_dim, - ) - return [ - ( - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.num_layers) - ] - - @register_in_tasks_manager( "mistral", *[ @@ -1550,11 +1187,12 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int ], library_name="transformers", ) -class MistralOpenVINOConfig(MistralOnnxConfig): +class MistralOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_key_value_heads="num_key_value_heads", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = ( - OVMistralDummyPastKeyValuesGenerator, - ) + TextDecoderOnnxConfig.DUMMY_INPUT_GENERATOR_CLASSES - DUMMY_PKV_GENERATOR_CLASS = OVMistralDummyPastKeyValuesGenerator + MistralDummyPastKeyValuesGenerator, + ) + TextDecoderOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator _MODEL_PATCHER = MistralModelPatcher @@ -1569,16 +1207,16 @@ class MistralOpenVINOConfig(MistralOnnxConfig): ], library_name="transformers", ) -class GPTNeoxOpenVINOConfig(GPTNeoXOnnxConfig): +class GPTNeoxOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = GptNeoxModelPatcher @register_in_tasks_manager( "gpt_neox_japanese", *["text-generation", "text-generation-with-past"], library_name="transformers" ) -class GPTNeoxJapaneseOpenVINOConfig(TextDecoderOnnxConfig): +class GPTNeoxJapaneseOpenVINOConfig(TextDecoderOpenVINOConfig): # GPTNeoxJapanese does not require position_ids input. - DEFAULT_ONNX_OPSET = 13 NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = GptNeoxModelPatcher @@ -1614,67 +1252,6 @@ class Gemma3TextOpenVINOConfig(Gemma2OpenVINOConfig): MIN_TRANSFORMERS_VERSION = "4.50.0" -class Gemma4DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.num_key_value_heads = normalized_config.num_key_value_heads - self.head_dim = normalized_config.head_dim - self.global_head_dim = getattr(normalized_config.config, "global_head_dim", self.head_dim) - self.layer_types = normalized_config.config.layer_types - self.num_kv_shared_layers = normalized_config.config.num_kv_shared_layers - self.sliding_window = normalized_config.config.sliding_window - # Full-attention layers use fewer KV heads than sliding-attention layers (e.g. 2 vs 8 for 26B-A4B) - self.num_global_key_value_heads = ( - getattr(normalized_config.config, "num_global_key_value_heads", None) or self.num_key_value_heads - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - # some layers do not produce their own KV-cache, they use the shared KV-cache - if self.num_kv_shared_layers > 0: - layer_types = self.layer_types[: -self.num_kv_shared_layers] - else: - layer_types = self.layer_types - past_kv_values = [] - for layer_type in layer_types: - if layer_type == "sliding_attention": - shape = ( - self.batch_size, - self.num_key_value_heads, - self.sliding_window, - self.head_dim, - ) - else: - shape = ( - self.batch_size, - self.num_global_key_value_heads, - self.sequence_length, - self.global_head_dim, - ) - past_kv_value = ( - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - ) - past_kv_values.append(past_kv_value) - - return past_kv_values - - @register_in_tasks_manager( "gemma4_text", *[ @@ -1714,49 +1291,8 @@ def add_past_key_values(self, inputs_or_outputs: dict[str, dict[int, str]], dire inputs_or_outputs[f"{name}.{i}.value"] = {0: "batch_size", 2: decoder_sequence_name} -class DeciDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - ) - self.num_key_value_heads_per_layer = normalized_config.num_key_value_heads_per_layer - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_values = [] - - for layer_id in range(self.num_layers): - shape = ( - self.batch_size, - self.num_key_value_heads_per_layer[layer_id], - self.sequence_length, - self.hidden_size // self.num_attention_heads, - ) - past_key_values.append( - ( - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(shape, framework=framework, dtype=float_dtype), - ) - ) - return past_key_values - - @register_in_tasks_manager("deci", *["text-generation", "text-generation-with-past"], library_name="transformers") -class DeciOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class DeciOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): MAX_TRANSFORMERS_VERSION = "4.57.6" DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, DeciDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = DeciDummyPastKeyValuesGenerator @@ -1764,9 +1300,19 @@ class DeciOpenVINOConfig(TextDecoderWithPositionIdsOnnxConfig): _MODEL_PATCHER = DeciLMModelPatcher +class CLIPNormalizedConfig(NormalizedTextAndVisionConfig): + TEXT_CONFIG = "text_config" + VISION_CONFIG = "vision_config" + + +class SiglipNormalizedConfig(CLIPNormalizedConfig): + pass + + @register_in_tasks_manager("clip", *["zero-shot-image-classification"], library_name="open_clip") -class OpenCLIPOpenVINOConfig(CLIPOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class OpenCLIPOpenVINOConfig(TextAndVisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -1795,20 +1341,26 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): return super().generate_dummy_inputs(framework, **kwargs) def generate_dummy_inputs_for_validation( - self, reference_model_inputs: Dict[str, Any], onnx_input_names: Optional[List[str]] = None + self, reference_model_inputs: Dict[str, Any], ov_input_names: Optional[List[str]] = None ) -> Dict[str, Any]: if "attention_mask" in reference_model_inputs: reference_model_inputs.pop("attention_mask") - if "image" in onnx_input_names and "pixel_values" in reference_model_inputs: + if "image" in ov_input_names and "pixel_values" in reference_model_inputs: reference_model_inputs["image"] = reference_model_inputs.pop("pixel_values") - if "text" in onnx_input_names and "input_ids" in reference_model_inputs: + if "text" in ov_input_names and "input_ids" in reference_model_inputs: reference_model_inputs["text"] = reference_model_inputs.pop("input_ids") return super().generate_dummy_inputs_for_validation(reference_model_inputs) @register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="open_clip") -class OpenCLIPTextOpenVINOConfig(CLIPTextOnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class OpenCLIPTextOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher @property def inputs(self) -> Dict[str, Dict[int, str]]: @@ -1837,9 +1389,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): @register_in_tasks_manager("clip_vision_model", *["feature-extraction"], library_name="open_clip") -class OpenCLIPVisualOpenVINOConfig(VisionOnnxConfig): - DEFAULT_ONNX_OPSET = 14 - +class OpenCLIPVisualOpenVINOConfig(VisionOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @property @@ -1863,24 +1413,113 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager( "clip", *["feature-extraction", "zero-shot-image-classification"], library_name="transformers" ) -class CLIPOpenVINOConfig(CLIPOnnxConfig): - pass - +class CLIPOpenVINOConfig(TextAndVisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = CLIPNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher -@register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="transformers") -@register_in_tasks_manager("clip-text", *["feature-extraction"], library_name="diffusers") -class CLIPTextOpenVINOConfig(CLIPTextOnnxConfig): - pass + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + inputs = {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + inputs.update( + { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "attention_mask": {0: "text_batch_size", 1: "sequence_length"}, + } + ) + return inputs -@register_in_tasks_manager("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") -class CLIPTextWithProjectionOpenVINOConfig(CLIPTextWithProjectionOnnxConfig): - pass + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + return { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + else: + return super().outputs -@register_in_tasks_manager("clip_vision_model", *["feature-extraction"], library_name="transformers") -class CLIPVisionModelOpenVINOConfig(CLIPVisionModelOnnxConfig): - pass +@register_in_tasks_manager("clip_text_model", *["feature-extraction"], library_name="transformers") +@register_in_tasks_manager("clip-text", *["feature-extraction"], library_name="diffusers") +class CLIPTextOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + "pooler_output": {0: "batch_size"}, + } + + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs + + +@register_in_tasks_manager("clip-text-with-projection", *["feature-extraction"], library_name="diffusers") +class CLIPTextWithProjectionOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = { + "text_embeds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} + + return common_outputs + + +@register_in_tasks_manager("clip_vision_model", *["feature-extraction"], library_name="transformers") +class CLIPVisionModelOpenVINOConfig(VisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return {"pixel_values": {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"}} + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + common_outputs = super().outputs + common_outputs["last_hidden_state"] = {0: "batch_size"} + common_outputs["pooler_output"] = {0: "batch_size"} + + return common_outputs @register_in_tasks_manager( @@ -1895,12 +1534,21 @@ class CLIPVisionModelOpenVINOConfig(CLIPVisionModelOnnxConfig): ], library_name="transformers", ) -class IBertOpenVINOConfig(IBertOnnxConfig): +class IBertOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = IBertModelPatcher + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} + # TODO: this is a very confusing class TBH, why not simply decompose the VLM into components, like diffusion models ? -class LMInputEmbedsConfigHelper(TextDecoderWithPositionIdsOnnxConfig): +class LMInputEmbedsConfigHelper(TextDecoderWithPositionIdsOpenVINOConfig): def __init__(self, export_config, patcher_cls=None, dummy_input_generator=None, inputs_update=None): self.orig_export_config = export_config if dummy_input_generator is not None: @@ -1908,7 +1556,6 @@ def __init__(self, export_config, patcher_cls=None, dummy_input_generator=None, dummy_input_generator, ) + export_config.DUMMY_INPUT_GENERATOR_CLASSES self.DUMMY_INPUT_GENERATOR_CLASSES = export_config.DUMMY_INPUT_GENERATOR_CLASSES - self.DEFAULT_ONNX_OPSET = export_config.DEFAULT_ONNX_OPSET self.DUMMY_PKV_GENERATOR_CLASS = export_config.DUMMY_PKV_GENERATOR_CLASS self._config = export_config._config self._normalized_config = export_config._normalized_config @@ -1969,7 +1616,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): return dummy_inputs -class InputEmbedOpenVINOConfig(TextDecoderOnnxConfig): +class InputEmbedOpenVINOConfig(TextDecoderOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig _MODEL_PATCHER = InputEmbeddingPatcher @@ -2048,7 +1695,7 @@ class VLMConfigBehavior(str, enum.Enum): LANGUAGE = "language" -class BaseVLMOpenVINOConfig(OnnxConfig): +class BaseVLMOpenVINOConfig(OpenVINOConfig): SUPPORTED_BEHAVIORS = [model_type.value for model_type in VLMConfigBehavior] NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator,) @@ -2189,36 +1836,7 @@ class LlavaNextOpenVINOConfig(LlavaOpenVINOConfig): _OV_2026_1_MODEL_TYPE = "llava_next" -class DummyLLavaMultiModalProjectorInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ["image_features"] - - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - self.task = task - - self.batch_size = batch_size - self.hidden_size = normalized_config.hidden_size - self.num_patches = (normalized_config.image_size // normalized_config.patch_size) ** 2 - self.normalized_config = normalized_config - - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - ): - shape = [self.batch_size, self.num_patches, self.hidden_size] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - -class LLavaMultimodalProjectorOpenVINOConfig(OnnxConfig): +class LLavaMultimodalProjectorOpenVINOConfig(OpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyLLavaMultiModalProjectorInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @@ -2526,95 +2144,16 @@ def rename_ambiguous_inputs(self, inputs): return super().rename_ambiguous_inputs(inputs) -class PooledProjectionsDummyInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ["pooled_projections"] - - def __init__( - self, - task: str, - normalized_config: NormalizedConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - self.task = task - self.batch_size = batch_size - self.pooled_projection_dim = normalized_config.config.pooled_projection_dim - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - shape = [self.batch_size, self.pooled_projection_dim] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - -class DummyTransformerTimestpsInputGenerator(DummyTimestepInputGenerator): - SUPPORTED_INPUT_NAMES = ("timestep", "text_embeds", "time_ids", "timestep_cond", "guidance") - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["timestep", "guidance"]: - shape = [self.batch_size] - return self.random_float_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) - return super().generate(input_name, framework, int_dtype, float_dtype) - - -class DummyUnetVisionInputGenerator(DummyVisionInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name not in ["sample", "latent_sample"]: - return super().generate(input_name, framework, int_dtype, float_dtype) - # add height and width discount for enable any resolution generation - return self.random_float_tensor( - shape=[self.batch_size, self.num_channels, self.height - 1, self.width - 1], - framework=framework, - dtype=float_dtype, - ) - - -class DummyUnetTimestepInputGenerator(DummyTimestepInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name != "timestep": - return super().generate(input_name, framework, int_dtype, float_dtype) - shape = [self.batch_size] - return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=int_dtype) - - -class DummySanaTimestepInputGenerator(DummyTimestepInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name != "timestep": - return super().generate(input_name, framework, int_dtype, float_dtype) - shape = [self.batch_size] - return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) - - -class DummyUnetEncoderInputGenerator(DummySeq2SeqDecoderTextInputGenerator): - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - num_choices: int = DEFAULT_DUMMY_SHAPES["num_choices"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - random_sequence_length_range: Optional[Tuple[int, int]] = None, - random_num_choices_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - super().__init__( - task, - normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - num_choices=num_choices, - random_batch_size_range=random_batch_size_range, - random_sequence_length_range=random_sequence_length_range, - random_num_choices_range=random_num_choices_range, - **kwargs, - ) - if hasattr(normalized_config.config, "model_max_length"): - self.sequence_length = normalized_config.config.model_max_length - - @register_in_tasks_manager("unet", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("unet-2d-condition", *["semantic-segmentation"], library_name="diffusers") -class UNetOpenVINOConfig(UNetOnnxConfig): +class UNetOpenVINOConfig(VisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + image_size="sample_size", + num_channels="in_channels", + hidden_size="cross_attention_dim", + vocab_size="norm_num_groups", + allow_new=True, + ) DUMMY_INPUT_GENERATOR_CLASSES = ( DummyUnetVisionInputGenerator, DummyUnetTimestepInputGenerator, @@ -2623,12 +2162,59 @@ class UNetOpenVINOConfig(UNetOnnxConfig): @property def inputs(self) -> Dict[str, Dict[int, str]]: - common_inputs = super().inputs - common_inputs["timestep"] = {0: "batch_size"} + common_inputs = { + "sample": {0: "batch_size", 2: "height", 3: "width"}, + "timestep": {0: "batch_size"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + } + + # TODO : add addition_embed_type == text_image, image and image_embeds + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + common_inputs["text_embeds"] = {0: "batch_size"} + common_inputs["time_ids"] = {0: "batch_size"} + + if getattr(self._normalized_config, "time_cond_proj_dim", None) is not None: + common_inputs["timestep_cond"] = {0: "batch_size"} + if hasattr(self._normalized_config.config, "model_max_length"): common_inputs["encoder_hidden_states"] = {0: "batch_size"} + return common_inputs + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "out_sample": {0: "batch_size", 2: "height", 3: "width"}, + } + + @property + def torch_to_ov_output_map(self) -> Dict[str, str]: + return { + "sample": "out_sample", + } + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + dummy_inputs["encoder_hidden_states"] = dummy_inputs["encoder_hidden_states"][0] + + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + dummy_inputs["added_cond_kwargs"] = { + "text_embeds": dummy_inputs.pop("text_embeds"), + "time_ids": dummy_inputs.pop("time_ids"), + } + + return dummy_inputs + + def ordered_inputs(self, model) -> Dict[str, Dict[int, str]]: + inputs = super().ordered_inputs(model=model) + # to fix mismatch between model forward signature and expected inputs + # a dictionary of additional embeddings `added_cond_kwargs` is expected depending on config.addition_embed_type + if getattr(self._normalized_config, "addition_embed_type", None) == "text_time": + inputs["text_embeds"] = self.inputs["text_embeds"] + inputs["time_ids"] = self.inputs["time_ids"] + + return inputs + @register_in_tasks_manager("sd3-transformer", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("sd3-transformer-2d", *["semantic-segmentation"], library_name="diffusers") @@ -2678,44 +2264,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: } -class DummySanaSeq2SeqDecoderTextWithEncMaskInputGenerator(DummySeq2SeqDecoderTextInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "decoder_input_ids", - "decoder_attention_mask", - "encoder_outputs", - "encoder_hidden_states", - "encoder_attention_mask", - ) - - -class DummySanaTransformerVisionInputGenerator(DummyUnetVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "pixel_values", - "pixel_mask", - "sample", - "latent_sample", - "guidance", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"] // 8, - height: int = DEFAULT_DUMMY_SHAPES["height"] // 8, - # Reduce img shape by 4 for FLUX to reduce memory usage on conversion - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width=width, height=height, **kwargs) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "guidance": - return self.random_float_tensor([self.batch_size], framework=framework, dtype=float_dtype) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("sana-transformer", *["semantic-segmentation"], library_name="diffusers") class SanaTransformerOpenVINOConfig(UNetOpenVINOConfig): NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( @@ -2748,17 +2296,62 @@ def rename_ambiguous_inputs(self, inputs): @register_in_tasks_manager("vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class VaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig): - pass +class VaeEncoderOpenVINOConfig(VisionOpenVINOConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "sample": {0: "batch_size", 2: "sample_height", 3: "sample_width"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + down_sampling_factor = 2 ** (len(self._normalized_config.down_block_types) - 1) + + return { + "latent_parameters": { + 0: "batch_size", + 2: f"sample_height / {down_sampling_factor}", + 3: f"sample_width / {down_sampling_factor}", + }, + } @register_in_tasks_manager("vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class VaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig): - pass +class VaeDecoderOpenVINOConfig(VisionOpenVINOConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "latent_sample": {0: "batch_size", 2: "latent_height", 3: "latent_width"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + up_sampling_factor = 2 ** (len(self._normalized_config.up_block_types) - 1) + + return { + "sample": { + 0: "batch_size", + 2: f"latent_height * {up_sampling_factor}", + 3: f"latent_width * {up_sampling_factor}", + }, + } @register_in_tasks_manager("dcae-encoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig): +class DcaeEncoderOpenVINOConfig(VisionOpenVINOConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) + @property def inputs(self) -> Dict[str, Dict[int, str]]: return { @@ -2773,7 +2366,10 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("dcae-decoder", *["semantic-segmentation"], library_name="diffusers") -class DcaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig): +class DcaeDecoderOpenVINOConfig(VisionOpenVINOConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) + @property def inputs(self) -> Dict[str, Dict[int, str]]: return { @@ -2787,76 +2383,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } -class DummyFluxTransformerInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "pixel_values", - "pixel_mask", - "sample", - "latent_sample", - "hidden_states", - "img_ids", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"] // 4, - height: int = DEFAULT_DUMMY_SHAPES["height"] // 4, - # Reduce img shape by 4 for FLUX to reduce memory usage on conversion - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - if getattr(normalized_config, "in_channels", None): - self.num_channels = normalized_config.in_channels // 4 - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["hidden_states", "sample"]: - shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels * 4] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - if input_name == "img_ids": - img_ids_height = self.height // 2 - img_ids_width = self.width // 2 - return self.random_int_tensor( - ( - [self.batch_size, img_ids_height * img_ids_width, 3] - if is_diffusers_version("<", "0.31.0") - else [img_ids_height * img_ids_width, 3] - ), - min_value=0, - max_value=min(img_ids_height, img_ids_width), - framework=framework, - dtype=float_dtype, - ) - - return super().generate(input_name, framework, int_dtype, float_dtype) - - -class DummyFluxTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "decoder_input_ids", - "decoder_attention_mask", - "encoder_outputs", - "encoder_hidden_states", - "txt_ids", - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "txt_ids": - import torch - - shape = ( - [self.batch_size, self.sequence_length, 3] - if is_diffusers_version("<", "0.31.0") - else [self.sequence_length, 3] - ) - dtype = DTYPE_MAPPER.pt(float_dtype) - return torch.full(shape, 0, dtype=dtype) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("flux-transformer", *["semantic-segmentation"], library_name="diffusers") @register_in_tasks_manager("flux-transformer-2d", *["semantic-segmentation"], library_name="diffusers") class FluxTransformerOpenVINOConfig(SD3TransformerOpenVINOConfig): @@ -2866,7 +2392,7 @@ class FluxTransformerOpenVINOConfig(SD3TransformerOpenVINOConfig): DummyFluxTextInputGenerator, PooledProjectionsDummyInputGenerator, ) - _MODEL_PATCHER = FluxTransfromerModelPatcher + _MODEL_PATCHER = FluxTransformerModelPatcher @property def inputs(self): @@ -2886,36 +2412,12 @@ def inputs(self): return common_inputs -class LTXVaeDummyInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("pixel_values", "pixel_mask", "sample", "latent_sample", "timestep") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - num_frames: int = 2, - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - self.num_frames = num_frames - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["sample", "latent_sample"]: - return self.random_float_tensor( - [self.batch_size, self.num_channels, self.num_frames, self.height, self.width] - ) - if input_name == "timestep": - return self.random_int_tensor([1], max_value=20, min_value=1, framework=framework, dtype=int_dtype) - - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("ltx-vae-encoder", *["semantic-segmentation"], library_name="diffusers") -class LTXVaeEncoderOpenVINOConfig(VaeEncoderOnnxConfig): +class LTXVaeEncoderOpenVINOConfig(VisionOpenVINOConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_channels="in_channels", image_size="sample_size", allow_new=True + ) DUMMY_INPUT_GENERATOR_CLASSES = (LTXVaeDummyInputGenerator,) @property @@ -2932,7 +2434,9 @@ def outputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("ltx-vae-decoder", *["semantic-segmentation"], library_name="diffusers") -class LTXVaeDecoderOpenVINOConfig(VaeDecoderOnnxConfig): +class LTXVaeDecoderOpenVINOConfig(VisionOpenVINOConfig): + ATOL_FOR_VALIDATION = 3e-4 # TODO: this only happens in test_export.py + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(num_channels="latent_channels", allow_new=True) DUMMY_INPUT_GENERATOR_CLASSES = (LTXVaeDummyInputGenerator,) @property @@ -2951,53 +2455,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } -class LTXTransformerDummyInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("hidden_states", "width", "height", "num_frames", "rope_interpolation_scale") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = 16, - height: int = 8, - num_frames: int = 2, - frame_rate: int = 10, - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - self.num_frames = num_frames - self.frame_rate = frame_rate - self.vae_spatial_compression_ratio = normalized_config.config.vae_spatial_compression_ratio - self.vae_temporal_compression_ratio = normalized_config.config.vae_temporal_compression_ratio - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - import torch - - if input_name == "hidden_states": - return self.random_float_tensor( - [self.batch_size, self.num_frames * self.height * self.width, self.num_channels] - ) - if input_name == "width": - return torch.tensor(self.width) - if input_name == "height": - return torch.tensor(self.height) - if input_name == "num_frames": - return torch.tensor(self.num_frames) - if input_name == "rope_interpolation_scale": - import torch - - return torch.tensor( - [ - self.vae_temporal_compression_ratio / self.frame_rate, - self.vae_spatial_compression_ratio, - self.vae_spatial_compression_ratio, - ] - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("ltx-video-transformer", *["semantic-segmentation"], library_name="diffusers") class LTXVideoTransformerOpenVINOConfig(SanaTransformerOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = ( @@ -3026,90 +2483,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } -class DummyMiniCPMVImageInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("pixel_values", "patch_attention_mask", "position_ids") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height) - self.patch_size = normalized_config.config.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "pixel_values": - return self.random_float_tensor( - shape=[ - self.batch_size, - self.num_channels, - self.patch_size, - (self.height * self.width) // self.patch_size, - ], - framework=framework, - dtype=float_dtype, - ) - - if input_name == "patch_attention_mask": - return self.random_int_tensor( - shape=[self.batch_size, 1, (self.height // self.patch_size) * (self.width // self.patch_size)], - framework=framework, - dtype=float_dtype, - min_value=0, - max_value=2, - ) - - if input_name == "position_ids": - return self.random_int_tensor( - shape=[self.batch_size, (self.height // self.patch_size) * (self.width // self.patch_size)], - max_value=self.patch_size, - ) - - -class DummyMiniCPMVResampleInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("image_feature", "pos_embed", "key_padding_mask") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height) - self.patch_size = normalized_config.config.patch_size - self.hidden_size = normalized_config.config.hidden_size - self.img_hidden_size = normalized_config.config.vision_config.hidden_size - self.feat_size = (normalized_config.config.vision_config.image_size // self.patch_size) * ( - normalized_config.config.vision_config.image_size // self.patch_size - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "image_feature": - return self.random_float_tensor( - shape=[self.batch_size, self.feat_size, self.img_hidden_size], framework=framework, dtype=float_dtype - ) - - if input_name == "key_padding_mask": - return self.constant_tensor( - shape=[self.batch_size, self.feat_size], - framework=framework, - value=1, - dtype=DTYPE_MAPPER.pt(float_dtype), - ) - - if input_name == "pos_embed": - return self.random_float_tensor(shape=[self.feat_size, self.batch_size, self.hidden_size]) - - class MiniCPMVConfigBehavior(str, enum.Enum): RESAMPLER = "resampler" LANGUAGE = "language" @@ -3259,55 +2632,6 @@ class Phi3VisionConfigBehavior(str, enum.Enum): TEXT_EMBEDDINGS = "text_embeddings" -class DummyPhi3VisionProjectionInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("input",) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = 336, - height: int = 336, - crop_size=336, - **kwargs, - ): - self.batch_size = batch_size - self._embed_layer_realization = ( - normalized_config.config.embd_layer["embedding_cls"] - if hasattr(normalized_config.config, "embd_layer") - else "image_audio" - ) - if not hasattr(normalized_config.config, "vision_config"): - self.image_dim_out = ( - normalized_config.config.img_processor.get( - "image_dim_out", normalized_config.config.img_processor.get("hidden_size") - ) - if normalized_config.config.img_processor is not None - else 1152 - ) - if "image_embd_layer" in normalized_config.config.embd_layer: - self.crop_size = normalized_config.config.embd_layer["image_embd_layer"].get("crop_size", crop_size) - else: - self.crop_size = normalized_config.config.embd_layer.get("crop_size", crop_size) - else: - self.image_dim_out = normalized_config.config.vision_config.hidden_size - self.crop_size = normalized_config.config.vision_config.crop_size - self.height = height - self.width = width - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - h = self.height // self.crop_size - w = self.width // self.crop_size - feat_size = (h * w + 1) * 144 + 1 + (h + 1) * 12 - if self._embed_layer_realization in ["linear", "image_audio"]: - shape = [self.batch_size, feat_size, self.image_dim_out] - else: - shape = [self.batch_size, feat_size, self.image_dim_out * 4] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) - - @register_in_tasks_manager("phi3_v", *["image-text-to-text"], library_name="transformers") class Phi3VisionOpenVINOConfig(BaseVLMOpenVINOConfig): SUPPORTED_BEHAVIORS = [model_type.value for model_type in Phi3VisionConfigBehavior] @@ -3426,129 +2750,16 @@ def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[ return super().patch_model_for_export(model, model_kwargs) -class DummyAudioPhi4MMInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ("audio_input", "audio_feature", "audio_mask") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - signal_length=498, - **kwargs, - ): - self.signal_length = signal_length - if hasattr(normalized_config.config, "audio_processor"): - self.audio_chunk_size = ( - signal_length // normalized_config.config.audio_processor["config"]["time_reduction"] + 1 - ) - self.input_size = normalized_config.config.audio_processor["config"]["input_size"] - self.attention_dim = normalized_config.config.audio_processor["config"]["attention_dim"] - else: - self.audio_chunk_size = signal_length // normalized_config.config.audio_config.time_reduction + 1 - self.input_size = normalized_config.config.audio_config.input_size - self.attention_dim = normalized_config.config.audio_config.hidden_size - self.batch_size = batch_size - self.task = task - self.normalized_config = normalized_config - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "audio_input": - return self.random_float_tensor( - [self.batch_size, self.signal_length, self.input_size], framework=framework, dtype=float_dtype - ) - - if input_name == "audio_feature": - return self.random_float_tensor( - [self.batch_size, self.audio_chunk_size, self.attention_dim], framework=framework, dtype=float_dtype - ) - - if input_name == "audio_mask": - return self.random_int_tensor( - [self.batch_size, self.audio_chunk_size, self.audio_chunk_size], - max_value=2, - framework=framework, - dtype="bool", - ) - - -class DummyVisionPositionIdsPhi4InputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("patch_position_ids", "patch_attention_mask") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - if hasattr(normalized_config.config, "vision_conifg"): - self.patch_size = getattr(normalized_config.config.vision_config, "patch_size", 14) - else: - self.patch_size = 14 - self.num_patches_per_side = self.height // self.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "patch_position_ids": - return self.get_vision_position_ids() - if input_name == "patch_attention_mask": - return self.random_int_tensor( - [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], - framework=framework, - dtype="bool", - max_value=2, - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - def get_vision_position_ids(self): - # Adopted from https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/models/phi4_multimodal/modeling_phi4_multimodal.py#L494-L512 - import torch - - batch_size = self.batch_size - max_im_h, max_im_w = self.height, self.width - max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size - boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side) - position_ids = torch.full( - size=( - batch_size, - max_nb_patches_h * max_nb_patches_w, - ), - fill_value=0, - ) - patch_attention_mask = torch.ones( - [self.batch_size, self.height // self.patch_size, self.width // self.patch_size], dtype=torch.int64 - ) - patch_attention_mask[0, self.height - 2 :] = 0 - - for batch_idx, p_attn_mask in enumerate(patch_attention_mask): - nb_patches_h = p_attn_mask[:, 0].sum() - nb_patches_w = p_attn_mask[0].sum() - - fractional_coords_h = torch.arange(0, 1 - 1e-6, 1 / nb_patches_h) - fractional_coords_w = torch.arange(0, 1 - 1e-6, 1 / nb_patches_w) - - bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) - bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) - - pos_ids = (bucket_coords_h[:, None] * self.num_patches_per_side + bucket_coords_w).flatten() - position_ids[batch_idx][p_attn_mask.view(-1)] = pos_ids - return position_ids - - -class Phi4MMConfigBehavior(str, enum.Enum): - AUDIO_EMBEDDINGS = "audio_embeddings" - AUDIO_ENCODER = "audio_encoder" - AUDIO_FORWARD_EMBEDDINGS = "audio_forward_embeddings" - AUDIO_VISION_PROJECTION = "audio_vision_projection" - AUDIO_SPEECH_PROJECTION = "audio_speech_projection" - LANGUAGE = "language" - TEXT_EMBEDDINGS = "text_embeddings" - VISION_PROJECTION = "vision_projection" - VISION_EMBEDDINGS = "vision_embeddings" +class Phi4MMConfigBehavior(str, enum.Enum): + AUDIO_EMBEDDINGS = "audio_embeddings" + AUDIO_ENCODER = "audio_encoder" + AUDIO_FORWARD_EMBEDDINGS = "audio_forward_embeddings" + AUDIO_VISION_PROJECTION = "audio_vision_projection" + AUDIO_SPEECH_PROJECTION = "audio_speech_projection" + LANGUAGE = "language" + TEXT_EMBEDDINGS = "text_embeddings" + VISION_PROJECTION = "vision_projection" + VISION_EMBEDDINGS = "vision_embeddings" @register_in_tasks_manager( @@ -3780,123 +2991,6 @@ def rename_ambiguous_inputs(self, inputs): return inputs -class DummyQwen2VLLMInputGenerator(DummyTextInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - generated_input = super().generate(input_name, framework, int_dtype, float_dtype) - if input_name == "position_ids": - return generated_input.unsqueeze(0).expand(3, -1, -1) - return generated_input - - -class DummyQwen3_5LMInputGenerator(DummyTextInputGenerator): - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - generated_input = super().generate(input_name, framework, int_dtype, float_dtype) - if input_name == "position_ids": - return generated_input.unsqueeze(0).expand(4, -1, -1) - return generated_input - - -class DummyQwen2VLVisionEmbedInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "hidden_states", - "attention_mask", - "window_attention_mask", - "window_index", - "rotary_pos_emb", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = 1, - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = 420, - height: int = 420, - **kwargs, - ): - self.batch_size = batch_size - self.height = height - self.width = width - self.num_channels = num_channels - self.temporal_patch_size = normalized_config.config.temporal_patch_size - self.patch_size = normalized_config.config.patch_size - if normalized_config.use_embed_dim: - self.embed_dim = ( - normalized_config.config.embed_dim - if hasattr(normalized_config.config, "embed_dim") - else normalized_config.hidden_size - ) - else: - self.embed_dim = self.num_channels * self.temporal_patch_size * self.patch_size * self.patch_size - self.num_heads = normalized_config.config.num_heads - self.spatial_merge_size = None - if hasattr(normalized_config.config, "spatial_merge_size"): - self.spatial_merge_size = normalized_config.config.spatial_merge_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size - grid_t = self.batch_size - - if input_name == "hidden_states": - return self.random_float_tensor( - [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype - ) - - if input_name in ["attention_mask", "window_attention_mask"]: - return self.random_mask_tensor( - [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype - ) - - if input_name == "rotary_pos_emb": - dim = self.embed_dim // self.num_heads // 2 - return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) - - if input_name == "window_index": - if self.spatial_merge_size is None: - raise ValueError( - "`spatial_merge_size` parameter is not found in model config. Can not generate dummy input data for `window_index` input" - ) - spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size - hidden_size = (grid_t * grid_h * grid_w) // spatial_merge_unit - return self.random_int_tensor([hidden_size], max_value=hidden_size) - - -class DummyQwen3VLVisionEmbedInputGenerator(DummyQwen2VLVisionEmbedInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "hidden_states", - "attention_mask", - "rotary_pos_emb", - "input", - ) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size - grid_t = self.batch_size - - if input_name == "hidden_states": - return self.random_float_tensor( - [grid_t * grid_h * grid_w, self.embed_dim], framework=framework, dtype=float_dtype - ) - - if input_name in ["attention_mask"]: - return self.random_mask_tensor( - [1, grid_t * grid_h * grid_w, grid_t * grid_h * grid_w], framework=framework, dtype=float_dtype - ) - - if input_name == "rotary_pos_emb": - dim = self.embed_dim // self.num_heads // 2 - return self.random_float_tensor([grid_h * grid_t * grid_w, dim], framework=framework, dtype=float_dtype) - - if input_name == "input": - return self.constant_tensor( - [4, DEFAULT_DUMMY_SHAPES["sequence_length"]], - framework=framework, - value=0, - dtype=DTYPE_MAPPER.pt(int_dtype), - ) - - class QwenVLConfigBehavior(str, enum.Enum): LANGUAGE = "language" VISION_EMBEDDINGS = "vision_embeddings" @@ -4255,39 +3349,50 @@ class GraniteMoEOpenVINOConfig(LlamaOpenVINOConfig): ], library_name="transformers", ) -class WhisperOpenVINOConfig(WhisperOnnxConfig): +class WhisperOpenVINOConfig(AudioToTextOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + feature_size="num_mel_bins", + allow_new=True, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "audio-classification": + return {"input_features": {0: "batch_size"}} -class Qwen3ASRDummySeq2SeqPastKeyValuesGenerator(DummySeq2SeqPastKeyValuesGenerator): - """Custom KV cache generator for Qwen3-ASR with GQA (num_key_value_heads != num_attention_heads). - Qwen3-ASR has no cross-attention, so only self-attention KV cache is generated (2 per layer).""" - - def __init__(self, task, normalized_config, **kwargs): - super().__init__(task, normalized_config, **kwargs) - # Override head count and head_dim for GQA - self.decoder_num_attention_heads = normalized_config.decoder_num_attention_heads - self.decoder_head_dim = getattr(normalized_config, "head_dim", None) - if self.decoder_head_dim is None: - self.decoder_head_dim = self.decoder_hidden_size // normalized_config.num_attention_heads - - def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="fp32"): - if input_name == "past_key_values": - decoder_shape = ( - self.batch_size, - self.decoder_num_attention_heads, - self.sequence_length, - self.decoder_head_dim, - ) - # Qwen3-ASR has no cross-attention, only self-attention KV cache - return [ - ( - self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(decoder_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.decoder_num_layers) - ] - return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) + common_inputs = super().inputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["input_features"] = {0: "batch_size"} # Remove unnecessary dynamic axis. + else: + # the dynamic encoder sequence length is only needed here because the input generator generates + # encoder_outputs with a seq_len=16 but the model expects at inference time seq_len=1500 + # TODO: this can be fixed by generating the correct inputs in the input generator + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + if is_transformers_version(">=", "4.43.0") and is_transformers_version("<", "4.46.0"): + # since https://github.com/huggingface/transformers/pull/31166 + if self._behavior is not ConfigBehavior.ENCODER and self.use_past_in_inputs: + common_inputs["cache_position"] = {0: "decoder_sequence_length"} + + return common_inputs + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "audio-classification": + return {"logits": {0: "batch_size"}} + + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. + return common_outputs @register_in_tasks_manager( @@ -4298,7 +3403,7 @@ def generate(self, input_name, framework="pt", int_dtype="int64", float_dtype="f ], library_name="transformers", ) -class Qwen3ASROpenVINOConfig(AudioToTextOnnxConfig): +class Qwen3ASROpenVINOConfig(AudioToTextOpenVINOConfig): """OpenVINO export config for Qwen3-ASR model.""" DUMMY_INPUT_GENERATOR_CLASSES = ( @@ -4382,7 +3487,20 @@ def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], dire *["feature-extraction", "feature-extraction-with-past", "text2text-generation", "text2text-generation-with-past"], library_name="transformers", ) -class T5OpenVINOConfig(T5OnnxConfig): +class T5OpenVINOConfig(TextSeq2SeqOpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + *TextSeq2SeqOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES[:-1], + T5DummySeq2SeqPastKeyValuesGenerator, + ) + DUMMY_PKV_GENERATOR_CLASS = T5DummySeq2SeqPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + hidden_size="d_model", + num_attention_heads="num_heads", + encoder_num_layers="num_layers", + decoder_num_layers="num_decoder_layers", + key_value_dim="d_kv", + allow_new=True, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher @@ -4419,9 +3537,124 @@ class LongT5OpenVINOConfig(T5OpenVINOConfig): ], library_name="transformers", ) -class BartOpenVINOConfig(BartOnnxConfig): +class BartOpenVINOConfig(TextSeq2SeqOpenVINOConfig): + PAD_ATTENTION_MASK_TO_PAST = True + + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", # Used for the text-generation task past key values input generation. + encoder_num_attention_heads="encoder_attention_heads", + decoder_num_attention_heads="decoder_attention_heads", + eos_token_id="eos_token_id", + ) + DUMMY_INPUT_GENERATOR_CLASSES = ( + BartDummyTextInputGenerator, + { + "feature-extraction": DummySeq2SeqDecoderTextInputGenerator, + "text-generation": DummyDecoderTextInputGenerator, + }, + { + "feature-extraction": DummySeq2SeqPastKeyValuesGenerator, + "text-generation": DummyPastKeyValuesGenerator, + }, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher + def _create_dummy_input_generator_classes(self, **kwargs): + dummy_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[0]( + self.task, self._normalized_config, **kwargs + ) + task = "feature-extraction" if self.task != "text-generation" else "text-generation" + dummy_decoder_text_input_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[1][task]( + self.task, self._normalized_config, **kwargs + ) + if self.task != "text-generation": + kwargs["encoder_sequence_length"] = dummy_text_input_generator.sequence_length + + dummy_seq2seq_past_key_values_generator = self.DUMMY_INPUT_GENERATOR_CLASSES[2][task]( + self.task, self._normalized_config, **kwargs + ) + dummy_inputs_generators = [ + dummy_text_input_generator, + dummy_decoder_text_input_generator, + dummy_seq2seq_past_key_values_generator, + ] + + return dummy_inputs_generators + + @property + def inputs_for_default_and_seq2seq_lm(self): + return super().inputs + + @property + def inputs_for_causal_lm(self): + if self.use_past_in_inputs: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "past_sequence_length + sequence_length"}, + } + for i in range(self._normalized_config.decoder_num_layers): + common_inputs[f"past_key_values.{i}.key"] = { + 0: "batch_size", + 2: "past_sequence_length", + } + common_inputs[f"past_key_values.{i}.value"] = { + 0: "batch_size", + 2: "past_sequence_length", + } + else: + common_inputs = { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + return common_inputs + + @property + def inputs_for_other_tasks(self): + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def inputs(self) -> dict: + inputs_properties = { + "feature-extraction": self.inputs_for_default_and_seq2seq_lm, + "text2text-generation": self.inputs_for_default_and_seq2seq_lm, + "text-generation": self.inputs_for_causal_lm, + "other": self.inputs_for_other_tasks, + } + return inputs_properties.get(self.task, inputs_properties["other"]) + + @property + def outputs(self) -> dict: + if self.task in ["feature-extraction", "text2text-generation"]: + common_outputs = super().outputs + else: + common_outputs = super(OpenVINOConfigWithPast, self).outputs + if self.use_past: + for i in range( + self._normalized_config.encoder_num_layers + if self.task != "text-generation" + else self._normalized_config.decoder_num_layers + ): + common_outputs[f"present.{i}.key"] = {0: "batch_size", 2: "past_sequence_length + sequence_length"} + common_outputs[f"present.{i}.value"] = { + 0: "batch_size", + 2: "past_sequence_length + sequence_length", + } + return common_outputs + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if self.task in ["feature-extraction", "text2text-generation"]: + flattened_output = super().flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output = super(OpenVINOSeq2SeqConfigWithPast, self).flatten_past_key_values( + flattened_output, name, idx, t + ) + @register_in_tasks_manager( "bigbird_pegasus", @@ -4573,54 +3806,6 @@ class Gemma4ConfigBehavior(str, enum.Enum): TEXT_EMBEDDINGS_PER_LAYER = "text_embeddings_per_layer" -class DummyGemma4VisionInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("pixel_values", "image_position_ids") - - def __init__(self, task, normalized_config, batch_size=DEFAULT_DUMMY_SHAPES["batch_size"], **kwargs): - super().__init__(task, normalized_config, batch_size, **kwargs) - self.patch_size = getattr(normalized_config, "patch_size", 16) - self.pooling_kernel_size = getattr(normalized_config, "pooling_kernel_size", 3) - # Gemma4 processor always pads pixel_values to max_soft_tokens * pooling_kernel_size^2 patches. - # The vision model's pooling uses shape-dependent Python operations that get baked in during tracing, - # so the dummy input must match the actual inference shapes. - max_soft_tokens = getattr(normalized_config, "image_seq_length", None) - if max_soft_tokens is None: - max_soft_tokens = getattr(normalized_config, "max_soft_tokens", 280) - self.num_patches = max_soft_tokens * self.pooling_kernel_size**2 - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "pixel_values": - # Gemma4 expects pre-patchified pixel_values: [batch, num_patches, 3 * patch_size^2] - return self.random_float_tensor( - shape=[self.batch_size, self.num_patches, 3 * self.patch_size**2], - framework=framework, - dtype=float_dtype, - ) - if input_name == "image_position_ids": - # Create position ids as a grid. The patch count = h_patches * w_patches - # where both are divisible by pooling_kernel_size for correct pooling. - k = self.pooling_kernel_size - total_pooled = self.num_patches // (k * k) - # Find roughly square grid for pooled side - pooled_side = int(math.sqrt(total_pooled)) - if pooled_side * pooled_side < total_pooled: - pooled_h = pooled_side - pooled_w = total_pooled // pooled_h - else: - pooled_h = pooled_w = pooled_side - h_patches = pooled_h * k - w_patches = pooled_w * k - pos_ids = torch.stack( - torch.meshgrid(torch.arange(h_patches), torch.arange(w_patches), indexing="ij"), dim=-1 - ).reshape(1, -1, 2) - # Pad to num_patches with -1 (padding position) - if pos_ids.shape[1] < self.num_patches: - pad = torch.full((1, self.num_patches - pos_ids.shape[1], 2), -1, dtype=pos_ids.dtype) - pos_ids = torch.cat([pos_ids, pad], dim=1) - return pos_ids.expand(self.batch_size, -1, -1).clone() - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("gemma4", *["image-text-to-text"], library_name="transformers") class Gemma4OpenVINOConfig(Gemma3OpenVINOConfig): SUPPORTED_BEHAVIORS = [model_type.value for model_type in Gemma4ConfigBehavior] @@ -4808,35 +3993,6 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return super().outputs -class DummyVisionPositionIdsInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("patch_attention_mask", "patch_position_ids") - - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, **kwargs) - self.patch_size = normalized_config.config.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "patch_attention_mask": - shape = [self.batch_size, self.height // self.patch_size, self.width // self.patch_size] - return self.random_int_tensor(shape, max_value=2, framework=framework, dtype="bool") - if input_name == "patch_position_ids": - max_nb_patches_h, max_nb_patches_w = self.height // self.patch_size, self.width // self.patch_size - shape = [self.batch_size, max_nb_patches_h * max_nb_patches_w] - return self.random_int_tensor( - shape, max_value=min(max_nb_patches_h, max_nb_patches_w), framework=framework, dtype=int_dtype - ) - return super().generate(input_name, framework, int_dtype, float_dtype) - - @register_in_tasks_manager("idefics3", *["image-text-to-text"], library_name="transformers") class Idefics3OpenVINOConfig(BaseVLMOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionPositionIdsInputGenerator) @@ -4917,7 +4073,7 @@ class SmolVLMOpenVINOConfig(Idefics3OpenVINOConfig): ], library_name="transformers", ) -class BlenderbotOpenVINOConfig(BlenderbotOnnxConfig): +class BlenderbotOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = BlenderbotModelPatcher @@ -4933,7 +4089,7 @@ class BlenderbotOpenVINOConfig(BlenderbotOnnxConfig): ], library_name="transformers", ) -class BlenderbotSmallOpenVINOConfig(BlenderbotSmallOnnxConfig): +class BlenderbotSmallOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = BlenderbotSmallModelPatcher @@ -4949,7 +4105,7 @@ class BlenderbotSmallOpenVINOConfig(BlenderbotSmallOnnxConfig): ], library_name="transformers", ) -class PegasusOpenVINOConfig(PegasusOnnxConfig): +class PegasusOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = PegasusModelPatcher def __init__(self, *args, **kwargs): @@ -4969,61 +4125,12 @@ def __init__(self, *args, **kwargs): ], library_name="transformers", ) -class MarianOpenVINOConfig(MarianOnnxConfig): +class MarianOpenVINOConfig(BartOpenVINOConfig): _MODEL_PATCHER = MarianModelPatcher # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" -class DummySpeechT5OpenVINOInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ( - "inputs_embeds", - "output_sequence", - "speaker_embeddings", - "spectrogram", - "raw_spectrogram", - "encoder_hidden_states", - ) - - def __init__( - self, - task: str, - normalized_config: NormalizedConfig, - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.task = task - self.batch_size = 1 - - self.sequence_length = sequence_length - self.speaker_embedding_dim = normalized_config.speaker_embedding_dim - self.num_mel_bins = normalized_config.num_mel_bins - self.reduction_factor = normalized_config.config.reduction_factor - self.hidden_size = normalized_config.config.hidden_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name in ["output_sequence", "inputs_embeds"]: - shape = [self.batch_size, self.sequence_length, self.num_mel_bins] - elif input_name == "speaker_embeddings": - shape = [self.batch_size, self.speaker_embedding_dim] - elif input_name == "raw_spectrogram": - shape = [self.sequence_length, self.batch_size, self.reduction_factor, self.num_mel_bins] - elif input_name == "encoder_hidden_states": - shape = [self.batch_size, self.sequence_length, self.hidden_size] - elif input_name == "spectrogram": - shape = [self.batch_size, self.sequence_length, self.num_mel_bins] - else: - raise ValueError(f"Unsupported input {input_name} for DummySpeechT5InputGenerator") - - return self.random_float_tensor( - shape=shape, - min_value=0, - max_value=1, - framework=framework, - dtype=float_dtype, - ) - - class SpeechT5ConfigBehavior(str, enum.Enum): ENCODER = "encoder" DECODER = "decoder" @@ -5036,13 +4143,27 @@ class SpeechT5ConfigBehavior(str, enum.Enum): *["text-to-audio", "text-to-audio-with-past"], library_name="transformers", ) -class SpeechT5OpenVINOConfig(SpeechT5OnnxConfig): +class SpeechT5OpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): DUMMY_INPUT_GENERATOR_CLASSES = ( DummyTextInputGenerator, DummySeq2SeqPastKeyValuesGenerator, - DummySpeechT5OpenVINOInputGenerator, + DummySpeechT5InputGenerator, + ) + _MODEL_PATCHER = SpeechT5ModelPatcher + + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + hidden_size="hidden_size", + num_attention_heads="encoder_attention_heads", + encoder_num_layers="encoder_layers", + decoder_num_layers="decoder_layers", + allow_new=True, ) - _MODEL_PATCHER = OVSpeechT5ModelPatcher + DUMMY_PKV_GENERATOR_CLASS = DummySeq2SeqPastKeyValuesGenerator + VARIANTS = { # noqa: RUF012 + "with-past": "The export follows the Transformers implementation using the KV cache.", + "without-past": "The same as `with-past`, just without KV cache support.", + } + DEFAULT_VARIANT = "with-past" def __init__( self, @@ -5064,9 +4185,10 @@ def __init__( use_past_in_inputs=use_past_in_inputs, behavior=behavior, preprocessors=preprocessors, - is_postnet_and_vocoder=False, ) + self.is_postnet_and_vocoder = False + def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): if direction not in ["inputs", "outputs"]: raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') @@ -5170,6 +4292,15 @@ def with_behavior( "self._behavior is neither encoder, decoder, postnet, or vocoder. This should not happen." ) + def overwrite_shape_and_generate_input( + self, dummy_input_gen: DummyInputGenerator, input_name: str, framework: str, input_shapes: dict + ): + dummy_input_gen.batch_size = 1 + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + return dummy_input + @register_in_tasks_manager( "llama4_text", *["text-generation", "text-generation-with-past"], library_name="transformers" @@ -5202,55 +4333,11 @@ def patch_model_for_export(self, model: PreTrainedModel, model_kwargs: Optional[ return Llama4ImageEmbeddingsModelPatcher(self, model, model_kwargs) -class MambaCacheDummyInputGenerator(DummyInputGenerator): - """ - Generates dummy past_ssm_states, past_conv_states and cache_position inputs for Mamba architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params", "cache_position") - - def __init__( - self, - task: str, - normalized_config, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.normalized_config = normalized_config - self.batch_size = batch_size - self.sequence_length = sequence_length - self.intermediate_size = self.normalized_config.config.intermediate_size - self.ssm_state_size = self.normalized_config.config.state_size - self.conv_kernel_size = self.normalized_config.config.conv_kernel - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "cache_params": - ssm_shape = [self.batch_size, self.intermediate_size, self.ssm_state_size] - conv_shape = [self.batch_size, self.intermediate_size, self.conv_kernel_size] - return [ - ( - self.random_float_tensor(ssm_shape, framework=framework, dtype=float_dtype), - self.random_float_tensor(conv_shape, framework=framework, dtype=float_dtype), - ) - for _ in range(self.normalized_config.num_layers) - ] - elif input_name == "cache_position": - return self.random_int_tensor( - shape=[self.conv_kernel_size], - max_value=self.sequence_length, - framework=framework, - dtype=int_dtype, - ) - - raise ValueError(f"Unsupported input name {input_name}") - - @register_in_tasks_manager( "falcon_mamba", *["text-generation", "text-generation-with-past"], library_name="transformers" ) @register_in_tasks_manager("mamba", *["text-generation", "text-generation-with-past"], library_name="transformers") -class MambaOpenVINOConfig(TextDecoderOnnxConfig): +class MambaOpenVINOConfig(TextDecoderOpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MambaCacheDummyInputGenerator) DUMMY_PKV_GENERATOR_CLASS = MambaCacheDummyInputGenerator NORMALIZED_CONFIG_CLASS = NormalizedTextConfig @@ -5317,7 +4404,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) return dummy_inputs @@ -5335,7 +4422,8 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): ], library_name="transformers", ) -class GPT2OpenVINOConfig(GPT2OnnxConfig): +class GPT2OpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_layers="n_layer", num_attention_heads="n_head") _MODEL_PATCHER = OVDecoderModelPatcher @@ -5348,84 +4436,33 @@ class GPT2OpenVINOConfig(GPT2OnnxConfig): "document-question-answering-with-past", ], ) -class VisionEncoderDecoderOpenVINOConfig(VisionEncoderDecoderOnnxConfig): +class VisionEncoderDecoderOpenVINOConfig(EncoderDecoderBaseOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator) _MODEL_PATCHER = OVSeq2SeqModelPatcher + @property + def inputs(self) -> dict: + common_inputs = {} -class Zamba2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy cache_params inputs for Zamba2 architectures. - """ + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["pixel_values"] = {0: "batch_size", 1: "num_channels", 2: "height", 3: "width"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} - SUPPORTED_INPUT_NAMES = ("cache_params",) + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") - def __init__( - self, - task: str, - normalized_config, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - **kwargs, - ) + return common_inputs - config = normalized_config.config - self.intermediate_size = int(config.mamba_expand * config.hidden_size) - self.conv_kernel_size = config.mamba_d_conv - self.mamba_d_state = config.mamba_d_state - if config.model_type == "zamba2": - self.n_mamba_heads = config.n_mamba_heads - self.mamba_ngroups = config.mamba_ngroups - self.mamba_headdim = config.mamba_headdim - self.head_dim = config.attention_head_dim - # in Zamba2, all layers contain Mamba block - # some of these layers are hybrid so they contain both attention and mamba blocks - self.num_attention_layers = len(config.hybrid_layer_ids) - self.num_mamba_layers = self.num_layers - logger.warning( - "The current support for the 'Zamba2' model type is experimental. " - "Performance is not optimal with high memory consumption. " - "Optimizations and improved support will be available in a future OpenVINO release." - ) + @property + def outputs(self) -> dict: + if self._behavior == ConfigBehavior.ENCODER: + return self._encoder_ov_config.outputs else: - # currently, this else-branch is applied for GraniteMoeHybrid models - self.n_mamba_heads = config.mamba_n_heads - self.mamba_ngroups = config.mamba_n_groups - self.mamba_headdim = config.mamba_d_head - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_attention_layers = config.layer_types.count("attention") - self.num_mamba_layers = config.layer_types.count("mamba") - self.num_attention_heads = config.num_key_value_heads - self.sequence_length = 0 - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_values = [] - for i in range(self.num_mamba_layers): - conv_state_shape = ( - self.batch_size, - self.intermediate_size + 2 * self.mamba_ngroups * self.mamba_d_state, - self.conv_kernel_size, - ) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - past_key_values.append(conv_state) - ssm_state_shape = (self.batch_size, self.n_mamba_heads, self.mamba_headdim, self.mamba_d_state) - ssm_state = self.random_float_tensor(ssm_state_shape, framework=framework, dtype=float_dtype) - past_key_values.append(ssm_state) - - for i in range(self.num_attention_layers): - kv_shape = (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - past_key_values.append(k) - past_key_values.append(v) - - return past_key_values + return super().outputs @register_in_tasks_manager("zamba2", *["text-generation", "text-generation-with-past"], library_name="transformers") @@ -5475,63 +4512,6 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return common_inputs -class Lfm2DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy past_key_values inputs for Lfm2 architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - **kwargs, - ) - config = normalized_config.config - self.num_conv_layers = config.layer_types.count("conv") - self.num_atten_layers = config.layer_types.count("full_attention") - self.batch_size = batch_size - self.normalized_config = normalized_config - self.hidden_size = self.normalized_config.hidden_size - self.conv_L_cache = self.normalized_config.conv_L_cache - self.num_key_value_heads = self.normalized_config.num_key_value_heads - self.num_hidden_layers = self.normalized_config.num_hidden_layers - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - past_key_values = [] - - for i in range(self.num_conv_layers): - conv_state_shape = (self.batch_size, self.hidden_size, self.conv_L_cache) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - past_key_values.append(conv_state) - - for i in range(self.num_atten_layers): - shape = ( - self.batch_size, - self.num_key_value_heads, - self.sequence_length, - self.hidden_size // self.num_attention_heads, - ) - - kv_shape = shape # (self.batch_size, self.num_attention_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - past_key_values.append(k) - past_key_values.append(v) - - return past_key_values - - @register_in_tasks_manager( "lfm2", *[ @@ -5630,8 +4610,15 @@ def inputs(self) -> Dict[str, Dict[int, str]]: @register_in_tasks_manager("audio-spectrogram-transformer", *["feature-extraction", "audio-classification"]) -class ASTOpenVINOConfig(ASTOnnxConfig): - pass +class ASTOpenVINOConfig(OpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + num_mel_bins="num_mel_bins", max_length="max_length", allow_new=True + ) + DUMMY_INPUT_GENERATOR_CLASSES = (ASTDummyAudioInputGenerator,) + + @property + def inputs(self) -> dict: + return {"input_values": {0: "batch_size"}} @register_in_tasks_manager( @@ -5652,12 +4639,19 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager("olmo2", *COMMON_TEXT_GENERATION_TASKS, library_name="transformers") -class Olmo2OOpenVINOConfig(Olmo2OnnxConfig): - pass +class Olmo2OOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, MistralDummyPastKeyValuesGenerator) + DUMMY_PKV_GENERATOR_CLASS = MistralDummyPastKeyValuesGenerator + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + MIN_TRANSFORMERS_VERSION = "4.47.0" @register_in_tasks_manager("opt", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "question-answering"]) -class OPTOpenVINOConfig(OPTOnnxConfig): +class OPTOpenVINOConfig( + TextDecoderWithPositionIdsOpenVINOConfig if is_transformers_version(">=", "4.46.0") else TextDecoderOpenVINOConfig +): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) _warn_potential_accuracy_issue_ov_2026_1("opt") @@ -5666,8 +4660,50 @@ def __init__(self, *args, **kwargs): @register_in_tasks_manager( "gpt_bigcode", *[*COMMON_TEXT_GENERATION_TASKS, "text-classification", "token-classification"] ) -class GPTBigCodeOpenVINOConfig(GPTBigCodeOnnxConfig): - pass +class GPTBigCodeOpenVINOConfig(TextDecoderWithPositionIdsOpenVINOConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GPTBigCodeDummyPastKeyValuesGenerator) + NORMALIZED_CONFIG_CLASS = NormalizedConfigManager.get_normalized_config_class("gpt_bigcode") + DUMMY_PKV_GENERATOR_CLASS = GPTBigCodeDummyPastKeyValuesGenerator + + def add_past_key_values(self, inputs_or_outputs: dict, direction: str): + if is_transformers_version(">=", "4.54"): + super().add_past_key_values(inputs_or_outputs, direction) + else: + if direction not in ["inputs", "outputs"]: + raise ValueError(f'direction must either be "inputs" or "outputs", but {direction} was given') + + if direction == "inputs": + decoder_sequence_name = "past_sequence_length" + name = "past_key_values" + else: + decoder_sequence_name = "past_sequence_length + sequence_length" + name = "present" + + if self._normalized_config.multi_query: + decoder_sequence_dim = 1 + else: + decoder_sequence_dim = 2 + + for i in range(self._normalized_config.num_layers): + inputs_or_outputs[f"{name}.{i}.key_value"] = { + 0: "batch_size", + decoder_sequence_dim: decoder_sequence_name, + } + + def flatten_past_key_values(self, flattened_output, name, idx, t): + if is_transformers_version(">=", "4.54"): + super().flatten_past_key_values(flattened_output, name, idx, t) + else: + flattened_output[f"{name}.{idx}.key_value"] = t + + +class _Pix2StructNormalizedConfig(NormalizedSeq2SeqConfig): + ENCODER_NUM_LAYERS = "vision_config.num_hidden_layers" + DECODER_NUM_LAYERS = "text_config.num_layers" + ENCODER_NUM_ATTENTION_HEADS = "vision_config.num_attention_heads" + DECODER_NUM_ATTENTION_HEADS = "text_config.num_heads" + HIDDEN_SIZE = "text_config.hidden_size" + VOCAB_SIZE = "text_config.vocab_size" @register_in_tasks_manager( @@ -5677,78 +4713,187 @@ class GPTBigCodeOpenVINOConfig(GPTBigCodeOnnxConfig): "image-to-text-with-past", ], ) -class Pix2StructOpenVINOConfig(Pix2StructOnnxConfig): +class Pix2StructOpenVINOConfig(OpenVINOSeq2SeqConfigWithPast): + PAD_ATTENTION_MASK_TO_PAST = True + NORMALIZED_CONFIG_CLASS = _Pix2StructNormalizedConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTextInputGenerator, + DummySeq2SeqDecoderTextInputGenerator, + DummySeq2SeqPastKeyValuesGenerator, + DummyPix2StructInputGenerator, + ) _MODEL_PATCHER = OVSeq2SeqModelPatcher + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if is_transformers_version("==", "4.46.0"): + logger.warning( + "Found transformers v4.46.0 while trying to export a Pix2Struct model, " + "this specific version of transformers is broken for this model. Please " + "upgrade to v4.46.1 or higher, or downgrade to v4.45.x.", + ) -@register_in_tasks_manager("bert", *COMMON_TEXT_TASKS) -class BertOpenVINOConfig(BertOnnxConfig): - pass + @property + def inputs(self): + common_inputs = {} + + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + common_inputs["flattened_patches"] = {0: "batch_size"} + else: + common_inputs["encoder_outputs"] = {0: "batch_size"} + common_inputs["attention_mask"] = {0: "batch_size"} + + if self._behavior in {ConfigBehavior.DECODER, ConfigBehavior.MONOLITH}: + common_inputs["decoder_input_ids"] = {0: "batch_size", 1: "decoder_sequence_length"} + if self.use_past_in_inputs: + self.add_past_key_values(common_inputs, direction="inputs") + decoder_attention_mask_dim = "past_decoder_sequence_length + decoder_sequence_length" + else: + decoder_attention_mask_dim = "decoder_sequence_length" + common_inputs["decoder_attention_mask"] = {0: "batch_size", 1: decoder_attention_mask_dim} + + return common_inputs + + @property + def outputs(self) -> dict: + common_outputs = super().outputs + if self._behavior in {ConfigBehavior.ENCODER, ConfigBehavior.MONOLITH}: + if self._behavior is ConfigBehavior.MONOLITH: + output_name = "encoder_last_hidden_state" + else: + output_name = "last_hidden_state" + common_outputs[output_name] = {0: "batch_size"} # Remove unnecessary dynamic axis. + + return common_outputs + + def _create_dummy_input_generator_classes(self, **kwargs): + if self._preprocessors is None or len(self._preprocessors) < 2: + raise ValueError( + f"Preprocessors for pix2struct need to be available for the OpenVINO export to infer input static shapes. Got: {self._preprocessors}" + ) + + dummy_inputs_generators = [] + dummy_inputs_generators.append( + self.DUMMY_INPUT_GENERATOR_CLASSES[0](self.task, self._normalized_config, **kwargs) + ) + encoder_sequence_length = self._preprocessors[1].image_processor.max_patches + kwargs["preprocessors"] = self._preprocessors + for cls_ in self.DUMMY_INPUT_GENERATOR_CLASSES[1:]: + dummy_inputs_generators.append( + cls_(self.task, self._normalized_config, encoder_sequence_length=encoder_sequence_length, **kwargs) + ) + + return dummy_inputs_generators + + def overwrite_shape_and_generate_input(self, dummy_input_gen, input_name: str, framework: str, input_shapes: dict): + if self._preprocessors is None or len(self._preprocessors) < 2: + raise ValueError( + f"Preprocessors for pix2struct need to be available for the OpenVINO export to infer input static shapes. Got: {self._preprocessors}" + ) + + if input_name in ["encoder_outputs", "attention_mask"]: + original_seq_length = dummy_input_gen.sequence_length + dummy_input_gen.sequence_length = self._preprocessors[1].image_processor.max_patches + dummy_input = dummy_input_gen.generate( + input_name, framework=framework, int_dtype=self.int_dtype, float_dtype=self.float_dtype + ) + dummy_input_gen.sequence_length = original_seq_length + else: + dummy_input = super().overwrite_shape_and_generate_input( + dummy_input_gen, input_name, framework, input_shapes + ) + + return dummy_input + + +@register_in_tasks_manager("bert", *COMMON_TEXT_TASKS) +class BertOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + "token_type_ids": dynamic_axis, + } @register_in_tasks_manager("albert", *COMMON_TEXT_TASKS) -class AlbertOpenVINOConfig(AlbertOnnxConfig): +class AlbertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("nystromformer", *COMMON_TEXT_TASKS) -class NystromformerOpenVINOConfig(NystromformerOnnxConfig): +class NystromformerOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("convbert", *COMMON_TEXT_TASKS) -class ConvBertOpenVINOConfig(ConvBertOnnxConfig): +class ConvBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("electra", *COMMON_TEXT_TASKS) -class ElectraOpenVINOConfig(ElectraOnnxConfig): +class ElectraOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("roformer", *COMMON_TEXT_TASKS) -class RoFormerOpenVINOConfig(RoFormerOnnxConfig): +class RoFormerOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("squeezebert", *COMMON_TEXT_TASKS) -class SqueezeBertOpenVINOConfig(SqueezeBertOnnxConfig): +class SqueezeBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("mobilebert", *COMMON_TEXT_TASKS) -class MobileBertOpenVINOConfig(MobileBertOnnxConfig): +class MobileBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("xlm", *COMMON_TEXT_TASKS) -class XLMOpenVINOConfig(XLMOnnxConfig): +class XLMOpenVINOConfig(BertOpenVINOConfig): # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" -@register_in_tasks_manager("xlm-roberta", *COMMON_TEXT_TASKS) -class XLMRobertaOpenVINOConfig(XLMRobertaOnnxConfig): - pass +@register_in_tasks_manager("distilbert", *COMMON_TEXT_TASKS) +class DistilBertOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + @property + def inputs(self) -> dict: + if self.task == "multiple-choice": + dynamic_axis = {0: "batch_size", 1: "num_choices", 2: "sequence_length"} + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} -@register_in_tasks_manager("distilbert", *COMMON_TEXT_TASKS) -class DistilBertOpenVINOConfig(DistilBertOnnxConfig): + +@register_in_tasks_manager("xlm-roberta", *COMMON_TEXT_TASKS) +class XLMRobertaOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("roberta", *COMMON_TEXT_TASKS) -class RobertaOpenVINOConfig(RobertaOnnxConfig): +class RobertaOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("camembert", *COMMON_TEXT_TASKS) -class CamembertOpenVINOConfig(CamembertOnnxConfig): +class CamembertOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("flaubert", *COMMON_TEXT_TASKS) -class FlaubertOpenVINOConfig(FlaubertOnnxConfig): +class FlaubertOpenVINOConfig(BertOpenVINOConfig): # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" @@ -5757,15 +4902,39 @@ class FlaubertOpenVINOConfig(FlaubertOnnxConfig): "deberta", *["feature-extraction", "fill-mask", "text-classification", "token-classification", "question-answering"], ) -class DebertaOpenVINOConfig(DebertaOnnxConfig): - pass +class DebertaOpenVINOConfig(BertOpenVINOConfig): + @property + def inputs(self) -> dict: + common_inputs = super().inputs + if self._config.type_vocab_size == 0: + common_inputs.pop("token_type_ids") + return common_inputs @register_in_tasks_manager("deberta-v2", *COMMON_TEXT_TASKS) -class DebertaV2OpenVINOConfig(DebertaV2OnnxConfig): +class DebertaV2OpenVINOConfig(DebertaOpenVINOConfig): pass +@register_in_tasks_manager("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) +class HubertOpenVINOConfig(AudioOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def outputs(self) -> dict: + outputs = super().outputs + + # Hubert output formula adapted from: + # https://github.com/huggingface/transformers/blob/v4.55.2/src/transformers/models/hubert/modeling_hubert.py#L721 + if self.task == "automatic-speech-recognition": + sequence_length = "sequence_length" + for kernel_size, stride in zip(self._config.conv_kernel, self._config.conv_stride): + sequence_length = f"( {sequence_length} - {kernel_size} ) // {stride} + 1" + outputs["logits"] = {0: "batch_size", 1: sequence_length} + + return outputs + + @register_in_tasks_manager( "data2vec-audio", *[ @@ -5776,95 +4945,190 @@ class DebertaV2OpenVINOConfig(DebertaV2OnnxConfig): "audio-xvector", ], ) -class Data2VecAudioOpenVINOConfig(Data2VecAudioOnnxConfig): +class Data2VecAudioOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("data2vec-text", *COMMON_TEXT_TASKS) -class Data2VecTextOpenVINOConfig(Data2VecTextOnnxConfig): +class Data2VecTextOpenVINOConfig(DistilBertOpenVINOConfig): # TODO (@echarlaix): add v5 support MAX_TRANSFORMERS_VERSION = "4.57.6" +@register_in_tasks_manager("vit", *["feature-extraction", "image-classification", "masked-im"]) +class ViTOpenVINOConfig(VisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig + + @property + def inputs(self) -> dict: + return {"pixel_values": {0: "batch_size", 2: "height", 3: "width"}} + + @property + def outputs(self) -> dict: + common_outputs = super().outputs + + if self.task == "feature-extraction": + common_outputs["last_hidden_state"] = {0: "batch_size"} + + return common_outputs + + @register_in_tasks_manager("data2vec-vision", *["feature-extraction", "image-classification"]) -class Data2VecVisionOpenVINOConfig(Data2VecVisionOnnxConfig): +class Data2VecVisionOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("perceiver", *["fill-mask", "text-classification", "image-classification"]) -class PerceiverOpenVINOConfig(PerceiverOnnxConfig): - pass +class PerceiverOpenVINOConfig(TextAndVisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + DUMMY_INPUT_GENERATOR_CLASSES = ( + PerceiverDummyInputGenerator, + *TextAndVisionOpenVINOConfig.DUMMY_INPUT_GENERATOR_CLASSES, + ) + + def __init__( + self, + config, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + preprocessors=None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.is_generating_dummy_inputs = False + + @property + def inputs_name(self): + if self.is_generating_dummy_inputs: + if self.task in ["fill-mask", "text-classification"]: + return "input_ids" + else: + return "pixel_values" + else: + return "inputs" + + @property + def inputs(self) -> dict: + if self.inputs_name in ["input_ids", "inputs"]: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } + else: + dynamic_axis = {0: "batch_size", 1: "sequence_length", 2: "width", 3: "height"} + return { + "pixel_values": dynamic_axis, + } + + @property + def outputs(self) -> dict: + outputs = super().outputs + + if "logits" in outputs: + outputs["logits"] = {0: "batch_size"} + + return outputs + + def generate_dummy_inputs(self, framework: str = "pt", **kwargs): + self.is_generating_dummy_inputs = True + dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) + dummy_inputs[self.inputs_name] = dummy_inputs.pop(self.inputs_name) + return dummy_inputs @register_in_tasks_manager("esm", *["feature-extraction", "fill-mask", "text-classification", "token-classification"]) -class EsmOpenVINOConfig(EsmOnnxConfig): - pass +class EsmOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + + @property + def inputs(self) -> dict: + dynamic_axis = {0: "batch_size", 1: "sequence_length"} + return { + "input_ids": dynamic_axis, + "attention_mask": dynamic_axis, + } @register_in_tasks_manager("mpnet", *COMMON_TEXT_TASKS) -class MPNetOpenVINOConfig(MPNetOnnxConfig): +class MPNetOpenVINOConfig(DistilBertOpenVINOConfig): pass @register_in_tasks_manager("beit", *["feature-extraction", "image-classification"]) -class BeitOpenVINOConfig(BeitOnnxConfig): +class BeitOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("deit", *["feature-extraction", "image-classification", "masked-im"]) -class DeiTOpenVINOConfig(DeiTOnnxConfig): +class DeiTOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("levit", *["feature-extraction", "image-classification"]) -class LevitOpenVINOConfig(LevitOnnxConfig): +class LevitOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("mobilevit", *["feature-extraction", "image-classification", "image-segmentation"]) -class MobileViTOpenVINOConfig(MobileViTOnnxConfig): +class MobileViTOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("mobilenet_v1", *["feature-extraction", "image-classification"]) -class MobileNetV1OpenVINOConfig(MobileNetV1OnnxConfig): - pass +class MobileNetV1OpenVINOConfig(ViTOpenVINOConfig): + @property + def inputs(self) -> dict: + return {"pixel_values": {0: "batch_size"}} @register_in_tasks_manager("mobilenet_v2", *["feature-extraction", "image-classification"]) -class MobileNetV2OpenVINOConfig(MobileNetV2OnnxConfig): +class MobileNetV2OpenVINOConfig(MobileNetV1OpenVINOConfig): pass @register_in_tasks_manager("poolformer", *["feature-extraction", "image-classification"]) -class PoolFormerOpenVINOConfig(PoolFormerOnnxConfig): - pass +class PoolFormerOpenVINOConfig(ViTOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @register_in_tasks_manager( "segformer", *["feature-extraction", "image-classification", "image-segmentation", "semantic-segmentation"] ) -class SegformerOpenVINOConfig(SegformerOnnxConfig): - pass +class SegformerOpenVINOConfig(ViTOpenVINOConfig): + @property + def outputs(self) -> dict: + outputs = super().outputs + + if self.task == "image-segmentation": + outputs["logits"] = {0: "batch_size"} + + return outputs @register_in_tasks_manager("swin", *["feature-extraction", "image-classification", "masked-im"]) -class SwinOpenVINOConfig(SwinOnnxConfig): +class SwinOpenVINOConfig(ViTOpenVINOConfig): pass -@register_in_tasks_manager("vit", *["feature-extraction", "image-classification", "masked-im"]) -class ViTOpenVINOConfig(ViTOnnxConfig): +@register_in_tasks_manager("donut-swin", *["feature-extraction"]) +class DonutSwinOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("convnext", *["feature-extraction", "image-classification"]) -class ConvNextOpenVINOConfig(ConvNextOnnxConfig): +class ConvNextOpenVINOConfig(ViTOpenVINOConfig): pass @register_in_tasks_manager("resnet", *["feature-extraction", "image-classification"]) -class ResNetOpenVINOConfig(ResNetOnnxConfig): +class ResNetOpenVINOConfig(ViTOpenVINOConfig): pass @@ -5878,7 +5142,7 @@ class ResNetOpenVINOConfig(ResNetOnnxConfig): "audio-xvector", ], ) -class Wav2Vec2OpenVINOConfig(Wav2Vec2OnnxConfig): +class Wav2Vec2OpenVINOConfig(HubertOpenVINOConfig): pass @@ -5892,29 +5156,24 @@ class Wav2Vec2OpenVINOConfig(Wav2Vec2OnnxConfig): "audio-xvector", ], ) -class Wav2Vec2ConformerOpenVINOConfig(Wav2Vec2ConformerOnnxConfig): - pass - - -@register_in_tasks_manager("hubert", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class HubertOpenVINOConfig(HubertOnnxConfig): +class Wav2Vec2ConformerOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("sew", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class SEWOpenVINOConfig(SEWOnnxConfig): +class SEWOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("sew-d", *["feature-extraction", "automatic-speech-recognition", "audio-classification"]) -class SEWDOpenVINOConfig(SEWDOnnxConfig): +class SEWDOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager( "unispeech", *["feature-extraction", "automatic-speech-recognition", "audio-classification"] ) -class UniSpeechOpenVINOConfig(UniSpeechOnnxConfig): +class UniSpeechOpenVINOConfig(HubertOpenVINOConfig): pass @@ -5928,7 +5187,7 @@ class UniSpeechOpenVINOConfig(UniSpeechOnnxConfig): "audio-xvector", ], ) -class UniSpeechSATOpenVINOConfig(UniSpeechSATOnnxConfig): +class UniSpeechSATOpenVINOConfig(HubertOpenVINOConfig): pass @@ -5942,117 +5201,174 @@ class UniSpeechSATOpenVINOConfig(UniSpeechSATOnnxConfig): "audio-xvector", ], ) -class WavLMOpenVINOConfig(WavLMOnnxConfig): +class WavLMOpenVINOConfig(HubertOpenVINOConfig): pass @register_in_tasks_manager("sam", *["feature-extraction"]) -class SamOpenVINOConfig(SamOnnxConfig): - pass +class SamOpenVINOConfig(OpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyVisionInputGenerator, DummyPointsGenerator, DummyVisionEmbeddingsGenerator) + VARIANTS = { # noqa: RUF012 + "monolith": "All the SAM model components are exported as a single model.", + "split": "The vision encoder is exported as a separate vision_encoder, and the prompt encoder and mask decoder are exported as a prompt_encoder_mask_decoder. This allows to encoder the image only once for multiple point queries.", + } + DEFAULT_VARIANT = "split" + _MODEL_PATCHER = SAMModelPatcher + + def __init__( + self, + config, + task: str = "feature-extraction", + int_dtype: str = "int64", + float_dtype: str = "fp32", + variant: str = "split", + vision_encoder=None, + preprocessors=None, + ): + super().__init__( + config=config, + task=task, + int_dtype=int_dtype, + float_dtype=float_dtype, + preprocessors=preprocessors, + ) + self.variant = variant + self.vision_encoder = vision_encoder + self._normalized_config.ENCODER_NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig(self._config.vision_config) + + @property + def inputs(self) -> dict: + if self.variant == "monolith": + inputs = { + "pixel_values": {0: "batch_size"}, + "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + } + else: + if self.vision_encoder: + inputs = {"pixel_values": {0: "batch_size"}} + else: + inputs = { + "image_positional_embeddings": {0: "batch_size"}, + "image_embeddings": {0: "batch_size"}, + "input_points": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + "input_labels": {0: "batch_size", 1: "point_batch_size", 2: "nb_points_per_image"}, + } + return inputs + + @property + def outputs(self) -> dict: + if self.variant == "split" and self.vision_encoder: + return {"image_embeddings": {0: "batch_size"}, "image_positional_embeddings": {0: "batch_size"}} + else: + return { + "iou_scores": {0: "batch_size", 1: "point_batch_size"}, + "pred_masks": {0: "batch_size", 1: "point_batch_size"}, + } @register_in_tasks_manager("siglip", *["feature-extraction", "zero-shot-image-classification"]) -class SiglipOpenVINOConfig(SiglipOnnxConfig): - pass +class SiglipOpenVINOConfig(TextAndVisionOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = SiglipNormalizedConfig + _MODEL_PATCHER = CLIPModelPatcher + + @property + def inputs(self) -> dict: + return { + "input_ids": {0: "text_batch_size", 1: "sequence_length"}, + "pixel_values": {0: "image_batch_size", 1: "num_channels", 2: "height", 3: "width"}, + # NOTE: No attention_mask + } + + @property + def outputs(self) -> dict: + if self.task in ["feature-extraction", "zero-shot-image-classification"]: + return { + "logits_per_image": {0: "image_batch_size", 1: "text_batch_size"}, + "logits_per_text": {0: "text_batch_size", 1: "image_batch_size"}, + "text_embeds": {0: "text_batch_size"}, + "image_embeds": {0: "image_batch_size"}, + } + else: + return super().outputs @register_in_tasks_manager( "transformer", *["feature-extraction", "sentence-similarity"], library_name="sentence_transformers" ) -class SentenceTransformersTransformerOpenVINOConfig(SentenceTransformersTransformerOnnxConfig): - pass +class SentenceTransformersTransformerOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + _MODEL_PATCHER = SentenceTransformersTransformerPatcher + + @property + def inputs(self) -> dict: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> dict: + return { + "token_embeddings": {0: "batch_size", 1: "sequence_length"}, + "sentence_embedding": {0: "batch_size"}, + } @register_in_tasks_manager("rembert", *COMMON_TEXT_TASKS) -class RemBertOpenVINOConfig(RemBertOnnxConfig): +class RemBertOpenVINOConfig(BertOpenVINOConfig): pass @register_in_tasks_manager("siglip-text-with-projection", *["feature-extraction"]) -class SiglipTextWithProjectionOpenVINOConfig(SiglipTextWithProjectionOnnxConfig): - pass - +class SiglipTextWithProjectionOpenVINOConfig(TextEncoderOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + vocab_size="vocab_size", + sequence_length="max_position_embeddings", + num_layers="num_hidden_layers", + allow_new=True, + ) + _MODEL_PATCHER = CLIPModelPatcher -@register_in_tasks_manager("siglip-text", *["feature-extraction"]) -class SiglipTextOpenVINOConfig(SiglipTextOnnxConfig): - pass + @property + def inputs(self) -> dict: + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + @property + def outputs(self) -> dict: + common_outputs = { + "text_embeds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} -class DummyVideoChatFlashQwenInputGenerator(DummyVisionInputGenerator): - SUPPORTED_INPUT_NAMES = ("hidden_states", "rotary_pos_emb") + return common_outputs - def __init__( - self, - task: str, - normalized_config: NormalizedVisionConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - num_channels: int = DEFAULT_DUMMY_SHAPES["num_channels"], - width: int = DEFAULT_DUMMY_SHAPES["width"], - height: int = DEFAULT_DUMMY_SHAPES["height"], - visual_seq_length: int = DEFAULT_DUMMY_SHAPES["visual_seq_length"], - **kwargs, - ): - super().__init__(task, normalized_config, batch_size, num_channels, width, height, visual_seq_length, **kwargs) - self.num_frames = normalized_config.config.mm_local_num_frames - self.embed_dim = normalized_config.config.mm_hidden_size - self.height = normalized_config.config.image_size - self.width = normalized_config.config.image_size - self.image_size = (self.height, self.width) - self.patch_size = normalized_config.config.patch_size - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "hidden_states": - return self.random_float_tensor( - shape=[ - self.batch_size, - self.num_channels, - self.num_frames, - self.height, - self.width, - ], - framework=framework, - dtype=float_dtype, - ) - elif input_name == "rotary_pos_emb": - grid_h, grid_w = self.height // self.patch_size, self.width // self.patch_size - grid_t = self.num_frames - # Source: https://huggingface.co/OpenGVLab/VideoChat-Flash-Qwen2_5-7B_InternVideo2-1B/blob/main/vision_tower_builder.py#L523 - # The first dimension of rotary_pos_emb is fixed to 1 in the original model. - # And the second dimension is the total number of tokens for all frames, which is calculated as grid_h * grid_w * grid_t plus 1 for the cls token. - return self.random_float_tensor( - [1, 1 + grid_h * grid_t * grid_w, self.embed_dim], framework=framework, dtype=float_dtype - ) +@register_in_tasks_manager("siglip-text", *["feature-extraction"]) +class SiglipTextOpenVINOConfig(SiglipTextWithProjectionOpenVINOConfig): + _MODEL_PATCHER = CLIPModelPatcher -class DummyVideoChatFlashQwenProjectorInputGenerator(DummyInputGenerator): - SUPPORTED_INPUT_NAMES = ["input"] + @property + def outputs(self) -> dict: + common_outputs = { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + "pooler_output": {0: "batch_size"}, + } - def __init__( - self, - task: str, - normalized_config: NormalizedTextConfig, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - random_batch_size_range: Optional[Tuple[int, int]] = None, - **kwargs, - ): - self.task = task - self.batch_size = batch_size - self.hidden_size = normalized_config.config.mm_hidden_size - self.num_patches = normalized_config.config.mm_projector_num_tome_tokens - self.normalized_config = normalized_config + if self._normalized_config.output_hidden_states: + for i in range(self._normalized_config.num_layers + 1): + common_outputs[f"hidden_states.{i}"] = {0: "batch_size", 1: "sequence_length"} - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - ): - shape = [self.batch_size, self.num_patches, self.hidden_size] - return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + return common_outputs -class VideoChatFlashQwenProjectorOpenVINOConfig(OnnxConfig): +class VideoChatFlashQwenProjectorOpenVINOConfig(OpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyVideoChatFlashQwenProjectorInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig @@ -6204,68 +5520,6 @@ class HunyuanV1DenseOpenVINOConfig(LlamaOpenVINOConfig): DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator -class Qwen3NextDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy cache_params inputs for Qwen3-Next architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - **kwargs, - ) - - config = normalized_config.config - self.num_full_attn_layers = config.layer_types.count("full_attention") - self.num_linear_attn_layers = config.layer_types.count("linear_attention") - self.conv_kernel_size = config.linear_conv_kernel_dim - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads - self.head_k_dim = config.linear_key_head_dim - self.head_v_dim = config.linear_value_head_dim - self.num_v_heads = config.linear_num_value_heads - self.num_k_heads = config.linear_num_key_heads - self.num_key_value_heads = config.num_key_value_heads - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - cache_params = [] - - for idx in range(self.num_linear_attn_layers): - d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) - conv_state_shape = ( - self.batch_size, - d_inner, - self.conv_kernel_size, - ) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(conv_state) - num_heads = self.num_v_heads - recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) - recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(recurrent_state) - - for idx in range(self.num_full_attn_layers): - kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - cache_params.append(k) - cache_params.append(v) - - return cache_params - - @register_in_tasks_manager( "qwen3_next", *["text-generation", "text-generation-with-past"], @@ -6335,7 +5589,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) return dummy_inputs @@ -6354,67 +5608,6 @@ class LFM2MoeOpenVINOConfig(LFM2OpenVINOConfig): _MODEL_PATCHER = Lfm2MoeModelPatcher -class Qwen3_5DummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): - """ - Generates dummy cache_params inputs for Qwen3.5 architectures. - """ - - SUPPORTED_INPUT_NAMES = ("cache_params",) - - def __init__( - self, - task: str, - normalized_config, - batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - super().__init__( - task=task, - normalized_config=normalized_config, - batch_size=batch_size, - sequence_length=sequence_length, - **kwargs, - ) - - config = normalized_config.config - self.num_full_attn_layers = config.layer_types.count("full_attention") - self.num_linear_attn_layers = config.layer_types.count("linear_attention") - self.conv_kernel_size = config.linear_conv_kernel_dim - self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) - self.head_k_dim = config.linear_key_head_dim - self.head_v_dim = config.linear_value_head_dim - self.num_v_heads = config.linear_num_value_heads - self.num_k_heads = config.linear_num_key_heads - self.num_key_value_heads = config.num_key_value_heads - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - cache_params = [] - - for idx in range(self.num_linear_attn_layers): - d_inner = self.num_k_heads * (2 * self.head_k_dim + self.head_v_dim * self.num_v_heads // self.num_k_heads) - conv_state_shape = ( - self.batch_size, - d_inner, - self.conv_kernel_size, - ) - conv_state = self.random_float_tensor(conv_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(conv_state) - num_heads = self.num_v_heads - recurrent_state_shape = (self.batch_size, num_heads, self.head_k_dim, self.head_v_dim) - recurrent_state = self.random_float_tensor(recurrent_state_shape, framework=framework, dtype=float_dtype) - cache_params.append(recurrent_state) - - for idx in range(self.num_full_attn_layers): - kv_shape = (self.batch_size, self.num_key_value_heads, self.sequence_length, self.head_dim) - k = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - v = self.random_float_tensor(kv_shape, framework=framework, dtype=float_dtype) - cache_params.append(k) - cache_params.append(v) - - return cache_params - - @register_in_tasks_manager( "qwen3_5_text", *["text-generation", "text-generation-with-past"], @@ -6483,7 +5676,7 @@ def generate_dummy_inputs(self, framework: str = "pt", **kwargs): break if not input_was_inserted: raise RuntimeError( - f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model ONNX config.' + f'Could not generate dummy input for "{input_name}". Try adding a proper dummy input generator to the model OpenVINO config.' ) return dummy_inputs @@ -6649,50 +5842,12 @@ def outputs(self) -> Dict[str, Dict[int, str]]: return super().outputs -class DummyKokoroInputGenerator(DummyInputGenerator): - """Generates dummy inputs for the Kokoro TTS model.""" - - SUPPORTED_INPUT_NAMES = ("input_ids", "ref_s", "speed") - - def __init__( - self, - task: str, - normalized_config: NormalizedConfig, - sequence_length: int = DEFAULT_DUMMY_SHAPES["sequence_length"], - **kwargs, - ): - self.task = task - self.batch_size = 1 - self.sequence_length = sequence_length - self.style_dim = getattr(normalized_config, "style_dim", 128) - - def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if input_name == "input_ids": - shape = [self.batch_size, self.sequence_length] - input_ids_value = self.random_int_tensor( - shape=shape, min_value=0, max_value=178, framework=framework, dtype=int_dtype - ) - input_ids_value[:, 0] = 0 - input_ids_value[:, -1] = 0 - return input_ids_value - elif input_name == "ref_s": - shape = [self.batch_size, self.style_dim * 2] - return self.random_float_tensor( - shape=shape, min_value=-1, max_value=1, framework=framework, dtype=float_dtype - ) - elif input_name == "speed": - return self.random_int_tensor(shape=[1], min_value=1, max_value=10, framework=framework, dtype=float_dtype) - else: - raise ValueError(f"Unsupported input {input_name} for DummyKokoroInputGenerator") - - @register_in_tasks_manager( "kokoro", *["text-to-audio"], library_name="kokoro", ) -class KokoroOpenVINOConfig(OnnxConfig): - DEFAULT_ONNX_OPSET = 14 +class KokoroOpenVINOConfig(OpenVINOConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyKokoroInputGenerator,) NORMALIZED_CONFIG_CLASS = NormalizedConfig _MODEL_PATCHER = KokoroModelPatcher @@ -6727,3 +5882,19 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "waveform": {0: "batch_size", 1: "audio_length"}, "phonemes": {0: "batch_size", 1: "phoneme_length"}, } + + +@register_in_tasks_manager( + "trocr", + *[ + "feature-extraction", + "feature-extraction-with-past", + ], +) +class TrOCROpenVINOConfig(TextSeq2SeqOpenVINOConfig): + NORMALIZED_CONFIG_CLASS = NormalizedSeq2SeqConfig.with_args( + decoder_num_layers="decoder_layers", + num_layers="decoder_layers", + decoder_num_attention_heads="decoder_attention_heads", + hidden_size="hidden_size", + ) diff --git a/optimum/exporters/openvino/model_patcher.py b/optimum/exporters/openvino/model_patcher.py index 6fcc9de4de..199bbce903 100644 --- a/optimum/exporters/openvino/model_patcher.py +++ b/optimum/exporters/openvino/model_patcher.py @@ -19,11 +19,13 @@ import math import types from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch import torch.nn.functional as F +import transformers from torch import nn +from transformers import PreTrainedModel from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache from transformers.configuration_utils import PretrainedConfig from transformers.generation import GenerationMixin @@ -32,7 +34,6 @@ BaseModelOutputWithPast, BaseModelOutputWithPooling, ) -from transformers.modeling_utils import PreTrainedModel from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import ( LlamaAttention, @@ -46,15 +47,14 @@ from transformers.processing_utils import Unpack from transformers.utils import ModelOutput -from optimum.exporters.onnx.base import OnnxConfig -from optimum.exporters.onnx.model_patcher import ( - UNSUPPORTED_OPS_PATCHING_SPEC, +from optimum.exporters.openvino._ov_ops import convert_recurrent_attention_cell +from optimum.exporters.openvino.base import OpenVINOConfig +from optimum.exporters.openvino.patching_utils import ( ModelPatcher, - gpt_oss_forward, + eager_mask_without_vmap, override_arguments, -) -from optimum.exporters.onnx.model_patcher import ( - sdpa_mask_without_vmap as _orig_sdpa_mask_without_vmap, + postprocess_past_key_values, + preprocess_past_key_values, ) from optimum.intel.utils.import_utils import ( is_diffusers_version, @@ -63,93 +63,174 @@ is_transformers_version, ) -from ._ov_ops import convert_recurrent_attention_cell + +if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention +if is_transformers_version(">=", "4.48"): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache if is_transformers_version(">=", "4.53"): - from transformers.masking_utils import ALL_MASK_ATTENTION_FUNCTIONS, eager_mask, sdpa_mask + from transformers.masking_utils import ( + ALL_MASK_ATTENTION_FUNCTIONS, + eager_mask, + sdpa_mask, + ) from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock + + if is_transformers_version(">=", "4.54"): from transformers.masking_utils import create_causal_mask + from transformers.utils import TransformersKwargs +else: + TransformersKwargs = object + + if is_transformers_version(">=", "4.56"): import transformers.masking_utils + + if is_transformers_version(">=", "4.57"): from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding + + if is_transformers_version(">=", "5"): from transformers.modeling_rope_utils import RotaryEmbeddingConfigMixin -if TYPE_CHECKING: - from transformers.cache_utils import Cache - from transformers.modeling_utils import PreTrainedModel - from optimum.exporters.onnx.config import OnnxConfig +logger = logging.getLogger(__name__) -if is_transformers_version(">=", "4.54"): - from transformers.utils import TransformersKwargs -else: - TransformersKwargs = object +class SAMModelPatcher(ModelPatcher): + def __init__( + self, + config: "OpenVINOConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + super().__init__(config, model, model_kwargs) + + def patched_forward( + pixel_values=None, + input_points=None, + input_labels=None, + image_embeddings=None, + image_positional_embeddings=None, + return_dict=True, + **kwargs, + ): + if config.variant == "monolith": + return self.orig_forward( + pixel_values=pixel_values, + input_points=input_points, + input_labels=input_labels, + image_embeddings=image_embeddings, + return_dict=return_dict, + **kwargs, + ) + elif config.variant == "split": + # return_dict = get_argument(args, kwargs, signature, "return_dict") + if config.vision_encoder: + # pixel_values = get_argument(args, kwargs, signature, "pixel_values") + image_positional_embeddings = model.get_image_wide_positional_embeddings() -logger = logging.getLogger(__name__) + # repeat with batch size + batch_size = pixel_values.shape[0] + image_positional_embeddings = image_positional_embeddings.repeat(batch_size, 1, 1, 1) + vision_outputs = model.vision_encoder( + pixel_values, + output_attentions=False, + output_hidden_states=False, + return_dict=return_dict, + ) + image_embeddings = vision_outputs[0] -def postprocess_past_key_values(past_key_values): - if isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): - if hasattr(past_key_values, "to_legacy_cache"): - past_key_values = past_key_values.to_legacy_cache() - elif isinstance(past_key_values, DynamicCache): - past_key_values = [(lay.keys, lay.values) for lay in past_key_values.layers] - elif isinstance(past_key_values, EncoderDecoderCache): - past_key_values = [ - (self_lay.keys, self_lay.values, cross_lay.keys, cross_lay.values) - for self_lay, cross_lay in zip( - past_key_values.self_attention_cache.layers, - past_key_values.cross_attention_cache.layers, - ) - ] - return past_key_values + if not return_dict: + return (image_embeddings, image_positional_embeddings) + else: + return { + "image_embeddings": image_embeddings, + "image_positional_embeddings": image_positional_embeddings, + } + else: + if input_points is None: + raise ValueError("input_points is required to export the prompt encoder / mask decoder.") + + sparse_embeddings, dense_embeddings = model.prompt_encoder( + input_points=input_points, + input_labels=input_labels, + input_boxes=None, # Not supported in the OpenVINO export + input_masks=None, # Not supported in the OpenVINO export + ) + outputs = model.mask_decoder( + image_embeddings=image_embeddings, + image_positional_embeddings=image_positional_embeddings, + sparse_prompt_embeddings=sparse_embeddings, + dense_prompt_embeddings=dense_embeddings, + multimask_output=True, # Not supported in the OpenVINO export + attention_similarity=None, # Not supported in the OpenVINO export + target_embedding=None, # Not supported in the OpenVINO export + ) + low_res_masks, iou_predictions = outputs[:2] + if not return_dict: + return (iou_predictions, low_res_masks) + else: + return {"iou_scores": iou_predictions, "pred_masks": low_res_masks} -def _get_model_attribute(model, name): - target = getattr(model, "model", model) if is_transformers_version(">=", "5") else model - return getattr(target, name) + self.patched_forward = patched_forward -# Compatibility wrapper for sdpa_mask_without_vmap from optimum. -# The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...), -# but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...). -def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs): - import inspect - - sig = inspect.signature(_orig_sdpa_mask_without_vmap) - if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None: - # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...) - cache_position = torch.arange(q_length, dtype=torch.long) + q_offset - kwargs.pop("q_offset", None) - kwargs.pop("allow_is_bidirectional_skip", None) - kwargs.pop("allow_torch_fix", None) - kwargs.pop("use_vmap", None) - kwargs.pop("device", None) - return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs) - else: - return _orig_sdpa_mask_without_vmap( - batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs - ) +class SentenceTransformersTransformerPatcher(ModelPatcher): + def __init__( + self, + config: "OpenVINOConfig", + model: PreTrainedModel, + model_kwargs: dict[str, Any], + ): + super().__init__(config, model, model_kwargs) + + def patched_forward(input_ids, attention_mask): + result = self.orig_forward({"input_ids": input_ids, "attention_mask": attention_mask}) + + if "input_ids" in result: + del result["input_ids"] + if "attention_mask" in result: + del result["attention_mask"] + if "all_layer_embeddings" in result: + del result["all_layer_embeddings"] + + return result + + self.patched_forward = patched_forward + + +class CLIPModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + self.original_sdpa_forward = CLIPSdpaAttention.forward + CLIPSdpaAttention.forward = CLIPAttention.forward + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + if is_transformers_version(">=", "4.43") and is_transformers_version("<", "4.48"): + CLIPSdpaAttention.forward = self.original_sdpa_forward -for idx, spec in enumerate(UNSUPPORTED_OPS_PATCHING_SPEC): - if spec.name in { - # onnx-exporter-specific fixes - "triu", - "tril", - "norm", - "unfold", - "movedim", - "rms_norm", - "repeat_interleave", - "scaled_dot_product_attention", - }: - UNSUPPORTED_OPS_PATCHING_SPEC.pop(idx) +# This is a traceable version of the original function, +# the original results in a constant integer due to the use of int(expr) +def _get_feat_extract_output_lengths_patched(self, input_lengths: torch.LongTensor): + output_conv1_length = (input_lengths - 127) // 64 + 1 + output_conv2_length = (output_conv1_length - 7) // 3 + 1 + output_conv3_length = (output_conv2_length - 3) // 2 + 1 + return output_conv3_length + + +def _get_model_attribute(model, name): + target = getattr(model, "model", model) if is_transformers_version(">=", "5") else model + return getattr(target, name) # Original code: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/integrations/moe.py#L98 @@ -376,22 +457,6 @@ def patch_cos_sin_cached_fp32(model): ) -# Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 -# Specifically for OpenVINO, we use torch.finfo(torch.float16).min instead of torch.finfo(dtype).min -def eager_mask_without_vmap(*args, **kwargs) -> Optional[torch.Tensor]: - kwargs.pop("allow_is_causal_skip", None) - dtype = kwargs.get("dtype", torch.float32) - mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) - # we use torch.finfo(torch.float16).min instead torch.finfo(dtype).min to avoid an overflow but not - # sure this is the right way to handle this, we are basically pretending that -65,504 is -inf - mask = torch.where( - mask, - torch.tensor(0.0, device=mask.device, dtype=dtype), - torch.tensor(torch.finfo(torch.float16).min, device=mask.device, dtype=dtype), - ) - return mask - - # Adapted from https://github.com/huggingface/transformers/blob/3c307e380ad07ca16903a39e09a47d532cb782d9/src/transformers/models/phimoe/modular_phimoe.py#L57 def _longrope_forward(self, x, position_ids=None, layer_type=None): # _compute_longrope_parameters https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/modeling_rope_utils.py#L391 @@ -697,7 +762,7 @@ def _glm4_core_attention_forward(self, query_layer, key_layer, value_layer, atte class ChatGLMModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -1081,7 +1146,7 @@ def _qwen_attention_forward( class QwenModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -1242,7 +1307,7 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids): class BaichuanModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3175,7 +3240,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Gemma2ModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -3364,7 +3429,7 @@ def __exit__(self, exc_type, exc_value, traceback): class IBertModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3382,7 +3447,7 @@ def __init__( class InternVLChatImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3405,7 +3470,7 @@ def __exit__(self, exc_type, exc_value, traceback): class InternVL2ChatLangModelPatcher(OVDecoderModelPatcher): - def __init__(self, config: "OnnxConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any]): + def __init__(self, config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any]): model_type = model.config.model_type patcher_for_model_type = { "llama": OVDecoderModelPatcher, @@ -3525,7 +3590,7 @@ def maira_vision_embed_forward(self, pixel_values): class LlavaImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3541,7 +3606,7 @@ def __exit__(self, exc_type, exc_value, traceback): class MairaImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3558,7 +3623,7 @@ def __exit__(self, exc_type, exc_value, traceback): class LlavaNextVideoImageEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3597,7 +3662,7 @@ def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: return emb.unsqueeze(1) -class FluxTransfromerModelPatcher(ModelPatcher): +class FluxTransformerModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() if is_diffusers_version("<", "0.31.0"): @@ -3775,7 +3840,7 @@ def _minicpmv_siglip_transformer_forward( class MiniCPMVResamplerModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3792,7 +3857,7 @@ def __exit__(self, exc_type, exc_value, traceback): class MiniCPMVImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3823,7 +3888,7 @@ def __exit__(self, exc_type, exc_value, traceback): class LlavaQwen2ImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3841,7 +3906,7 @@ def __exit__(self, exc_type, exc_value, traceback): class InputEmbeddingPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -3866,7 +3931,7 @@ def phi3_vision_embeddings_forward(self, pixel_values: torch.FloatTensor): class Phi3VisionImageEmbeddingsPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -4318,7 +4383,7 @@ def deepseek_moe_infer(self, x, topk_ids, topk_weight): class Qwen2VLLanguageModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -4361,7 +4426,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen3VLLanguageModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: Union["PreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -4527,7 +4592,7 @@ def block_forward( class Qwen2VLVisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -4561,7 +4626,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen2_5_VLVisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -4623,7 +4688,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen3VLVisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: Union["PreTrainedModel"], model_kwargs: Dict[str, Any] = None, ): @@ -4736,7 +4801,7 @@ def __exit__(self, exc_type, exc_value, traceback): class OVSeq2SeqModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -4782,7 +4847,7 @@ def patched_forward(*args, **kwargs): outputs = self.super_patched_forward(*args, **kwargs) - # the optimum-onnx seq2seq model patcher only converts to tuple starting from 4.48 + # the seq2seq model patcher only converts to tuple starting from 4.48 if isinstance(outputs.get("past_key_values"), (DynamicCache, EncoderDecoderCache)): outputs["past_key_values"] = postprocess_past_key_values(outputs["past_key_values"]) @@ -4861,7 +4926,7 @@ def __exit__(self, exc_type, exc_value, traceback): class MiniCPMModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -4876,7 +4941,7 @@ def __init__( class CommonImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -4960,7 +5025,7 @@ def _gemma3_mm_update_causal_mask( class Gemma3LMModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -5200,8 +5265,6 @@ def gemma4_lm_forward( logits_to_keep: Union[int, torch.Tensor] = 0, **lm_kwargs, ): - from optimum.exporters.onnx.model_patcher import preprocess_past_key_values - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states @@ -5396,7 +5459,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Idefics3ImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6182,7 +6245,7 @@ def speecht5_decoder_layer_forward( return outputs -class OVSpeechT5ModelPatcher(ModelPatcher): +class SpeechT5ModelPatcher(ModelPatcher): def __enter__(self): if self.real_config._behavior != "vocoder": super().__enter__() @@ -6212,7 +6275,7 @@ def __exit__(self, exc_type, exc_value, traceback): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -6327,7 +6390,7 @@ def patched_vocoder_forward(spectrogram: torch.FloatTensor): class Phi4MMLanguageModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6368,7 +6431,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Phi4MMAudioForwardEmbeddingsPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6392,7 +6455,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Phi4MMAudioEncoderPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6433,7 +6496,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Phi4MMVisionEmbeddingsPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -6742,7 +6805,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Llama4ImageEmbeddingsModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any], ): @@ -7107,7 +7170,7 @@ def mamba_mixer_forward( class MambaPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -7614,7 +7677,7 @@ def segment_sum(input_tensor): class Zamba2ModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -7827,7 +7890,7 @@ def lfm2_short_conv_forward_patched( class Lfm2ModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -7986,6 +8049,26 @@ def __exit__(self, exc_type, exc_value, traceback): conv_layer.slow_forward = conv_layer._orig_forward +# Copied from https://github.com/huggingface/transformers/blob/v4.56.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L81 +def gpt_oss_forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor: + batch_size = hidden_states.shape[0] + hidden_states = hidden_states.reshape(-1, self.hidden_size) + num_experts = routing_weights.shape[1] + hidden_states = hidden_states.repeat(num_experts, 1) + hidden_states = hidden_states.view(num_experts, -1, self.hidden_size) + gate_up = torch.bmm(hidden_states, self.gate_up_proj) + self.gate_up_proj_bias[..., None, :] + gate, up = gate_up[..., ::2], gate_up[..., 1::2] + gate = gate.clamp(min=None, max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + glu = gate * torch.sigmoid(gate * self.alpha) + next_states = torch.bmm(((up + 1) * glu), self.down_proj) + next_states = next_states + self.down_proj_bias[..., None, :] + next_states = next_states.view(num_experts, batch_size, -1, self.hidden_size) + next_states = next_states * routing_weights.transpose(0, 1).view(num_experts, batch_size, -1)[..., None] + next_states = next_states.sum(dim=0) + return next_states + + class GptOssModelPatcher(OVDecoderModelPatcher): def __enter__(self): super().__enter__() @@ -8052,7 +8135,7 @@ def granite_moe_hybrid_update_causal_mask( class GraniteMoeHybridModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -8234,7 +8317,7 @@ def __enter__(self): if self.real_config._behavior == "encoder" and self._model.config.attention_type == "block_sparse": logger.warning( - "BigBirdPegasus model is using block sparse attention, which is not supported in ONNX export. " + "BigBirdPegasus model is using block sparse attention, which is not supported in OpenVINO export. " "The model will be exported with original full attention." ) self._model.set_attention_type("original_full") @@ -8334,7 +8417,7 @@ def __exit__(self, exc_type, exc_value, traceback): class VideoChatFlashQwenVisionEmbeddingModelPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -8884,7 +8967,7 @@ def forward( class Qwen3NextModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -9062,32 +9145,6 @@ def __exit__(self, exc_type, exc_value, traceback): del sparse_moe_block.down_projs, sparse_moe_block.gate_projs, sparse_moe_block.up_projs -class Gemma4PerLayerInputsGetterModelPatcher(ModelPatcher): - def __init__( - self, - config: "OnnxConfig", - model: Union["PreTrainedModel"], - model_kwargs: Dict[str, Any] = None, - ): - model.__orig_forward = model.forward - - def per_layer_inputs_forward(self, input_ids: torch.Tensor) -> torch.Tensor: - per_layer_inputs_mask = torch.logical_and(input_ids >= 0, input_ids < self.vocab_size_per_layer_input) - per_layer_inputs_tokens = torch.where(per_layer_inputs_mask, input_ids, torch.zeros_like(input_ids)) - per_layer_inputs = self.language_model.get_per_layer_inputs(per_layer_inputs_tokens, None) - return per_layer_inputs - - model.forward = types.MethodType(per_layer_inputs_forward, model) - super().__init__(config, model, model_kwargs) - - def __enter__(self): - super().__enter__() - - def __exit__(self, exc_type, exc_value, traceback): - super().__exit__(exc_type, exc_value, traceback) - self._model.forward = self._model.__orig_forward - - # OpenVINO has a bug due to which Clamp(-inf, inf) doesn't work correctly: CVS-185473. # When min == -inf and max == inf, Clamp is equivalent to an identity operation and # can be removed from the model, which serves as a workaround for the issue. @@ -9136,8 +9193,6 @@ def patched_encoder_forward(inputs_embeds, attention_mask=None, pixel_position_i **kwargs, ) - from transformers.modeling_outputs import BaseModelOutputWithPast - return BaseModelOutputWithPast(last_hidden_state=hidden_states) self._orig_encoder_forward = orig_encoder_forward @@ -9366,7 +9421,7 @@ def apply_mask_to_padding_states(hidden_states, attention_mask): class Qwen3_5ModelPatcher(OVDecoderModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -9551,7 +9606,7 @@ def __exit__(self, exc_type, exc_value, traceback): class Qwen3_5VisionEmbMergerPatcher(ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Dict[str, Any] = None, ): @@ -9625,7 +9680,7 @@ def patched_qwen3_5_moe_sparse_moe_block(self, hidden_states: torch.Tensor) -> t class Qwen3_5MoeModelPatcher(Qwen3_5ModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): @@ -9661,7 +9716,7 @@ class Qwen3ASRModelPatcher(OVSeq2SeqModelPatcher): def __init__( self, - config: "OnnxConfig", + config: "OpenVINOConfig", model: "PreTrainedModel", model_kwargs: Optional[Dict[str, Any]] = None, ): diff --git a/optimum/exporters/openvino/patching_utils.py b/optimum/exporters/openvino/patching_utils.py new file mode 100644 index 0000000000..d8ce0a2c05 --- /dev/null +++ b/optimum/exporters/openvino/patching_utils.py @@ -0,0 +1,446 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import functools +import inspect +import logging +import types +from typing import Any, Callable + +import torch +import transformers +from transformers import PreTrainedModel +from transformers.cache_utils import DynamicCache, EncoderDecoderCache +from transformers.modeling_outputs import BaseModelOutput + +from optimum.exporters.base import ExporterConfig +from optimum.intel.utils.import_utils import is_transformers_version + + +if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + from optimum.exporters.openvino._traceable_cache import TraceableCache + + +if is_transformers_version(">=", "4.48"): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache +if is_transformers_version(">=", "4.53"): + from transformers.masking_utils import ( + ALL_MASK_ATTENTION_FUNCTIONS, + _ignore_causal_mask_sdpa, + and_masks, + causal_mask_function, + eager_mask, + padding_mask_function, + prepare_padding_mask, + sdpa_mask, + ) + + +if is_transformers_version(">=", "4.53.1"): + from transformers.masking_utils import find_packed_sequence_indices + + +if is_transformers_version(">=", "4.54"): + from transformers.utils import TransformersKwargs + + from optimum.exporters.openvino._traceable_decorator import traceable_check_model_inputs +else: + TransformersKwargs = object + + +if is_transformers_version(">=", "4.56"): + import transformers.masking_utils + from transformers.cache_utils import DynamicLayer + + +logger = logging.getLogger(__name__) + + +def override_arguments(args, kwargs, forward_signature, model_kwargs: dict[str, Any]): + """Override the args and kwargs with the argument values from model_kwargs, following the signature forward_signature corresponding to args and kwargs.""" + args = list(args) + + for argument in model_kwargs: + if argument in forward_signature.parameters: + argument_index = list(forward_signature.parameters.keys()).index(argument) + if argument in kwargs or len(args) <= argument_index: + kwargs[argument] = model_kwargs[argument] + else: + args[argument_index] = model_kwargs[argument] + else: + kwargs[argument] = model_kwargs[argument] + + return args, kwargs + + +def preprocess_encoder_outputs(encoder_outputs): + if is_transformers_version(">=", "4.54") and isinstance(encoder_outputs, (list, tuple)): + encoder_outputs = BaseModelOutput(*encoder_outputs) + + return encoder_outputs + + +def preprocess_past_key_values(past_key_values): + if ( + is_transformers_version(">=", "4.48") + and isinstance(past_key_values, (list, tuple)) + and isinstance(past_key_values[0], (list, tuple)) + ): + if len(past_key_values[0]) == 2: + if hasattr(DynamicCache, "from_legacy_cache"): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + else: + past_key_values = DynamicCache(past_key_values) + elif len(past_key_values[0]) == 4: + if hasattr(EncoderDecoderCache, "from_legacy_cache"): + past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values) + else: + past_key_values = EncoderDecoderCache( + DynamicCache([layer[:2] for layer in past_key_values]), + DynamicCache([layer[2:] for layer in past_key_values]), + ) + else: + raise ValueError( + f"past_key_values should have either 2 or 4 elements, but it has {len(past_key_values[0])} elements." + ) + + return past_key_values + + +def postprocess_past_key_values(past_key_values): + if isinstance(past_key_values, (EncoderDecoderCache, DynamicCache)): + if hasattr(past_key_values, "to_legacy_cache"): + past_key_values = past_key_values.to_legacy_cache() + elif isinstance(past_key_values, DynamicCache): + past_key_values = [(lay.keys, lay.values) for lay in past_key_values.layers] + elif isinstance(past_key_values, EncoderDecoderCache): + past_key_values = [ + (self_lay.keys, self_lay.values, cross_lay.keys, cross_lay.values) + for self_lay, cross_lay in zip( + past_key_values.self_attention_cache.layers, + past_key_values.cross_attention_cache.layers, + ) + ] + return past_key_values + + +@dataclasses.dataclass +class PatchingSpec: + """Data class that holds patching specifications. + + Args: + o: Module / object where the op to patch is located + name: Name of the op to monkey patch + custom_op: Custom op that patches the original op + orig_op: Original op that is being patched + op_wrapper: Wrapper (optional) that wraps both the original and custom ops. + It is useful for ops that are class or static methods for instance. + """ + + o: Any + name: str + custom_op: Callable + orig_op: Callable | None = None + op_wrapper: Callable | None = None + + +# A patched version of https://github.com/huggingface/transformers/blob/v4.53.2/src/transformers/masking_utils.py#L602 +# That returns a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. +def find_packed_sequence_indices_patched(position_ids: torch.Tensor) -> torch.Tensor: + return torch.zeros_like(position_ids) + + +if is_transformers_version(">=", "4.53"): + _prepare_padding_mask_slice = "_slice" in inspect.signature(prepare_padding_mask).parameters +else: + _prepare_padding_mask_slice = False + + +# Custom vectorized implementation of sdpa_mask without using vmap +def _orig_sdpa_mask_without_vmap( + batch_size: int, + cache_position: torch.Tensor, + kv_length: int, + kv_offset: int = 0, + mask_function: Callable | None = None, + attention_mask: torch.Tensor | None = None, + local_size: int | None = None, + allow_is_causal_skip: bool = True, + **kwargs, +) -> torch.Tensor | None: + if mask_function is None: + mask_function = causal_mask_function + + q_length = cache_position.shape[0] + # Potentially pad the 2D mask, and slice it correctly + if _prepare_padding_mask_slice: + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset, _slice=False) + else: + padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset) + + # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument + if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, kv_offset, local_size): + return None + + # Potentially add the padding 2D mask + if padding_mask is not None: + mask_function = and_masks(mask_function, padding_mask_function(padding_mask)) + + # Create broadcatable indices + device = cache_position.device + q_indices = cache_position[None, None, :, None] + head_indices = torch.arange(1, dtype=torch.long, device=device)[None, :, None, None] + batch_indices = torch.arange(batch_size, dtype=torch.long, device=device)[:, None, None, None] + kv_indices = torch.arange(kv_length, dtype=torch.long, device=device)[None, None, None, :] + kv_offset + # Apply mask function element-wise through broadcasting + causal_mask = mask_function(batch_indices, head_indices, q_indices, kv_indices) + # Expand the mask to match batch size and query length if they weren't used in the mask function + causal_mask = causal_mask.expand(batch_size, -1, q_length, kv_length) + + return causal_mask + + +# Compatibility wrapper for sdpa_mask_without_vmap from optimum. +# The installed optimum version expects (batch_size, cache_position: Tensor, kv_length, ...), +# but transformers >= 5.5 passes (batch_size, q_length: int, kv_length: int, q_offset: int, ...). +def sdpa_mask_without_vmap(batch_size, q_length=None, kv_length=None, q_offset=0, kv_offset=0, **kwargs): + import inspect + + sig = inspect.signature(_orig_sdpa_mask_without_vmap) + if is_transformers_version(">=", "5.5") and "cache_position" in sig.parameters and q_length is not None: + # Old optimum signature: (batch_size, cache_position, kv_length, kv_offset, ...) + cache_position = torch.arange(q_length, dtype=torch.long) + q_offset + kwargs.pop("q_offset", None) + kwargs.pop("allow_is_bidirectional_skip", None) + kwargs.pop("allow_torch_fix", None) + kwargs.pop("use_vmap", None) + kwargs.pop("device", None) + return _orig_sdpa_mask_without_vmap(batch_size, cache_position, kv_length, kv_offset=kv_offset, **kwargs) + else: + return _orig_sdpa_mask_without_vmap( + batch_size, q_length=q_length, kv_length=kv_length, q_offset=q_offset, kv_offset=kv_offset, **kwargs + ) + + +# Adapted from https://github.com/huggingface/transformers/blob/v4.53.0/src/transformers/masking_utils.py#L433 +# Specifically for OpenVINO, we use torch.finfo(torch.float16).min instead of torch.finfo(dtype).min +def eager_mask_without_vmap(*args, **kwargs): + kwargs.pop("allow_is_causal_skip", None) + dtype = kwargs.get("dtype", torch.float32) + mask = sdpa_mask_without_vmap(*args, allow_is_causal_skip=False, **kwargs) + # we use torch.finfo(torch.float16).min instead torch.finfo(dtype).min to avoid an overflow but not + # sure this is the right way to handle this, we are basically pretending that -65,504 is -inf + mask = torch.where( + mask, + torch.tensor(0.0, device=mask.device, dtype=dtype), + torch.tensor(torch.finfo(torch.float16).min, device=mask.device, dtype=dtype), + ) + return mask + + +def patched_dynamic_layer_update( + self, key_states: torch.Tensor, value_states: torch.Tensor, cache_kwargs: dict[str, Any] | None = None +) -> tuple[torch.Tensor, torch.Tensor]: + if self.keys is None: + self.keys = key_states + self.values = value_states + self.device = key_states.device + self.dtype = key_states.dtype + self.is_initialized = True + else: + self.keys = torch.cat([self.keys, key_states], dim=-2) + self.values = torch.cat([self.values, value_states], dim=-2) + return self.keys, self.values + + +UNSUPPORTED_OPS_PATCHING_SPEC = [ + # TracerWarning: Using len to get tensor shape might cause the trace to be incorrect. Recommended usage would be tensor.shape[0]. Passing a tensor of different shape might lead to errors or silently give incorrect results. + PatchingSpec(torch.Tensor, "__len__", lambda x: x.shape[0], torch.Tensor.__len__), +] + + +class ModelPatcher: + def __init__( + self, + config: ExporterConfig, + model: PreTrainedModel, + model_kwargs: dict[str, Any] | None = None, + ): + self._model = model + + patching_specs = config.PATCHING_SPECS or [] + patching_specs.extend(UNSUPPORTED_OPS_PATCHING_SPEC) + + self._patching_specs = [] + for spec in patching_specs: + final_spec = spec + if spec.orig_op is None: + final_spec = dataclasses.replace(spec, orig_op=getattr(spec.o, spec.name)) + self._patching_specs.append(final_spec) + + self.orig_forward_name = "forward" if hasattr(self._model, "forward") else "call" + self.orig_forward = getattr(self._model, self.orig_forward_name) + + if is_transformers_version(">=", "4.54") and hasattr(self.orig_forward, "__wrapped__"): + # the original check_model_inputs has some failing cases that we fix in traceable_check_model_inputs + # we fix those issues in a PR in transformers https://github.com/huggingface/transformers/pull/40811 + # issues are: support for positional args (use_cache for instance) and fix for _CAN_RECORD_REGISTRY + # explicitly mapping to None for some models + self.orig_forward = types.MethodType( + traceable_check_model_inputs(self.orig_forward.__wrapped__), self._model + ) + + self.real_config = config + use_cache = getattr(self.real_config, "use_past", False) + self.model_kwargs = model_kwargs if model_kwargs is not None else {} + + for module in self._model.modules(): + if hasattr(module, "config") and hasattr(module.config, "use_cache"): + module.config.use_cache = use_cache + + @functools.wraps(self.orig_forward) + def patched_forward(*args, **kwargs): + signature = inspect.signature(self.orig_forward) + args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) + + # Transformers doesn't always respect the config.use_cache attribute + # there are cases where setting use_cache to true in every config and + # subconfig of a model still doesn't enable past_key_values in the outputs (gemma3) + # Explicitly setting the use_cache argument of the forward method seems to be the most reliable way + if "use_cache" in signature.parameters: + use_cache_index = list(signature.parameters.keys()).index("use_cache") + if use_cache_index < len(args): + args[use_cache_index] = use_cache + elif "use_cache" in kwargs: + kwargs["use_cache"] = use_cache + + if "past_key_values" in signature.parameters: + # Most models require past_key_values to be a cache instance instead of a tuple now + pkv_index = list(signature.parameters.keys()).index("past_key_values") + if pkv_index < len(args) and args[pkv_index] is not None: + args[pkv_index] = preprocess_past_key_values(args[pkv_index]) + elif kwargs.get("past_key_values") is not None: + kwargs["past_key_values"] = preprocess_past_key_values(kwargs["past_key_values"]) + + if "encoder_outputs" in signature.parameters: + # Some encoder-decoder models started to not accept encoder_outputs as tuple (e.g. moonshine) + encoder_outputs_index = list(signature.parameters.keys()).index("encoder_outputs") + if encoder_outputs_index < len(args) and args[encoder_outputs_index] is not None: + args[encoder_outputs_index] = preprocess_encoder_outputs(args[encoder_outputs_index]) + elif kwargs.get("encoder_outputs") is not None: + kwargs["encoder_outputs"] = preprocess_encoder_outputs(kwargs["encoder_outputs"]) + + outputs = self.orig_forward(*args, **kwargs) + + # This code block handles different cases of the filtered_outputs input to align it with the expected + # format of outputs. It is common for the output type of a model to vary, such as tensor, list, + # tuple, etc. For Transformers models, the output is encapsulated in a ModelOutput object that + # contains the output names of the model. In the case of Timm classification models, the output + # is of type tensor. By default, it is assumed that the output names mentioned in the ONNX config + # match the outputs in order. + filtered_outputs = {} + output_names = list(config.outputs.keys()) + if isinstance(outputs, dict): + for name, value in outputs.items(): + ov_output_name = config.torch_to_ov_output_map.get(name, name) + if ( + ov_output_name in output_names + or (use_cache and name.startswith("past_key_values")) + or any(key.startswith(ov_output_name) for key in output_names) + ): + filtered_outputs[name] = value + elif isinstance(outputs, (list, tuple)): + filtered_outputs = dict(zip(output_names, outputs)) + else: + if len(output_names) > 1: + num_outputs = len(output_names) + output_names_str = ", ".join(output_names) + raise ValueError( + f"{config.__class__.__name__} expects the model to return {num_outputs} outputs: {output_names_str}, " + f"but the it returned a single output of type {type(outputs)}. Please make sure either that the model " + "returns all the expected outputs, or that the ONNX config is correctly defined with the expected outputs." + ) + output_name = output_names[0] + filtered_outputs[output_name] = outputs + + if filtered_outputs.get("past_key_values") is not None: + filtered_outputs["past_key_values"] = postprocess_past_key_values(filtered_outputs["past_key_values"]) + + return filtered_outputs + + self.patched_forward = patched_forward + + def patch_ops(self): + for spec in self._patching_specs: + custom_op = spec.custom_op if spec.op_wrapper is None else spec.op_wrapper(spec.custom_op) + setattr(spec.o, spec.name, custom_op) + + def restore_ops(self): + for spec in self._patching_specs: + orig_op = spec.orig_op if spec.op_wrapper is None else spec.op_wrapper(spec.orig_op) + setattr(spec.o, spec.name, orig_op) + + def __enter__(self): + self.patch_ops() + setattr(self._model, self.orig_forward_name, self.patched_forward) + + # This is a workaround for the Cache class in transformers, we replace it + # with traceable cache is because the original one used in transformers + # inherited from nn.Module (for a couple versions), which can't be traced as input. + if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + self.original_cache_class = transformers.cache_utils.Cache + transformers.cache_utils.Cache = TraceableCache + + # This is a workaround for mask generation in transformers >= 4.53. + # The masking process uses vmap which is not traceable by TorchScript. + if is_transformers_version(">=", "4.53"): + ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask_without_vmap) + ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask_without_vmap) + + # This is a workaround for the find_packed_sequence_indices function in transformers which + # should only return a tensor of zeros with the same shape as position_ids indicating no packed sequence indices. + # The function uses torch.diff which is not traceable by TorchScript. + if is_transformers_version(">=", "4.53.1"): + self.original_find_packed_sequence_indices = find_packed_sequence_indices + transformers.masking_utils.find_packed_sequence_indices = find_packed_sequence_indices_patched + + # Starting from transformers 4.56.0, DynamicCache uses DynamicLayer which has an update method + # that uses torch.cat to concatenate an empty tensor with the key/value states during the first call. + # This causes issues during TorchScript tracing. + if is_transformers_version(">=", "4.56"): + self.original_dynamic_layer_update = DynamicLayer.update + DynamicLayer.update = patched_dynamic_layer_update + + def __exit__(self, exc_type, exc_value, traceback): + self.restore_ops() + setattr(self._model, self.orig_forward_name, self.orig_forward) + + if is_transformers_version(">=", "4.44") and is_transformers_version("<", "4.50"): + transformers.cache_utils.Cache = self.original_cache_class + + if is_transformers_version(">=", "4.53"): + ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa", sdpa_mask) + ALL_MASK_ATTENTION_FUNCTIONS.register("eager", eager_mask) + + if is_transformers_version(">=", "4.53.1"): + transformers.masking_utils.find_packed_sequence_indices = self.original_find_packed_sequence_indices + + if is_transformers_version(">=", "4.56"): + DynamicLayer.update = self.original_dynamic_layer_update + + def __call__(self, *args, **kwargs): + if getattr(self._model, self.orig_forward_name) is self.orig_forward: + logger.warning("Running the non-patched model") + return self._model(*args, **kwargs) diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index a1756e27a6..77ec6c7cd7 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -24,7 +24,7 @@ from openvino import Dimension, PartialShape, Symbol from openvino.utils.types import get_element_type -from optimum.exporters.onnx.base import OnnxConfig +from optimum.exporters.openvino.base import OpenVINOConfig from optimum.exporters.tasks import TasksManager from optimum.intel.utils.import_utils import is_safetensors_available from optimum.utils import is_diffusers_available @@ -87,7 +87,7 @@ def flattenize_inputs(inputs: List[Any]): def _get_input_info( - model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any] + model: Union["PreTrainedModel", "ModelMixin"], config: OpenVINOConfig, dummy_inputs: Dict[str, Any] ) -> List[InputInfo]: sig = inspect.signature(model.forward) if hasattr(model, "forward") else inspect.signature(model.call) inputs = config.ordered_inputs(model) @@ -129,7 +129,7 @@ def _get_input_info( def _get_dynamic_shapes_info( - model: Union["PreTrainedModel", "ModelMixin"], config: OnnxConfig, dummy_inputs: Dict[str, Any] + model: Union["PreTrainedModel", "ModelMixin"], config: OpenVINOConfig, dummy_inputs: Dict[str, Any] ) -> List[InputInfo]: import torch @@ -248,7 +248,7 @@ def _get_open_clip_submodels_fn_and_export_configs( library_name: str = "open_clip", task: Optional[str] = None, preprocessors: List = None, - custom_export_configs: Dict[str, "OnnxConfig"] = None, + custom_export_configs: Dict[str, "OpenVINOConfig"] = None, fn_get_submodels: Callable = None, ): custom_export = {} @@ -295,7 +295,7 @@ def _get_kokoro_submodels_fn_and_export_configs( library_name: str = "kokoro", task: Optional[str] = None, preprocessors: List = None, - custom_export_configs: Dict[str, "OnnxConfig"] = None, + custom_export_configs: Dict[str, "OpenVINOConfig"] = None, fn_get_submodels: Callable = None, ): export_config_constructor = TasksManager.get_exporter_config_constructor( @@ -350,8 +350,7 @@ def _get_kokoro_submodels(model): "qwen3_5_moe_text", ] -# All transformers, diffusers, timm and sentence transformers models that are supported via optimum-onnx OnnxConfigs but that have currently no test -# TODO: add tests for all models that are compatible and remove support for all others +# All transformers, diffusers, timm and sentence transformers models that were supported via optimum-onnx OnnxConfigs for which support is now removed ONNX_SUPPORTED_ARCHITECTURES = { "big_bird", "chinese_clip", @@ -363,7 +362,6 @@ def _get_kokoro_submodels(model): "default-timm-config", "detr", "dinov2", - "donut-swin", "dpt", "efficientnet", "encoder-decoder", @@ -401,7 +399,6 @@ def _get_kokoro_submodels(model): "swin2sr", "swinv2", "table-transformer", - "trocr", "visual_bert", "vit_mae", "vit_msn", diff --git a/optimum/intel/openvino/modeling.py b/optimum/intel/openvino/modeling.py index 9fc568e042..91cf0149ec 100644 --- a/optimum/intel/openvino/modeling.py +++ b/optimum/intel/openvino/modeling.py @@ -571,16 +571,16 @@ def from_pretrained( "To load a timm model, please make sure to upgrade your `timm` version to at least 0.9.0, you can upgrade it by running `pip install --upgrade timm`" ) - from .modeling_timm import TimmConfig, TimmForImageClassification, TimmOnnxConfig + from .modeling_timm import TimmConfig, TimmForImageClassification, TimmOpenVINOConfig config = TimmConfig.from_pretrained(model_id, **kwargs) # If locally saved timm model, directly load if local_timm_model: return super()._from_pretrained(model_id=model_id, config=config) model = TimmForImageClassification.from_pretrained(model_id, **kwargs) - onnx_config = TimmOnnxConfig(model.config) + openvino_config = TimmOpenVINOConfig(model.config) - return cls._to_load(model=model, config=config, onnx_config=onnx_config, stateful=False, **kwargs) + return cls._to_load(model=model, config=config, exporter_config=openvino_config, stateful=False, **kwargs) else: return super().from_pretrained( model_id=model_id, @@ -716,7 +716,7 @@ def forward( @add_start_docstrings( """ - Onnx Model with a language modeling head on top for Connectionist Temporal Classification (CTC). + OpenVINO Model with a language modeling head on top for Connectionist Temporal Classification (CTC). """, MODEL_START_DOCSTRING, ) @@ -796,7 +796,7 @@ def forward( @add_start_docstrings( """ - Onnx Model with an XVector feature extraction head on top for tasks like Speaker Verification. + OpenVINO Model with an XVector feature extraction head on top for tasks like Speaker Verification. """, MODEL_START_DOCSTRING, ) diff --git a/optimum/intel/openvino/modeling_base.py b/optimum/intel/openvino/modeling_base.py index fd4231418e..40c9314d76 100644 --- a/optimum/intel/openvino/modeling_base.py +++ b/optimum/intel/openvino/modeling_base.py @@ -905,7 +905,7 @@ def _to_load( cls, model, config: PretrainedConfig, - onnx_config: ExportConfig, + exporter_config: ExportConfig, token: Optional[Union[bool, str]] = None, revision: Optional[str] = None, force_download: bool = False, @@ -924,11 +924,10 @@ def _to_load( ) compile_only = False - # Export the model to the ONNX format + # Export the model to the OpenVINO format export( model=model, - config=onnx_config, - opset=onnx_config.DEFAULT_ONNX_OPSET, + config=exporter_config, output=save_dir_path / cls._all_ov_model_paths["model"], stateful=stateful, ) diff --git a/optimum/intel/openvino/modeling_timm.py b/optimum/intel/openvino/modeling_timm.py index f2566c8c4a..d80e9b5c84 100644 --- a/optimum/intel/openvino/modeling_timm.py +++ b/optimum/intel/openvino/modeling_timm.py @@ -40,7 +40,7 @@ from transformers.modeling_outputs import ImageClassifierOutput from transformers.utils import TensorType -from optimum.exporters.onnx.config import VisionOnnxConfig +from optimum.exporters.openvino.config import VisionOpenVINOConfig from optimum.utils import NormalizedVisionConfig from .utils import _is_timm_ov_dir @@ -79,7 +79,7 @@ def from_pretrained( return cls.from_dict(config_dict, **kwargs) -class TimmOnnxConfig(VisionOnnxConfig): +class TimmOpenVINOConfig(VisionOpenVINOConfig): DEFAULT_TIMM_ONNX_OPSET = 13 outputs = OrderedDict([("logits", {0: "batch_size"})]) NORMALIZED_CONFIG_CLASS = NormalizedVisionConfig diff --git a/optimum/intel/openvino/quantization.py b/optimum/intel/openvino/quantization.py index 30005894f5..eb421646dc 100644 --- a/optimum/intel/openvino/quantization.py +++ b/optimum/intel/openvino/quantization.py @@ -1307,7 +1307,6 @@ def quantize( ] = None, save_directory: Optional[Union[str, Path]] = None, ov_config: OVConfig = None, - file_name: Optional[str] = None, batch_size: int = 1, data_collator: Optional[DataCollator] = None, remove_unused_columns: bool = False, @@ -1326,8 +1325,6 @@ def quantize( ov_config (`OVConfig`, *optional*): The configuration containing the parameters related to quantization. If not provided, 8-bit symmetric weight-only quantization will be applied. - file_name (`str`, *optional*): - The model file name to use when saving the model. Overwrites the default file name `"model.onnx"`. batch_size (`int`, defaults to 1): The number of calibration samples to load per batch. data_collator (`DataCollator`, *optional*): diff --git a/setup.py b/setup.py index 66947d1943..b6f2110b0b 100644 --- a/setup.py +++ b/setup.py @@ -28,7 +28,7 @@ INSTALL_REQUIRE = [ "torch>=2.1", - "optimum-onnx@git+https://github.com/huggingface/optimum-onnx.git@transformers-v5", + "optimum@git+https://github.com/huggingface/optimum.git", "transformers>=4.45,<5.1", "setuptools", "huggingface-hub>=0.23.2,<2.0", @@ -69,6 +69,7 @@ "decord", "imageio", "kokoro", + "protobuf", ] QUALITY_REQUIRE = ["black~=23.1", "ruff==0.4.4"] diff --git a/tests/openvino/test_export.py b/tests/openvino/test_export.py index fbca6ace56..45b606c7e1 100644 --- a/tests/openvino/test_export.py +++ b/tests/openvino/test_export.py @@ -24,11 +24,11 @@ MODEL_NAMES, OPENVINO_DEVICE, REMOTE_CODE_MODELS, + SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED, ) -from optimum.exporters.onnx.constants import SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED -from optimum.exporters.onnx.model_configs import BertOnnxConfig from optimum.exporters.openvino import export_from_model, main_export +from optimum.exporters.openvino.model_configs import BertOpenVINOConfig from optimum.exporters.tasks import TasksManager from optimum.intel import ( OVFluxPipeline, @@ -369,7 +369,7 @@ def test_compare_openvino_onnx_supported_architectures(self): class CustomExportModelTest(unittest.TestCase): def test_custom_export_config_model(self): - class BertOnnxConfigWithPooler(BertOnnxConfig): + class BertOpenVINOConfigWithPooler(BertOpenVINOConfig): @property def outputs(self): if self.task == "feature-extraction-with-pooler": @@ -386,7 +386,7 @@ def outputs(self): model_id = "sentence-transformers/all-MiniLM-L6-v2" config = AutoConfig.from_pretrained(model_id) - custom_export_configs = {"model": BertOnnxConfigWithPooler(config, task=custom_task)} + custom_export_configs = {"model": BertOpenVINOConfigWithPooler(config, task=custom_task)} with TemporaryDirectory() as tmpdirname: main_export( diff --git a/tests/openvino/test_quantization.py b/tests/openvino/test_quantization.py index 979e1c246d..d286718c26 100644 --- a/tests/openvino/test_quantization.py +++ b/tests/openvino/test_quantization.py @@ -2076,7 +2076,7 @@ def preprocess_function(examples, tokenizer): tokenizer = AutoTokenizer.from_pretrained(model_name) quantizer = OVQuantizer.from_pretrained(transformers_model, device=OPENVINO_DEVICE) calibration_dataset = quantizer.get_calibration_dataset( - "squadshifts", + "ludwigschmidt/squadshifts", dataset_config_name="new_wiki", preprocess_function=partial(preprocess_function, tokenizer=tokenizer), num_samples=10, diff --git a/tests/openvino/utils_tests.py b/tests/openvino/utils_tests.py index 8a842c3159..5edbd3a8d5 100644 --- a/tests/openvino/utils_tests.py +++ b/tests/openvino/utils_tests.py @@ -583,6 +583,13 @@ def _create_tiny_kokoro_model(): REMOTE_CODE_MODELS += ("afmoe",) +SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED = [ + "bart", + "musicgen", + "whisper", +] + + def get_num_quantized_nodes(model): num_fake_nodes = 0 types_map = {