Accelerate your training and inference workflows with AWS Trainium and AWS Inferentia
Accelerate your training and inference workflows with Google TPUs
-> [!TIP] -> Some packages provide hardware-agnostic features (e.g. INC interface in Optimum Intel). - - ## Open-source integrations 🤗 Optimum also supports a variety of open-source frameworks to make model optimization very easy. diff --git a/examples/onnxruntime/training/docker/Dockerfile-ort-nightly-cu118 b/examples/onnxruntime/training/docker/Dockerfile-ort-nightly-cu118 index 3e6841453b..386477503c 100644 --- a/examples/onnxruntime/training/docker/Dockerfile-ort-nightly-cu118 +++ b/examples/onnxruntime/training/docker/Dockerfile-ort-nightly-cu118 @@ -74,4 +74,4 @@ RUN $PYTHON_EXE -m torch_ort.configure WORKDIR . -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] diff --git a/examples/onnxruntime/training/docker/Dockerfile-ort1.17.1-cu118 b/examples/onnxruntime/training/docker/Dockerfile-ort1.17.1-cu118 index ee58b71114..8cb79a73cf 100644 --- a/examples/onnxruntime/training/docker/Dockerfile-ort1.17.1-cu118 +++ b/examples/onnxruntime/training/docker/Dockerfile-ort1.17.1-cu118 @@ -76,4 +76,4 @@ RUN pip uninstall nvidia-nccl-cu12 -y WORKDIR . -CMD ["/bin/bash"] \ No newline at end of file +CMD ["/bin/bash"] diff --git a/examples/onnxruntime/training/image-classification/run_image_classification.py b/examples/onnxruntime/training/image-classification/run_image_classification.py index c5d5aabe27..c2bcb86aa0 100644 --- a/examples/onnxruntime/training/image-classification/run_image_classification.py +++ b/examples/onnxruntime/training/image-classification/run_image_classification.py @@ -333,6 +333,7 @@ def compute_metrics(p): token=model_args.token, trust_remote_code=model_args.trust_remote_code, ignore_mismatched_sizes=model_args.ignore_mismatched_sizes, + attn_implementation="eager", ) image_processor = AutoImageProcessor.from_pretrained( model_args.image_processor_name or model_args.model_name_or_path, diff --git a/examples/onnxruntime/training/language-modeling/run_clm.py b/examples/onnxruntime/training/language-modeling/run_clm.py index 10c0622ec9..9481e182a1 100644 --- a/examples/onnxruntime/training/language-modeling/run_clm.py +++ b/examples/onnxruntime/training/language-modeling/run_clm.py @@ -442,9 +442,12 @@ def main(): trust_remote_code=model_args.trust_remote_code, torch_dtype=torch_dtype, low_cpu_mem_usage=model_args.low_cpu_mem_usage, + attn_implementation="eager", ) else: - model = AutoModelForCausalLM.from_config(config, trust_remote_code=model_args.trust_remote_code) + model = AutoModelForCausalLM.from_config( + config, trust_remote_code=model_args.trust_remote_code, attn_implementation="eager" + ) n_params = sum({p.data_ptr(): p.numel() for p in model.parameters()}.values()) logger.info(f"Training new model from scratch - Total size={n_params/2**20:.2f}M params") diff --git a/examples/onnxruntime/training/language-modeling/run_mlm.py b/examples/onnxruntime/training/language-modeling/run_mlm.py index d032210fa5..e25c70a297 100755 --- a/examples/onnxruntime/training/language-modeling/run_mlm.py +++ b/examples/onnxruntime/training/language-modeling/run_mlm.py @@ -430,10 +430,13 @@ def main(): token=model_args.token, trust_remote_code=model_args.trust_remote_code, low_cpu_mem_usage=model_args.low_cpu_mem_usage, + attn_implementation="eager", ) else: logger.info("Training new model from scratch") - model = AutoModelForMaskedLM.from_config(config, trust_remote_code=model_args.trust_remote_code) + model = AutoModelForMaskedLM.from_config( + config, trust_remote_code=model_args.trust_remote_code, attn_implementation="eager" + ) # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch # on a small vocab and want a smaller embedding size, remove this test. diff --git a/examples/onnxruntime/training/question-answering/run_qa.py b/examples/onnxruntime/training/question-answering/run_qa.py index 08b581a1a8..c63f7f6a19 100644 --- a/examples/onnxruntime/training/question-answering/run_qa.py +++ b/examples/onnxruntime/training/question-answering/run_qa.py @@ -364,6 +364,7 @@ def main(): revision=model_args.model_revision, token=model_args.token, trust_remote_code=model_args.trust_remote_code, + attn_implementation="eager", ) # Tokenizer check: this script requires a fast tokenizer. diff --git a/examples/onnxruntime/training/summarization/run_summarization.py b/examples/onnxruntime/training/summarization/run_summarization.py index 83ec61f225..c6a80e626d 100644 --- a/examples/onnxruntime/training/summarization/run_summarization.py +++ b/examples/onnxruntime/training/summarization/run_summarization.py @@ -458,6 +458,7 @@ def main(): revision=model_args.model_revision, token=model_args.token, trust_remote_code=model_args.trust_remote_code, + attn_implementation="eager", ) if model.config.decoder_start_token_id is None and isinstance(tokenizer, (MBartTokenizer, MBartTokenizerFast)): diff --git a/examples/onnxruntime/training/text-classification/run_classification.py b/examples/onnxruntime/training/text-classification/run_classification.py index 6600e26c36..1edcc3a999 100755 --- a/examples/onnxruntime/training/text-classification/run_classification.py +++ b/examples/onnxruntime/training/text-classification/run_classification.py @@ -527,6 +527,7 @@ def main(): token=model_args.token, trust_remote_code=model_args.trust_remote_code, ignore_mismatched_sizes=model_args.ignore_mismatched_sizes, + attn_implementation="eager", ) model.config.pad_token_id = model.config.eos_token_id diff --git a/examples/onnxruntime/training/text-classification/run_glue.py b/examples/onnxruntime/training/text-classification/run_glue.py index f3f04657af..27e14199b4 100644 --- a/examples/onnxruntime/training/text-classification/run_glue.py +++ b/examples/onnxruntime/training/text-classification/run_glue.py @@ -404,6 +404,7 @@ def main(): token=model_args.token, trust_remote_code=model_args.trust_remote_code, ignore_mismatched_sizes=model_args.ignore_mismatched_sizes, + attn_implementation="eager", ) # Preprocessing the raw_datasets diff --git a/examples/onnxruntime/training/token-classification/run_ner.py b/examples/onnxruntime/training/token-classification/run_ner.py index 55ddfa2cf0..102249fc51 100644 --- a/examples/onnxruntime/training/token-classification/run_ner.py +++ b/examples/onnxruntime/training/token-classification/run_ner.py @@ -405,6 +405,7 @@ def get_label_list(labels): token=model_args.token, trust_remote_code=model_args.trust_remote_code, ignore_mismatched_sizes=model_args.ignore_mismatched_sizes, + attn_implementation="eager", ) if tokenizer.pad_token is None: diff --git a/examples/onnxruntime/training/translation/run_translation.py b/examples/onnxruntime/training/translation/run_translation.py index 0b6a36d12f..f54246be33 100644 --- a/examples/onnxruntime/training/translation/run_translation.py +++ b/examples/onnxruntime/training/translation/run_translation.py @@ -408,6 +408,7 @@ def main(): revision=model_args.model_revision, token=model_args.token, trust_remote_code=model_args.trust_remote_code, + attn_implementation="eager", ) # Set decoder_start_token_id diff --git a/optimum/exporters/onnx/__main__.py b/optimum/exporters/onnx/__main__.py index 6a2cc6834a..20bea423cb 100644 --- a/optimum/exporters/onnx/__main__.py +++ b/optimum/exporters/onnx/__main__.py @@ -29,6 +29,7 @@ from ...utils import DEFAULT_DUMMY_SHAPES, logging from ...utils.save_utils import maybe_load_preprocessors from ..tasks import TasksManager +from ..utils import DisableCompileContextManager from .constants import SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED from .convert import onnx_export_from_model @@ -255,7 +256,7 @@ def main_export( if task == "auto": try: - task = TasksManager.infer_task_from_model(model_name_or_path) + task = TasksManager.infer_task_from_model(model_name_or_path, library_name=library_name) except KeyError as e: raise KeyError( f"The task could not be automatically inferred. Please provide the argument --task with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}" @@ -300,22 +301,23 @@ def main_export( if model_type in SDPA_ARCHS_ONNX_EXPORT_NOT_SUPPORTED and _transformers_version >= version.parse("4.35.99"): loading_kwargs["attn_implementation"] = "eager" - model = TasksManager.get_model_from_task( - task, - model_name_or_path, - subfolder=subfolder, - revision=revision, - cache_dir=cache_dir, - token=token, - local_files_only=local_files_only, - force_download=force_download, - trust_remote_code=trust_remote_code, - framework=framework, - torch_dtype=torch_dtype, - device=device, - library_name=library_name, - **loading_kwargs, - ) + with DisableCompileContextManager(): + model = TasksManager.get_model_from_task( + task, + model_name_or_path, + subfolder=subfolder, + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=local_files_only, + force_download=force_download, + trust_remote_code=trust_remote_code, + framework=framework, + torch_dtype=torch_dtype, + device=device, + library_name=library_name, + **loading_kwargs, + ) needs_pad_token_id = task == "text-classification" and getattr(model.config, "pad_token_id", None) is None diff --git a/optimum/exporters/onnx/_traceable_cache.py b/optimum/exporters/onnx/_traceable_cache.py new file mode 100644 index 0000000000..052cb04b12 --- /dev/null +++ b/optimum/exporters/onnx/_traceable_cache.py @@ -0,0 +1,95 @@ +import logging +from typing import Any, Dict, Optional, Tuple + +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: Optional[Dict[str, Any]] = 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: Optional[int] = 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 specifc 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) -> Optional[int]: + 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) -> Optional[int]: + """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: Optional[int] = 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/onnx/base.py b/optimum/exporters/onnx/base.py index 8cd94194ff..43468a15c0 100644 --- a/optimum/exporters/onnx/base.py +++ b/optimum/exporters/onnx/base.py @@ -27,16 +27,12 @@ from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union import numpy as np -import onnx from transformers.utils import is_accelerate_available, is_torch_available -from ...onnx import remove_duplicate_weights_from_tied_info - if is_torch_available(): import torch.nn as nn -from ...onnx import merge_decoders from ...utils import ( DEFAULT_DUMMY_SHAPES, DummyInputGenerator, @@ -48,12 +44,19 @@ from ...utils import TORCH_MINIMUM_VERSION as GLOBAL_MIN_TORCH_VERSION from ...utils import TRANSFORMERS_MINIMUM_VERSION as GLOBAL_MIN_TRANSFORMERS_VERSION from ...utils.doc import add_dynamic_docstring -from ...utils.import_utils import check_if_transformers_greater, is_onnx_available, is_onnxruntime_available +from ...utils.import_utils import ( + is_onnx_available, + is_onnxruntime_available, + is_torch_version, + is_transformers_version, +) from ..base import ExportConfig from .constants import ONNX_DECODER_MERGED_NAME, ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME from .model_patcher import ModelPatcher, Seq2SeqModelPatcher +# TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization + if is_accelerate_available(): from accelerate.utils import find_tied_parameters @@ -151,14 +154,14 @@ class OnnxConfig(ExportConfig, ABC): "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", 1: "num_labels", 2: "height", 3: "width"}}), + "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"}} ), "mask-generation": OrderedDict({"logits": {0: "batch_size"}}), "masked-im": OrderedDict( - {"reconstruction" if check_if_transformers_greater("4.29.0") else "logits": {0: "batch_size"}} + {"reconstruction" if is_transformers_version(">=", "4.29.0") else "logits": {0: "batch_size"}} ), "multiple-choice": OrderedDict({"logits": {0: "batch_size", 1: "num_choices"}}), "object-detection": OrderedDict( @@ -177,6 +180,7 @@ class OnnxConfig(ExportConfig, ABC): "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( @@ -319,6 +323,7 @@ def fix_dynamic_axes( input_shapes = {} dummy_inputs = self.generate_dummy_inputs(framework="np", **input_shapes) dummy_inputs = self.generate_dummy_inputs_for_validation(dummy_inputs, onnx_input_names=onnx_input_names) + dummy_inputs = self.rename_ambiguous_inputs(dummy_inputs) onnx_inputs = {} for name, value in dummy_inputs.items(): @@ -376,7 +381,7 @@ def is_transformers_support_available(self) -> bool: `bool`: Whether the install version of Transformers is compatible with the model. """ - return check_if_transformers_greater(self.MIN_TRANSFORMERS_VERSION) + return is_transformers_version(">=", self.MIN_TRANSFORMERS_VERSION.base_version) @property def is_torch_support_available(self) -> bool: @@ -387,9 +392,8 @@ def is_torch_support_available(self) -> bool: `bool`: Whether the installed version of PyTorch is compatible with the model. """ if is_torch_available(): - from ...utils import torch_version + return is_torch_version(">=", self.MIN_TORCH_VERSION.base_version) - return torch_version >= self.MIN_TORCH_VERSION return False @property @@ -541,6 +545,10 @@ def post_process_exported_models( first_key = next(iter(models_and_onnx_configs)) if is_torch_available() and isinstance(models_and_onnx_configs[first_key][0], nn.Module): if is_accelerate_available(): + import onnx + + from ...onnx import remove_duplicate_weights_from_tied_info + logger.info("Deduplicating shared (tied) weights...") for subpath, key in zip(onnx_files_subpaths, models_and_onnx_configs): torch_model = models_and_onnx_configs[key][0] @@ -933,6 +941,8 @@ def post_process_exported_models( decoder_with_past_path = Path(path, onnx_files_subpaths[2]) decoder_merged_path = Path(path, ONNX_DECODER_MERGED_NAME + ".onnx") try: + from ...onnx import merge_decoders + # The decoder with past does not output the cross attention past key values as they are constant, # hence the need for strict=False merge_decoders( diff --git a/optimum/exporters/onnx/config.py b/optimum/exporters/onnx/config.py index 9e808e392b..69366d6be1 100644 --- a/optimum/exporters/onnx/config.py +++ b/optimum/exporters/onnx/config.py @@ -20,7 +20,6 @@ from transformers.utils import is_tf_available -from ...onnx import merge_decoders from ...utils import ( DummyAudioInputGenerator, DummyBboxInputGenerator, @@ -38,6 +37,9 @@ from .model_patcher import DecoderModelPatcher +# TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization + + if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel @@ -129,6 +131,8 @@ def post_process_exported_models( # Attempt to merge only if the decoder-only was exported separately without/with past if self.use_past is True and len(models_and_onnx_configs) == 2: + from ...onnx import merge_decoders + decoder_path = Path(path, onnx_files_subpaths[0]) decoder_with_past_path = Path(path, onnx_files_subpaths[1]) decoder_merged_path = Path(path, ONNX_DECODER_MERGED_NAME + ".onnx") diff --git a/optimum/exporters/onnx/convert.py b/optimum/exporters/onnx/convert.py index 2661d83597..0cca6d129f 100644 --- a/optimum/exporters/onnx/convert.py +++ b/optimum/exporters/onnx/convert.py @@ -22,7 +22,7 @@ from inspect import signature from itertools import chain from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import numpy as np import onnx @@ -35,9 +35,9 @@ DEFAULT_DUMMY_SHAPES, ONNX_WEIGHTS_NAME, TORCH_MINIMUM_VERSION, - check_if_transformers_greater, is_diffusers_available, is_torch_onnx_support_available, + is_transformers_version, logging, require_numpy_strictly_lower, ) @@ -45,6 +45,7 @@ from ...utils.save_utils import maybe_save_preprocessors from ..error_utils import AtolError, MinimumVersionError, OutputMatchError, ShapeError from ..tasks import TasksManager +from ..utils import check_dummy_inputs_are_allowed from .base import OnnxConfig from .constants import UNPICKABLE_ARCHS from .model_configs import SpeechT5OnnxConfig @@ -56,6 +57,8 @@ ) +# TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization + if is_torch_available(): import torch import torch.nn as nn @@ -75,30 +78,6 @@ class DynamicAxisNameError(ValueError): pass -def check_dummy_inputs_are_allowed( - model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], dummy_input_names: Iterable[str] -): - """ - Checks that the dummy inputs from the ONNX config is a subset of the allowed inputs for `model`. - Args: - model (`Union[transformers.PreTrainedModel, transformers.TFPreTrainedModel`]): - The model instance. - model_inputs (`Iterable[str]`): - The model input names. - """ - - forward = model.forward if is_torch_available() and isinstance(model, nn.Module) else model.call - forward_parameters = signature(forward).parameters - forward_inputs_set = set(forward_parameters.keys()) - dummy_input_names = set(dummy_input_names) - - # We are fine if config_inputs has more keys than model_inputs - if not dummy_input_names.issubset(forward_inputs_set): - raise ValueError( - f"Config dummy inputs are not a subset of the model inputs: {dummy_input_names} vs {forward_inputs_set}" - ) - - def validate_models_outputs( models_and_onnx_configs: Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"] @@ -533,7 +512,7 @@ def export_pytorch( model_kwargs = model_kwargs or {} # num_logits_to_keep was added in transformers 4.45 and isn't added as inputs when exporting the model - if check_if_transformers_greater("4.44.99") and "num_logits_to_keep" in signature(model.forward).parameters.keys(): + if is_transformers_version(">=", "4.44.99") and "num_logits_to_keep" in signature(model.forward).parameters.keys(): model_kwargs["num_logits_to_keep"] = 0 with torch.no_grad(): @@ -872,17 +851,16 @@ def export( ) if is_torch_available() and isinstance(model, nn.Module): - from ...utils import torch_version + from ...utils.import_utils import _torch_version if not is_torch_onnx_support_available(): raise MinimumVersionError( - f"Unsupported PyTorch version, minimum required is {TORCH_MINIMUM_VERSION}, got: {torch_version}" + f"Unsupported PyTorch version, minimum required is {TORCH_MINIMUM_VERSION}, got: {_torch_version}" ) if not config.is_torch_support_available: raise MinimumVersionError( - f"Unsupported PyTorch version for this model. Minimum required is {config.MIN_TORCH_VERSION}," - f" got: {torch.__version__}" + f"Unsupported PyTorch version for this model. Minimum required is {config.MIN_TORCH_VERSION}, got: {_torch_version}" ) export_output = export_pytorch( @@ -1126,7 +1104,7 @@ def onnx_export_from_model( if isinstance(atol, dict): atol = atol[task.replace("-with-past", "")] - if check_if_transformers_greater("4.44.99"): + if is_transformers_version(">=", "4.44.99"): misplaced_generation_parameters = model.config._get_non_default_generation_parameters() if ( isinstance(model, GenerationMixin) @@ -1183,6 +1161,10 @@ def onnx_export_from_model( if tokenizer_2 is not None: tokenizer_2.save_pretrained(output.joinpath("tokenizer_2")) + tokenizer_3 = getattr(model, "tokenizer_3", None) + if tokenizer_3 is not None: + tokenizer_3.save_pretrained(output.joinpath("tokenizer_3")) + model.save_config(output) if float_dtype == "bf16": diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py index 9e57128c27..f420ab39c6 100644 --- a/optimum/exporters/onnx/model_configs.py +++ b/optimum/exporters/onnx/model_configs.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """Model specific ONNX configurations.""" + import random from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union @@ -20,17 +21,20 @@ from packaging import version from transformers.utils import is_tf_available -from ...onnx import merge_decoders from ...utils import ( DEFAULT_DUMMY_SHAPES, BloomDummyPastKeyValuesGenerator, DummyAudioInputGenerator, DummyCodegenDecoderTextInputGenerator, + DummyDecisionTransformerInputGenerator, DummyDecoderTextInputGenerator, DummyEncodecInputGenerator, + DummyFluxTransformerTextInputGenerator, + DummyFluxTransformerVisionInputGenerator, DummyInputGenerator, DummyIntGenerator, DummyPastKeyValuesGenerator, + DummyPatchTSTInputGenerator, DummyPix2StructInputGenerator, DummyPointsGenerator, DummySeq2SeqDecoderTextInputGenerator, @@ -38,6 +42,9 @@ DummySpeechT5InputGenerator, DummyTextInputGenerator, DummyTimestepInputGenerator, + DummyTransformerTextInputGenerator, + DummyTransformerTimestepInputGenerator, + DummyTransformerVisionInputGenerator, DummyVisionEmbeddingsGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator, DummyVisionInputGenerator, @@ -45,6 +52,8 @@ FalconDummyPastKeyValuesGenerator, GemmaDummyPastKeyValuesGenerator, GPTBigCodeDummyPastKeyValuesGenerator, + LongformerDummyTextInputGenerator, + MCTCTDummyAudioInputGenerator, MistralDummyPastKeyValuesGenerator, NormalizedConfig, NormalizedEncoderDecoderConfig, @@ -52,9 +61,11 @@ NormalizedTextAndVisionConfig, NormalizedTextConfig, NormalizedTextConfigWithGQA, + NormalizedTimeSeriesForecastingConfig, NormalizedVisionConfig, - check_if_transformers_greater, is_diffusers_available, + is_diffusers_version, + is_transformers_version, logging, ) from ...utils.normalized_config import NormalizedConfigManager @@ -74,6 +85,7 @@ from .model_patcher import ( CLIPModelPatcher, FalconModelPatcher, + MgpstrModelPatcher, MistralModelPatcher, MusicgenModelPatcher, SAMModelPatcher, @@ -85,6 +97,9 @@ ) +# TODO : moved back onnx imports applied in https://github.com/huggingface/optimum/pull/2114/files after refactorization + + if TYPE_CHECKING: from transformers import PretrainedConfig from transformers.modeling_utils import PreTrainedModel @@ -154,6 +169,27 @@ class SplinterOnnxConfig(BertOnnxConfig): DEFAULT_ONNX_OPSET = 11 +class RemBertOnnxConfig(BertOnnxConfig): + DEFAULT_ONNX_OPSET = 11 + + +class LongformerOnnxConfig(BertOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (LongformerDummyTextInputGenerator,) + DEFAULT_ONNX_OPSET = 14 + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + inputs = super().inputs + + inputs["global_attention_mask"] = inputs["attention_mask"] + + return inputs + + +class MegatronBertOnnxConfig(BertOnnxConfig): + DEFAULT_ONNX_OPSET = 11 + + class DistilBertOnnxConfig(BertOnnxConfig): DEFAULT_ONNX_OPSET = 14 # now uses F.scaled_dot_product_attention by default for transformers>=4.46.0 @@ -166,6 +202,10 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return {"input_ids": dynamic_axis, "attention_mask": dynamic_axis} +class ModernBertOnnxConfig(DistilBertOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.48.0") + + class MPNetOnnxConfig(DistilBertOnnxConfig): DEFAULT_ONNX_OPSET = 12 # For lower opsets, results in: Type 'tensor(int64)' of input parameter (/0/auto_model/encoder/Add_1_output_0) of operator (Min) in node (/0/auto_model/encoder/Min) is invalid. @@ -256,6 +296,30 @@ class ImageGPTOnnxConfig(GPT2OnnxConfig): pass +class DecisionTransformerOnnxConfig(OnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (DummyDecisionTransformerInputGenerator,) + NORMALIZED_CONFIG_CLASS = NormalizedConfig + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return { + "states": {0: "batch_size", 1: "sequence_length"}, + "actions": {0: "batch_size", 1: "sequence_length"}, + "timesteps": {0: "batch_size", 1: "sequence_length"}, + "returns_to_go": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "state_preds": {0: "batch_size", 1: "sequence_length"}, + "action_preds": {0: "batch_size", 1: "sequence_length"}, + "return_preds": {0: "batch_size", 1: "sequence_length"}, + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + + class GPTNeoOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): DEFAULT_ONNX_OPSET = 14 NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args(num_attention_heads="num_heads") @@ -267,7 +331,7 @@ class GPTNeoXOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): # OPT does not take position_ids as input for transfomers < v4.46, needs it for transformers >= v4.46 -if check_if_transformers_greater("4.45.99"): +if is_transformers_version(">=", "4.45.99"): class OPTOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): DEFAULT_ONNX_OPSET = 14 # uses SDPA in Transformers, hence opset>=14. @@ -288,6 +352,15 @@ class LlamaOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig +class OlmoOnnxConfig(LlamaOnnxConfig): + ATOL_FOR_VALIDATION = 1e-4 + MIN_TRANSFORMERS_VERSION = version.parse("4.40.0") + + +class Olmo2OnnxConfig(OlmoOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.47.0") + + class Qwen2OnnxConfig(LlamaOnnxConfig): MIN_TRANSFORMERS_VERSION = version.parse("4.37.0") @@ -295,7 +368,12 @@ class Qwen2OnnxConfig(LlamaOnnxConfig): class GemmaOnnxConfig(LlamaOnnxConfig): DUMMY_INPUT_GENERATOR_CLASSES = (DummyTextInputGenerator, GemmaDummyPastKeyValuesGenerator) DUMMY_PKV_GENERATOR_CLASS = GemmaDummyPastKeyValuesGenerator - pass + MIN_TRANSFORMERS_VERSION = version.parse("4.38.0") + + +class GraniteOnnxConfig(LlamaOnnxConfig): + MIN_TRANSFORMERS_VERSION = version.parse("4.45.0") + MIN_TORCH_VERSION = version.parse("2.5.0") class PhiOnnxConfig(TextDecoderWithPositionIdsOnnxConfig): @@ -313,8 +391,7 @@ class Phi3OnnxConfig(PhiOnnxConfig): MIN_TRANSFORMERS_VERSION = version.parse("4.41.0") def __init__(self, *args, **kwargs): - # TODO : replace check_if_transformers_greater with is_transformers_available - if check_if_transformers_greater("4.46.0") and not check_if_transformers_greater("4.46.1"): + if is_transformers_version("==", "4.46.0"): logger.error( "Found transformers v4.46.0 while trying to exporting a Phi3 model, this specific version of transformers is not supported. " "Please upgrade to v4.46.1 or higher, or downgrade your transformers version" @@ -343,6 +420,8 @@ def patch_model_for_export( class MPTOnnxConfig(TextDecoderOnnxConfig): # MPT does not require position_ids input. DEFAULT_ONNX_OPSET = 13 + # TODO: fix inference for transformers < v4.41 for beam_search > 1 + MIN_TRANSFORMERS_VERSION = version.parse("4.41.0") NORMALIZED_CONFIG_CLASS = NormalizedTextConfig.with_args( num_attention_heads="n_heads", hidden_size="d_model", num_layers="n_layers" ) @@ -358,7 +437,7 @@ class BloomOnnxConfig(TextDecoderOnnxConfig): DEFAULT_ONNX_OPSET = 14 # Bloom uses aten::triu that requires opset>=14, and F.scaled_dot_product_attention def add_past_key_values(self, inputs_or_outputs: Dict[str, Dict[int, str]], direction: str): - if check_if_transformers_greater("4.44"): + if is_transformers_version(">=", "4.44"): super().add_past_key_values(inputs_or_outputs, direction) else: if direction not in ["inputs", "outputs"]: @@ -712,7 +791,6 @@ def flatten_past_key_values(self, flattened_output, name, idx, t): class BartOnnxConfig(M2M100OnnxConfig): DEFAULT_ONNX_OPSET = 14 # Bart now uses F.scaled_dot_product_attention by default for torch>=2.1.1. MIN_TORCH_VERSION = version.parse("2.1.2") - pass class MBartOnnxConfig(BartOnnxConfig): @@ -727,21 +805,19 @@ class BlenderbotSmallOnnxConfig(BartOnnxConfig): pass -# big_bird and bigbird_pegasus are unsupported for now as block sparse attention is written in pure python and numpy in transformers. -# Thus, the case attention_type == "block_sparse" is unusable. -# Even with rewritting this part in pure PyTorch, torch.onnx.export is then prohibitively slow. -# References: https://github.com/pytorch/pytorch/issues/63734 & https://github.com/pytorch/pytorch/issues/94821 -""" class BigBirdOnnxConfig(DistilBertOnnxConfig): pass + class BigBirdPegasusOnnxConfig(BartOnnxConfig): - def generate_dummy_inputs_for_validation(self, reference_model_inputs: Dict[str, Any]) -> Dict[str, Any]: - if self._behavior is ConfigBehavior.ENCODER: - # TODO: check why the attention mask is not present in the exported model - reference_model_inputs.pop("attention_mask") - return super().generate_dummy_inputs_for_validation(reference_model_inputs) -""" + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + inputs = super().inputs + if self._config.attention_type == "block_sparse": + # BigBirdPegasusEncoder creates its own attention_mask internally + # https://github.com/huggingface/transformers/blob/v4.48.0/src/transformers/models/bigbird_pegasus/modeling_bigbird_pegasus.py#L1875 + inputs.pop("attention_mask", None) + return inputs class PegasusOnnxConfig(BartOnnxConfig): @@ -764,8 +840,10 @@ def inputs(self) -> Dict[str, Dict[int, str]]: @property def outputs(self) -> Dict[str, Dict[int, str]]: common_outputs = super().outputs + if self.task == "feature-extraction": common_outputs["last_hidden_state"] = {0: "batch_size"} + return common_outputs @@ -783,7 +861,7 @@ class DeiTOnnxConfig(ViTOnnxConfig): class BeitOnnxConfig(ViTOnnxConfig): - DEFAULT_ONNX_OPSET = 11 + DEFAULT_ONNX_OPSET = 14 # now uses F.scaled_dot_product_attention by default for torch>=2.1.1. class ConvNextOnnxConfig(ViTOnnxConfig): @@ -794,6 +872,65 @@ class ConvNextV2OnnxConfig(ViTOnnxConfig): DEFAULT_ONNX_OPSET = 11 +class HieraOnnxConfig(ViTOnnxConfig): + DEFAULT_ONNX_OPSET = 11 + + +class PvtOnnxConfig(ViTOnnxConfig): + DEFAULT_ONNX_OPSET = 11 + + +class VitMAEOnnxConfig(ViTOnnxConfig): + # torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::scaled_dot_product_attention' to ONNX opset version 11 is not supported. + # Support for this operator was added in version 14, try exporting with this version. + DEFAULT_ONNX_OPSET = 14 + + +class VitMSNOnnxConfig(ViTOnnxConfig): + # torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::scaled_dot_product_attention' to ONNX opset version 11 is not supported. + # Support for this operator was added in version 14, try exporting with this version. + DEFAULT_ONNX_OPSET = 14 + + +class Dinov2DummyInputGenerator(DummyVisionInputGenerator): + 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=task, + normalized_config=normalized_config, + batch_size=batch_size, + num_channels=num_channels, + width=width, + height=height, + **kwargs, + ) + + from transformers.onnx.utils import get_preprocessor + + preprocessor = get_preprocessor(normalized_config._name_or_path) + if preprocessor is not None and hasattr(preprocessor, "crop_size"): + self.height = preprocessor.crop_size.get("height", self.height) + self.width = preprocessor.crop_size.get("width", self.width) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + input_ = super().generate( + input_name=input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype + ) + return input_ + + +class Dinov2OnnxConfig(ViTOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = (Dinov2DummyInputGenerator,) + + class MobileViTOnnxConfig(ViTOnnxConfig): ATOL_FOR_VALIDATION = 1e-4 DEFAULT_ONNX_OPSET = 11 @@ -835,6 +972,10 @@ class SwinOnnxConfig(ViTOnnxConfig): DEFAULT_ONNX_OPSET = 11 +class SwinV2OnnxConfig(SwinOnnxConfig): + pass + + class Swin2srOnnxConfig(SwinOnnxConfig): pass @@ -854,7 +995,14 @@ class PoolFormerOnnxConfig(ViTOnnxConfig): class SegformerOnnxConfig(YolosOnnxConfig): - pass + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + outputs = super().outputs + + if self.task == "image-segmentation": + outputs["logits"] = {0: "batch_size"} + + return outputs class MobileNetV1OnnxConfig(ViTOnnxConfig): @@ -870,6 +1018,28 @@ class MobileNetV2OnnxConfig(MobileNetV1OnnxConfig): pass +class MaskFormerOnnxConfig(ViTOnnxConfig): + # torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::einsum' to ONNX opset version 11 is not supported. + # Support for this operator was added in version 12, try exporting with this version. + DEFAULT_ONNX_OPSET = 12 + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "image-segmentation": + return { + "class_queries_logits": {0: "batch_size", 1: "num_queries"}, + "masks_queries_logits": {0: "batch_size", 1: "num_queries", 2: "height", 3: "width"}, + } + else: + return super().outputs + + @property + def torch_to_onnx_output_map(self) -> Dict[str, str]: + return { + "transformer_decoder_last_hidden_state": "last_hidden_state", + } + + class DonutSwinOnnxConfig(ViTOnnxConfig): DEFAULT_ONNX_OPSET = 11 @@ -890,6 +1060,21 @@ def torch_to_onnx_input_map(self) -> Dict[str, str]: return {"x": "pixel_values"} +class MgpstrOnnxConfig(ViTOnnxConfig): + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "char_logits": {0: "batch_size"}, + "bpe_logits": {0: "batch_size"}, + "wp_logits": {0: "batch_size"}, + } + + def patch_model_for_export( + self, model: Union["PreTrainedModel", "TFPreTrainedModel"], model_kwargs: Optional[Dict[str, Any]] = None + ) -> "ModelPatcher": + return MgpstrModelPatcher(self, model, model_kwargs=model_kwargs) + + class SentenceTransformersTransformerOnnxConfig(TextEncoderOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedTextConfig DEFAULT_ONNX_OPSET = 14 # Some bottleneck transformers models require a specific ONNX opset to be successfully exported. We put a rather high opset here for the export to work for all architectures. @@ -1032,22 +1217,13 @@ def outputs(self) -> Dict[str, Dict[int, str]]: "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 - def generate_dummy_inputs(self, framework: str = "pt", **kwargs): - dummy_inputs = super().generate_dummy_inputs(framework=framework, **kwargs) - - # TODO: fix should be by casting inputs during inference and not export - if framework == "pt": - import torch - - dummy_inputs["input_ids"] = dummy_inputs["input_ids"].to(dtype=torch.int32) - return dummy_inputs - def patch_model_for_export( self, model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], @@ -1056,8 +1232,41 @@ def patch_model_for_export( return CLIPModelPatcher(self, model, model_kwargs=model_kwargs) +class SiglipNormalizedConfig(CLIPNormalizedConfig): + pass + + +class SiglipOnnxConfig(CLIPOnnxConfig): + NORMALIZED_CONFIG_CLASS = SiglipNormalizedConfig + # torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::scaled_dot_product_attention' to ONNX opset version 13 is not supported. + # Support for this operator was added in version 14, try exporting with this version. + DEFAULT_ONNX_OPSET = 14 + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + 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 + } + + +class SiglipTextWithProjectionOnnxConfig(CLIPTextWithProjectionOnnxConfig): + pass + + +class SiglipTextOnnxConfig(CLIPTextOnnxConfig): + pass + + +class SiglipVisionModelOnnxConfig(CLIPVisionModelOnnxConfig): + # torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::scaled_dot_product_attention' to ONNX opset version 11 is not supported. + # Support for this operator was added in version 14, try exporting with this version. + DEFAULT_ONNX_OPSET = 14 + + class UNetOnnxConfig(VisionOnnxConfig): - ATOL_FOR_VALIDATION = 1e-3 + ATOL_FOR_VALIDATION = 1e-4 # The ONNX export of a CLIPText architecture, an other Stable Diffusion component, needs the Trilu # operator support, available since opset 14 DEFAULT_ONNX_OPSET = 14 @@ -1080,17 +1289,19 @@ class UNetOnnxConfig(VisionOnnxConfig): def inputs(self) -> Dict[str, Dict[int, str]]: common_inputs = { "sample": {0: "batch_size", 2: "height", 3: "width"}, - "timestep": {0: "steps"}, + "timestep": {}, # a scalar with no dimension "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, } - # TODO : add text_image, image and image_embeds + # TODO : add addition_embed_type == text_image, image and image_embeds + # https://github.com/huggingface/diffusers/blob/9366c8f84bfe47099ff047272661786ebb54721d/src/diffusers/models/unets/unet_2d_condition.py#L671 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"} + return common_inputs @property @@ -1129,7 +1340,7 @@ def ordered_inputs(self, model) -> Dict[str, Dict[int, str]]: class VaeEncoderOnnxConfig(VisionOnnxConfig): - ATOL_FOR_VALIDATION = 1e-4 + ATOL_FOR_VALIDATION = 3e-4 # The ONNX export of a CLIPText architecture, an other Stable Diffusion component, needs the Trilu # operator support, available since opset 14 DEFAULT_ONNX_OPSET = 14 @@ -1177,6 +1388,101 @@ def outputs(self) -> Dict[str, Dict[int, str]]: } +class T5EncoderOnnxConfig(TextEncoderOnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTextConfig + ATOL_FOR_VALIDATION = 1e-4 + DEFAULT_ONNX_OPSET = 12 # int64 was supported since opset 12 + + @property + def inputs(self): + return { + "input_ids": {0: "batch_size", 1: "sequence_length"}, + } + + @property + def outputs(self): + return { + "last_hidden_state": {0: "batch_size", 1: "sequence_length"}, + } + + +class SD3TransformerOnnxConfig(VisionOnnxConfig): + ATOL_FOR_VALIDATION = 1e-4 + # The ONNX export of a CLIPText architecture, an other Stable Diffusion component, needs the Trilu + # operator support, available since opset 14 + DEFAULT_ONNX_OPSET = 14 + + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTransformerTimestepInputGenerator, + DummyTransformerVisionInputGenerator, + DummyTransformerTextInputGenerator, + ) + + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + image_size="sample_size", + num_channels="in_channels", + vocab_size="attention_head_dim", + hidden_size="joint_attention_dim", + projection_size="pooled_projection_dim", + allow_new=True, + ) + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + common_inputs = { + "hidden_states": {0: "batch_size", 2: "height", 3: "width"}, + "encoder_hidden_states": {0: "batch_size", 1: "sequence_length"}, + "pooled_projections": {0: "batch_size"}, + "timestep": {0: "step"}, + } + + return common_inputs + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + return { + "out_hidden_states": {0: "batch_size", 2: "height", 3: "width"}, + } + + @property + def torch_to_onnx_output_map(self) -> Dict[str, str]: + return { + "sample": "out_hidden_states", + } + + +class FluxTransformerOnnxConfig(SD3TransformerOnnxConfig): + DUMMY_INPUT_GENERATOR_CLASSES = ( + DummyTransformerTimestepInputGenerator, + DummyFluxTransformerVisionInputGenerator, + DummyFluxTransformerTextInputGenerator, + ) + + @property + def inputs(self): + common_inputs = super().inputs + common_inputs["hidden_states"] = {0: "batch_size", 1: "packed_height_width"} + common_inputs["txt_ids"] = ( + {0: "sequence_length"} if is_diffusers_version(">=", "0.31.0") else {0: "batch_size", 1: "sequence_length"} + ) + common_inputs["img_ids"] = ( + {0: "packed_height_width"} + if is_diffusers_version(">=", "0.31.0") + else {0: "batch_size", 1: "packed_height_width"} + ) + + if getattr(self._normalized_config, "guidance_embeds", False): + common_inputs["guidance"] = {0: "batch_size"} + + return common_inputs + + @property + def outputs(self): + return { + "out_hidden_states": {0: "batch_size", 1: "packed_height_width"}, + } + + class GroupViTOnnxConfig(CLIPOnnxConfig): pass @@ -1292,13 +1598,12 @@ class Data2VecTextOnnxConfig(DistilBertOnnxConfig): class Data2VecVisionOnnxConfig(ViTOnnxConfig): - DEFAULT_ONNX_OPSET = 11 + DEFAULT_ONNX_OPSET = 14 # now uses F.scaled_dot_product_attention by default for torch>=2.1.1. class Data2VecAudioOnnxConfig(AudioOnnxConfig): - NORMALIZED_CONFIG_CLASS = NormalizedConfig - ATOL_FOR_VALIDATION = 1e-4 DEFAULT_ONNX_OPSET = 14 # now uses F.scaled_dot_product_attention by default for torch>=2.1.1. + NORMALIZED_CONFIG_CLASS = NormalizedConfig class PerceiverDummyInputGenerator(DummyVisionInputGenerator): @@ -1385,6 +1690,17 @@ def inputs(self) -> Dict[str, Dict[int, str]]: "pixel_values": dynamic_axis, } + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + outputs = super().outputs + + if "logits" in outputs: + # default is {0: "batch_size", 1: "sequence_length"} where sequence_length is dynamic axis + # but perceiver always return the same max sequence length in the second dimension + 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) @@ -1454,23 +1770,16 @@ def inputs(self) -> Dict[str, Dict[int, str]]: return {"input_values": {0: "batch_size"}} -# TODO: currently disabled because an operator seems not supported by ONNX. -# class MCTCTDummyAudioInputGenerator(DummyAudioInputGenerator): -# def generate(self, input_name: str, framework: str = "pt"): -# shape = [self.batch_size, self.sequence_length, self.normalized_config.input_features_per_channel] -# if input_name == "input_features": -# return self.random_float_tensor(shape, min_value=-1, max_value=1, framework=framework) -# return super().generate(input_name, framework=framework) -# -# -# class MCTCTOnnxConfig(OnnxConfig): -# NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args(input_features_per_channel="input_feat_per_channel", allow_new=True) -# DUMMY_INPUT_GENERATOR_CLASSES = (MCTCTDummyAudioInputGenerator,) -# DEFAULT_ONNX_OPSET = 13 -# -# @property -# def inputs(self) -> Dict[str, Dict[int, str]]: -# return {"input_features": {0: "batch_size", 1: "sequence_classification"}} +class MCTCTOnnxConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedConfig.with_args( + input_features_per_channel="input_feat_per_channel", allow_new=True + ) + DUMMY_INPUT_GENERATOR_CLASSES = (MCTCTDummyAudioInputGenerator,) + DEFAULT_ONNX_OPSET = 13 + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return {"input_features": {0: "batch_size", 1: "sequence_classification"}} class WhisperOnnxConfig(AudioToTextOnnxConfig): @@ -1494,7 +1803,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: common_inputs["input_features"] = {0: "batch_size"} # Remove unnecessary dynamic axis. if self._behavior is not ConfigBehavior.ENCODER and self.use_past_in_inputs: - if check_if_transformers_greater("4.43.0"): + if is_transformers_version(">=", "4.43.0"): # since https://github.com/huggingface/transformers/pull/31166 common_inputs["cache_position"] = {0: "decoder_sequence_length"} @@ -1744,6 +2053,8 @@ def post_process_exported_models( decoder_with_past_path = Path(path, onnx_files_subpaths[3]) decoder_merged_path = Path(path, ONNX_DECODER_MERGED_NAME + ".onnx") try: + from ...onnx import merge_decoders + # The decoder with past does not output the cross attention past key values as they are constant, # hence the need for strict=False merge_decoders( @@ -2063,6 +2374,7 @@ def inputs(self) -> Dict[str, Dict[int, str]]: if self.use_past_in_inputs: self.add_past_key_values(common_inputs, direction="inputs") + if self._behavior is ConfigBehavior.DECODER: common_inputs["encoder_outputs"] = {0: "batch_size", 1: "encoder_sequence_length"} @@ -2179,12 +2491,7 @@ class Pix2StructOnnxConfig(OnnxSeq2SeqConfigWithPast): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # TODO : replace check_if_transformers_greater with is_transformers_available - if ( - check_if_transformers_greater("4.46.0") - and not check_if_transformers_greater("4.46.1") - and self._behavior is ConfigBehavior.DECODER - ): + if is_transformers_version("==", "4.46.0") and self._behavior is ConfigBehavior.DECODER: logger.error( "Found transformers v4.46.0 while trying to exporting a Pix2Struct model, this specific version of transformers is not supported. " "Please upgrade to v4.46.1 or higher, or downgrade your transformers version" @@ -2343,3 +2650,24 @@ class EncoderDecoderOnnxConfig(EncoderDecoderBaseOnnxConfig): NORMALIZED_CONFIG_CLASS = NormalizedEncoderDecoderConfig DEFAULT_ONNX_OPSET = 14 # uses SDPA in Transformers, hence opset>=14. + + +class PatchTSTOnnxConfig(OnnxConfig): + NORMALIZED_CONFIG_CLASS = NormalizedTimeSeriesForecastingConfig + DUMMY_INPUT_GENERATOR_CLASSES = (DummyPatchTSTInputGenerator,) + ATOL_FOR_VALIDATION = 1e-4 + + @property + def inputs(self) -> Dict[str, Dict[int, str]]: + return {"past_values": {0: "batch_size", 1: "sequence_length"}} + + @property + def outputs(self) -> Dict[str, Dict[int, str]]: + if self.task == "feature-extraction": + return {"last_hidden_state": {0: "batch_size"}} + else: + return super().outputs + + +class PatchTSMixerOnnxConfig(PatchTSTOnnxConfig): + pass diff --git a/optimum/exporters/onnx/model_patcher.py b/optimum/exporters/onnx/model_patcher.py index fdfb0e280f..53476ff206 100644 --- a/optimum/exporters/onnx/model_patcher.py +++ b/optimum/exporters/onnx/model_patcher.py @@ -20,35 +20,34 @@ import types from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +import torch import transformers -from packaging import version from transformers.models.speecht5.modeling_speecht5 import SpeechT5EncoderWithSpeechPrenet -from transformers.utils import is_torch_available +from ...utils import is_transformers_version, logging +from ._traceable_cache import TraceableCache -if is_torch_available(): - import torch -from ...configuration_utils import _transformers_version -from ...utils import logging - - -if _transformers_version > version.parse("4.34.99"): +if is_transformers_version(">=", "4.35"): from transformers.modeling_attn_mask_utils import AttentionMaskConverter -if _transformers_version >= version.parse("4.36"): +if is_transformers_version(">=", "4.36"): from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask_for_sdpa -else: - _prepare_4d_causal_attention_mask_for_sdpa = None - AttentionMaskConverter = None - -if _transformers_version >= version.parse("4.42"): +if is_transformers_version(">=", "4.43"): + from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention +if is_transformers_version(">=", "4.42"): from transformers.cache_utils import SlidingWindowCache, StaticCache +if is_transformers_version(">=", "4.48"): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + from transformers.integrations.sdpa_attention import repeat_kv, sdpa_attention_forward + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + if TYPE_CHECKING: from transformers import PreTrainedModel, TFPreTrainedModel from .base import OnnxConfig + logger = logging.get_logger(__name__) @@ -113,6 +112,54 @@ class PatchingSpec: op_wrapper: Optional[Callable] = None +# An ONNX-export-compatible version of `tensor.unfold`. Without this, we get: +# torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator Unfold, input size not accessible. +# See https://github.com/pytorch/pytorch/issues/81871 for more information +def onnx_compatible_unfold(input_tensor, dimension, size, step): + """ + Custom implementation of torch.unfold without using torch.unfold. + + Args: + input_tensor (torch.Tensor): The input tensor. + dimension (int): The dimension to unfold. + size (int): The size of each slice. + step (int): The step size between slices. + + Returns: + torch.Tensor: The unfolded tensor. + """ + # Check if dimension is within the valid range + if not (-input_tensor.dim() <= dimension < input_tensor.dim()): + raise ValueError( + f"Dimension out of range (expected to be in range of [{-input_tensor.dim()}, {input_tensor.dim() - 1}], but got {dimension})" + ) + + # Normalize negative dimension + dimension = dimension % input_tensor.dim() + + # Compute the shape of the unfolded output + input_size = input_tensor.size(dimension) + num_slices = (input_size - size) // step + 1 + + # Permute dimension to the end for easier indexing + input_tensor = input_tensor.transpose(dimension, -1) + + # Extract slices + slices = [] + for i in range(num_slices): + start = i * step + end = start + size + slices.append(input_tensor[..., start:end]) + + # Stack slices and permute dimensions back + result = torch.stack(slices, dim=-2).transpose(dimension, -2) + return result + + +UNSUPPORTED_OPS_PATCHING_SPEC = [PatchingSpec(torch.Tensor, "unfold", onnx_compatible_unfold, torch.Tensor.unfold)] +CACHE_PATCHING_SPEC = [PatchingSpec(transformers.cache_utils, "Cache", TraceableCache, transformers.cache_utils.Cache)] + + class ModelPatcher: def __init__( self, @@ -122,9 +169,12 @@ def __init__( ): self._model = model - patching_specs = config.PATCHING_SPECS + patching_specs = config.PATCHING_SPECS or [] + patching_specs.extend(UNSUPPORTED_OPS_PATCHING_SPEC) + patching_specs.extend(CACHE_PATCHING_SPEC) + self._patching_specs = [] - for spec in patching_specs if patching_specs is not None else []: + 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)) @@ -148,6 +198,39 @@ def patched_forward(*args, **kwargs): signature = inspect.signature(self.orig_forward) args, kwargs = override_arguments(args, kwargs, signature, model_kwargs=self.model_kwargs) + if is_transformers_version(">=", "4.48"): + if "past_key_values" in signature.parameters: + pkv_index = list(signature.parameters.keys()).index("past_key_values") + + if ( + pkv_index < len(args) # pkv is in args + and isinstance(args[pkv_index], (list, tuple)) + and isinstance(args[pkv_index][0], (list, tuple)) + ): + if len(args[pkv_index][0]) == 2: + args[pkv_index] = DynamicCache.from_legacy_cache(args[pkv_index]) + elif len(args[pkv_index][0]) == 4: + args[pkv_index] = EncoderDecoderCache.from_legacy_cache(args[pkv_index]) + else: + raise ValueError( + f"past_key_values should have either 2 or 4 elements, but it has {len(args[pkv_index][0])} elements" + ) + elif ( + "past_key_values" in kwargs # pkv is in kwargs + and isinstance(kwargs["past_key_values"], (list, tuple)) + and isinstance(kwargs["past_key_values"][0], (list, tuple)) + ): + if len(kwargs["past_key_values"][0]) == 2: + kwargs["past_key_values"] = DynamicCache.from_legacy_cache(kwargs["past_key_values"]) + elif len(kwargs["past_key_values"][0]) == 4: + kwargs["past_key_values"] = EncoderDecoderCache.from_legacy_cache( + kwargs["past_key_values"] + ) + else: + raise ValueError( + f"past_key_values should have either 2 or 4 elements, but it has {len(kwargs['past_key_values'][0])} elements" + ) + outputs = self.orig_forward(*args, **kwargs) # This code block handles different cases of the filterd_outputs input to align it with the expected @@ -168,7 +251,7 @@ def patched_forward(*args, **kwargs): filterd_outputs[name] = value elif isinstance(outputs, (list, tuple)): outputs_list = list(config.outputs.keys()) - dict(zip(outputs_list, outputs)) + filterd_outputs = dict(zip(outputs_list, outputs)) else: if len(config.outputs) > 1: num_outputs = len(config.outputs) @@ -181,6 +264,11 @@ def patched_forward(*args, **kwargs): filterd_outputs[name] = outputs name = list(config.outputs.keys())[0] filterd_outputs[name] = outputs + + if is_transformers_version(">=", "4.48"): + if isinstance(filterd_outputs.get("past_key_values"), (DynamicCache, EncoderDecoderCache)): + filterd_outputs["past_key_values"] = outputs["past_key_values"].to_legacy_cache() + return filterd_outputs self.patched_forward = patched_forward @@ -210,6 +298,18 @@ def __call__(self, *args, **kwargs): class Seq2SeqModelPatcher(ModelPatcher): + def __enter__(self): + super().__enter__() + if is_transformers_version(">=", "4.48"): + # this is required when gpt2 is used as decoder in any + # encoder-decoder model with cross attention blocks + ALL_ATTENTION_FUNCTIONS["sdpa"] = patched_sdpa_attention_forward + + def __exit__(self, exc_type, exc_value, traceback): + super().__exit__(exc_type, exc_value, traceback) + if is_transformers_version(">=", "4.48"): + ALL_ATTENTION_FUNCTIONS["sdpa"] = sdpa_attention_forward + def __init__( self, config: "OnnxConfig", @@ -261,6 +361,51 @@ def patched_forward(*args, **kwargs): self.patched_forward = patched_forward +def patched_sdpa_attention_forward( + module: torch.nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + dropout: float = 0.0, + scaling: Optional[float] = None, + is_causal: Optional[bool] = None, + **kwargs, +) -> Tuple[torch.Tensor, None]: + if hasattr(module, "num_key_value_groups"): + key = repeat_kv(key, module.num_key_value_groups) + value = repeat_kv(value, module.num_key_value_groups) + causal_mask = attention_mask + if attention_mask is not None: + causal_mask = causal_mask[:, :, :, : key.shape[-2]] + # SDPA with memory-efficient backend is bugged with non-contiguous inputs and custom attn_mask for some torch versions + # Reference: https://github.com/pytorch/pytorch/issues/112577. + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + if is_causal is None: + is_causal = causal_mask is None and query.shape[2] > 1 + + # Shapes (e.g. query.shape[2]) are tensors during jit tracing, resulting in `is_causal` being a tensor. + # We convert it to a bool for the SDPA kernel that only accepts bools. + if torch.jit.is_tracing() and isinstance(is_causal, torch.Tensor): + is_causal = is_causal.item() + + attn_output = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=causal_mask, + dropout_p=dropout, + scale=scaling, + is_causal=is_causal, + ) + attn_output = attn_output.transpose(1, 2).contiguous() + return attn_output, None + + class VisionEncoderDecoderPatcher(Seq2SeqModelPatcher): def __init__( self, @@ -275,14 +420,17 @@ def __init__( model.decoder.model.decoder.config.use_cache = True -def _unmask_unattended_patched_legacy( - expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float] -): - return expanded_mask +if is_transformers_version(">=", "4.39"): + def _unmask_unattended_patched(expanded_mask: torch.Tensor, min_dtype: float): + return expanded_mask -def _unmask_unattended_patched(expanded_mask: torch.Tensor, min_dtype: float): - return expanded_mask +else: + + def _unmask_unattended_patched( + expanded_mask: torch.Tensor, attention_mask: torch.Tensor, unmasked_value: Union[bool, float] + ): + return expanded_mask def _make_causal_mask_patched( @@ -317,14 +465,6 @@ def _make_causal_mask_patched( return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) -_make_causal_mask_patched_staticmethod = staticmethod(_make_causal_mask_patched) - -if _transformers_version >= version.parse("4.39.0"): - _unmask_unattended_patched_staticmethod = staticmethod(_unmask_unattended_patched) -else: - _unmask_unattended_patched_staticmethod = staticmethod(_unmask_unattended_patched_legacy) - - # Adapted from _prepare_4d_causal_attention_mask def _prepare_4d_causal_attention_mask_for_sdpa_patched( attention_mask: Optional[torch.Tensor], @@ -363,28 +503,22 @@ def _prepare_4d_causal_attention_mask_for_sdpa_patched( class DecoderModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() - if AttentionMaskConverter is not None: - # TODO: Remove this _make_causal_mask patch if once transformers if much above 4.35 - AttentionMaskConverter._make_causal_mask = _make_causal_mask_patched_staticmethod - - if _transformers_version >= version.parse("4.36"): - AttentionMaskConverter._unmask_unattended = _unmask_unattended_patched_staticmethod + if is_transformers_version(">=", "4.35"): + AttentionMaskConverter._make_causal_mask = staticmethod(_make_causal_mask_patched) - if _transformers_version >= version.parse("4.36"): + if is_transformers_version(">=", "4.36"): + AttentionMaskConverter._unmask_unattended = staticmethod(_unmask_unattended_patched) patch_everywhere( "_prepare_4d_causal_attention_mask_for_sdpa", _prepare_4d_causal_attention_mask_for_sdpa_patched ) def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) - if AttentionMaskConverter is not None: - # TODO: Remove this _make_causal_mask patch if once transformers if much above 4.35 - AttentionMaskConverter._make_causal_mask = staticmethod(self.original_make_causal) + if is_transformers_version(">=", "4.35"): + AttentionMaskConverter._make_causal_mask = staticmethod(self.original_make_causal_mask) - if _transformers_version >= version.parse("4.36"): - AttentionMaskConverter._unmask_unattended = staticmethod(self.original_unmask_unattended) - - if _transformers_version >= version.parse("4.36"): + if is_transformers_version(">=", "4.36"): + AttentionMaskConverter._unmask_unattended = staticmethod(self.original_unmask_unattended) patch_everywhere( "_prepare_4d_causal_attention_mask_for_sdpa", self.original_prepare_4d_causal_attention_mask_for_sdpa ) @@ -397,13 +531,12 @@ def __init__( ): super().__init__(config, model, model_kwargs) - if _transformers_version >= version.parse("4.36"): - self.original_prepare_4d_causal_attention_mask_for_sdpa = _prepare_4d_causal_attention_mask_for_sdpa - self.original_unmask_unattended = AttentionMaskConverter._unmask_unattended + if is_transformers_version(">=", "4.35"): + self.original_make_causal_mask = AttentionMaskConverter._make_causal_mask - # TODO: Remove this if once transformers if much above 4.35 - if AttentionMaskConverter is not None: - self.original_make_causal = AttentionMaskConverter._make_causal_mask + if is_transformers_version(">=", "4.36"): + self.original_unmask_unattended = AttentionMaskConverter._unmask_unattended + self.original_prepare_4d_causal_attention_mask_for_sdpa = _prepare_4d_causal_attention_mask_for_sdpa def falcon_build_alibi_tensor_patched( @@ -509,6 +642,32 @@ def patched_forward(*args, **kwargs): self.patched_forward = patched_forward +class MgpstrModelPatcher(ModelPatcher): + def __init__( + self, + config: "OnnxConfig", + model: Union["PreTrainedModel", "TFPreTrainedModel"], + model_kwargs: Optional[Dict[str, Any]] = None, + ): + super().__init__(config, model, model_kwargs) + + @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) + + # logits is a tuple, so we unpack it and return them as separate outputs + char_logits, bpe_logits, wp_logits = self.orig_forward(*args, **kwargs).logits + + return { + "char_logits": char_logits, + "bpe_logits": bpe_logits, + "wp_logits": wp_logits, + } + + self.patched_forward = patched_forward + + class SAMModelPatcher(ModelPatcher): def __init__( self, @@ -758,14 +917,22 @@ def patched_forward( class SentenceTransformersTransformerPatcher(ModelPatcher): def __enter__(self): super().__enter__() - if _transformers_version >= version.parse("4.42") and self.real_config._config.model_type == "mistral": + if ( + is_transformers_version(">=", "4.42") + and is_transformers_version("<", "4.48") + and self.real_config._config.model_type == "mistral" + ): self._model[0].auto_model._update_causal_mask = types.MethodType( _update_causal_mask_patched, self._model[0].auto_model ) def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) - if _transformers_version >= version.parse("4.42") and self.real_config._config.model_type == "mistral": + if ( + is_transformers_version(">=", "4.42") + and is_transformers_version("<", "4.48") + and self.real_config._config.model_type == "mistral" + ): self._model[0].auto_model._update_causal_mask = types.MethodType( self._update_causal_mask_original, self._model[0].auto_model ) @@ -778,7 +945,11 @@ def __init__( ): super().__init__(config, model, model_kwargs) - if _transformers_version >= version.parse("4.42") and self.real_config._config.model_type == "mistral": + if ( + is_transformers_version(">=", "4.42") + and is_transformers_version("<", "4.48") + and self.real_config._config.model_type == "mistral" + ): self._update_causal_mask_original = self._model[0].auto_model._update_causal_mask def patched_forward(input_ids, attention_mask): @@ -1057,36 +1228,25 @@ def _update_causal_mask_patched( padding_mask, min_dtype ) - # if ( - # self.config._attn_implementation == "sdpa" - # and attention_mask is not None - # and attention_mask.device.type == "cuda" - # and not output_attentions - # ): - # # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when - # # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. - # # Details: https://github.com/pytorch/pytorch/issues/110213 - # causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) return causal_mask -class MistralModelPatcher(ModelPatcher): +class MistralModelPatcher(DecoderModelPatcher): def __enter__(self): super().__enter__() - if AttentionMaskConverter is not None: - # TODO: Remove this _make_causal_mask patch if once transformers if much above 4.35 - AttentionMaskConverter._make_causal_mask = _make_causal_mask_patched_staticmethod - - if _transformers_version >= version.parse("4.36"): - AttentionMaskConverter._unmask_unattended = _unmask_unattended_patched_staticmethod - - if _transformers_version >= version.parse("4.36"): - patch_everywhere( - "_prepare_4d_causal_attention_mask_for_sdpa", _prepare_4d_causal_attention_mask_for_sdpa_patched - ) - if _transformers_version >= version.parse("4.42"): + if is_transformers_version(">=", "4.42") and is_transformers_version("<", "4.48"): if hasattr(self._model, "model"): self._model.model._update_causal_mask = types.MethodType( _update_causal_mask_patched, self._model.model @@ -1096,19 +1256,8 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, traceback): super().__exit__(exc_type, exc_value, traceback) - if AttentionMaskConverter is not None: - # TODO: Remove this _make_causal_mask patch if once transformers if much above 4.35 - AttentionMaskConverter._make_causal_mask = staticmethod(self.original_make_causal) - if _transformers_version >= version.parse("4.36"): - AttentionMaskConverter._unmask_unattended = staticmethod(self.original_unmask_unattended) - - if _transformers_version >= version.parse("4.36"): - patch_everywhere( - "_prepare_4d_causal_attention_mask_for_sdpa", self.original_prepare_4d_causal_attention_mask_for_sdpa - ) - - if _transformers_version >= version.parse("4.42"): + if is_transformers_version(">=", "4.42") and is_transformers_version("<", "4.48"): if hasattr(self._model, "model"): self._model.model._update_causal_mask = types.MethodType( self._update_causal_mask_original, self._model.model @@ -1124,15 +1273,7 @@ def __init__( ): super().__init__(config, model, model_kwargs) - if _transformers_version >= version.parse("4.36"): - self.original_prepare_4d_causal_attention_mask_for_sdpa = _prepare_4d_causal_attention_mask_for_sdpa - self.original_unmask_unattended = AttentionMaskConverter._unmask_unattended - - # TODO: Remove this if once transformers if much above 4.35 - if AttentionMaskConverter is not None: - self.original_make_causal = AttentionMaskConverter._make_causal_mask - - if _transformers_version >= version.parse("4.42"): + if is_transformers_version(">=", "4.42") and is_transformers_version("<", "4.48"): if hasattr(self._model, "model"): self._update_causal_mask_original = self._model.model._update_causal_mask else: @@ -1142,15 +1283,11 @@ def __init__( class CLIPModelPatcher(ModelPatcher): def __enter__(self): super().__enter__() - - if _transformers_version >= version.parse("4.43"): - from transformers.models.clip.modeling_clip import CLIPAttention, CLIPSdpaAttention - - self.original_sdpa_forward, CLIPSdpaAttention.forward = CLIPSdpaAttention.forward, CLIPAttention.forward + if is_transformers_version(">=", "4.43"): + 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 _transformers_version >= version.parse("4.43"): - from transformers.models.clip.modeling_clip import CLIPSdpaAttention - + if is_transformers_version(">=", "4.43"): CLIPSdpaAttention.forward = self.original_sdpa_forward diff --git a/optimum/exporters/onnx/utils.py b/optimum/exporters/onnx/utils.py index 56249bbf5c..3659480abf 100644 --- a/optimum/exporters/onnx/utils.py +++ b/optimum/exporters/onnx/utils.py @@ -20,14 +20,13 @@ from packaging import version from transformers.utils import is_tf_available, is_torch_available -from ...utils import ( - DIFFUSERS_MINIMUM_VERSION, - ORT_QUANTIZE_MINIMUM_VERSION, - check_if_diffusers_greater, +from ...utils import DIFFUSERS_MINIMUM_VERSION, ORT_QUANTIZE_MINIMUM_VERSION, logging +from ...utils.import_utils import ( + _diffusers_version, is_diffusers_available, - logging, + is_diffusers_version, + is_transformers_version, ) -from ...utils.import_utils import _diffusers_version, check_if_transformers_greater from ..utils import ( _get_submodels_and_export_configs, ) @@ -52,7 +51,7 @@ if is_diffusers_available(): - if not check_if_diffusers_greater(DIFFUSERS_MINIMUM_VERSION.base_version): + if not is_diffusers_version(">=", DIFFUSERS_MINIMUM_VERSION.base_version): raise ImportError( f"We found an older version of diffusers {_diffusers_version} but we require diffusers to be >= {DIFFUSERS_MINIMUM_VERSION}. " "Please update diffusers by running `pip install --upgrade diffusers`" @@ -86,10 +85,11 @@ "phi", "phi3", "qwen2", + "granite", } -if check_if_transformers_greater("4.45.99"): +if is_transformers_version(">=", "4.45.99"): MODEL_TYPES_REQUIRING_POSITION_IDS.add("opt") diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index a489f34fb0..47a6ae08ca 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -209,20 +209,27 @@ class TasksManager: "feature-extraction": "AutoModel", "fill-mask": "AutoModelForMaskedLM", "image-classification": "AutoModelForImageClassification", - "image-segmentation": ("AutoModelForImageSegmentation", "AutoModelForSemanticSegmentation"), + "image-segmentation": ( + "AutoModelForImageSegmentation", + "AutoModelForSemanticSegmentation", + "AutoModelForInstanceSegmentation", + "AutoModelForUniversalSegmentation", + ), "image-to-image": "AutoModelForImageToImage", - "image-to-text": "AutoModelForVision2Seq", + "image-to-text": ("AutoModelForVision2Seq", "AutoModel"), "mask-generation": "AutoModel", "masked-im": "AutoModelForMaskedImageModeling", "multiple-choice": "AutoModelForMultipleChoice", "object-detection": "AutoModelForObjectDetection", "question-answering": "AutoModelForQuestionAnswering", + "reinforcement-learning": "AutoModel", "semantic-segmentation": "AutoModelForSemanticSegmentation", "text-to-audio": ("AutoModelForTextToSpectrogram", "AutoModelForTextToWaveform"), "text-generation": "AutoModelForCausalLM", "text2text-generation": "AutoModelForSeq2SeqLM", "text-classification": "AutoModelForSequenceClassification", "token-classification": "AutoModelForTokenClassification", + "visual-question-answering": "AutoModelForVisualQuestionAnswering", "zero-shot-image-classification": "AutoModelForZeroShotImageClassification", "zero-shot-object-detection": "AutoModelForZeroShotObjectDetection", } @@ -306,6 +313,7 @@ class TasksManager: "vision2seq-lm": "image-to-text", "zero-shot-classification": "text-classification", "image-feature-extraction": "feature-extraction", + "pretraining": "feature-extraction", # for backward compatibility and testing (where # model task and model type are still the same) "stable-diffusion": "text-to-image", @@ -314,6 +322,8 @@ class TasksManager: } _CUSTOM_CLASSES = { + ("pt", "patchtsmixer", "time-series-forecasting"): ("transformers", "PatchTSMixerForPrediction"), + ("pt", "patchtst", "time-series-forecasting"): ("transformers", "PatchTSTForPrediction"), ("pt", "pix2struct", "image-to-text"): ("transformers", "Pix2StructForConditionalGeneration"), ("pt", "pix2struct", "visual-question-answering"): ("transformers", "Pix2StructForConditionalGeneration"), ("pt", "visual-bert", "question-answering"): ("transformers", "VisualBertForQuestionAnswering"), @@ -335,7 +345,11 @@ class TasksManager: } _DIFFUSERS_SUPPORTED_MODEL_TYPE = { - "clip-text-model": supported_tasks_mapping( + "t5-encoder": supported_tasks_mapping( + "feature-extraction", + onnx="T5EncoderOnnxConfig", + ), + "clip-text": supported_tasks_mapping( "feature-extraction", onnx="CLIPTextOnnxConfig", ), @@ -343,7 +357,15 @@ class TasksManager: "feature-extraction", onnx="CLIPTextWithProjectionOnnxConfig", ), - "unet": supported_tasks_mapping( + "flux-transformer-2d": supported_tasks_mapping( + "semantic-segmentation", + onnx="FluxTransformerOnnxConfig", + ), + "sd3-transformer-2d": supported_tasks_mapping( + "semantic-segmentation", + onnx="SD3TransformerOnnxConfig", + ), + "unet-2d-condition": supported_tasks_mapping( "semantic-segmentation", onnx="UNetOnnxConfig", ), @@ -418,31 +440,35 @@ class TasksManager: onnx="BertOnnxConfig", tflite="BertTFLiteConfig", ), - # For big-bird and bigbird-pegasus being unsupported, refer to model_configs.py - # "big-bird": supported_tasks_mapping( - # "feature-extraction", - # "fill-mask", - # # the logic for text-generation is not supported for big-bird - # # "text-generation", - # "text-classification", - # "multiple-choice", - # "token-classification", - # "question-answering", - # onnx="BigBirdOnnxConfig", - # # TODO: check model_config.py to know why it cannot be enabled yet. - # # tflite="BigBirdTFLiteConfig", - # ), - # "bigbird-pegasus": supported_tasks_mapping( - # "feature-extraction", - # "feature-extraction-with-past", - # "text-generation", - # "text-generation-with-past", - # "text2text-generation", - # "text2text-generation-with-past", - # "text-classification", - # "question-answering", - # onnx="BigBirdPegasusOnnxConfig", - # ), + "rembert": supported_tasks_mapping( + "fill-mask", + "feature-extraction", + "text-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx="RemBertOnnxConfig", + ), + "big-bird": supported_tasks_mapping( + "feature-extraction", + "fill-mask", + "text-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx="BigBirdOnnxConfig", + ), + "bigbird-pegasus": supported_tasks_mapping( + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", + "text2text-generation", + "text2text-generation-with-past", + "text-classification", + "question-answering", + onnx="BigBirdPegasusOnnxConfig", + ), "blenderbot": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -562,6 +588,11 @@ class TasksManager: onnx="DebertaV2OnnxConfig", tflite="DebertaV2TFLiteConfig", ), + "decision-transformer": supported_tasks_mapping( + "feature-extraction", + "reinforcement-learning", + onnx="DecisionTransformerOnnxConfig", + ), "deit": supported_tasks_mapping( "feature-extraction", "image-classification", @@ -574,6 +605,11 @@ class TasksManager: "image-segmentation", onnx="DetrOnnxConfig", ), + "dinov2": supported_tasks_mapping( + "feature-extraction", + "image-classification", + onnx="Dinov2OnnxConfig", + ), "distilbert": supported_tasks_mapping( "feature-extraction", "fill-mask", @@ -705,6 +741,11 @@ class TasksManager: "feature-extraction", onnx="GroupViTOnnxConfig", ), + "hiera": supported_tasks_mapping( + "feature-extraction", + "image-classification", + onnx="HieraOnnxConfig", + ), "hubert": supported_tasks_mapping( "feature-extraction", "automatic-speech-recognition", @@ -761,15 +802,15 @@ class TasksManager: "text2text-generation-with-past", onnx="LongT5OnnxConfig", ), - # "longformer": supported_tasks_mapping( - # "feature-extraction", - # "fill-mask", - # "multiple-choice", - # "question-answering", - # "text-classification", - # "token-classification", - # onnx_config_cls="models.longformer.LongformerOnnxConfig", - # ), + "longformer": supported_tasks_mapping( + "feature-extraction", + "fill-mask", + "multiple-choice", + "question-answering", + "text-classification", + "token-classification", + onnx="LongformerOnnxConfig", + ), "marian": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -786,6 +827,11 @@ class TasksManager: "question-answering", onnx="MarkupLMOnnxConfig", ), + "maskformer": supported_tasks_mapping( + "feature-extraction", + "image-segmentation", + onnx="MaskFormerOnnxConfig", + ), "mbart": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -797,6 +843,11 @@ class TasksManager: "question-answering", onnx="MBartOnnxConfig", ), + "mgp-str": supported_tasks_mapping( + "feature-extraction", + "image-to-text", + onnx="MgpstrOnnxConfig", + ), "mistral": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -805,12 +856,11 @@ class TasksManager: "text-classification", onnx="MistralOnnxConfig", ), - # TODO: enable once the missing operator is supported. - # "mctct": supported_tasks_mapping( - # "feature-extraction", - # "automatic-speech-recognition", - # onnx="MCTCTOnnxConfig", - # ), + "mctct": supported_tasks_mapping( + "feature-extraction", + "automatic-speech-recognition", + onnx="MCTCTOnnxConfig", + ), "mobilebert": supported_tasks_mapping( "feature-extraction", "fill-mask", @@ -821,6 +871,15 @@ class TasksManager: onnx="MobileBertOnnxConfig", tflite="MobileBertTFLiteConfig", ), + "megatron-bert": supported_tasks_mapping( + "feature-extraction", + "fill-mask", + "text-classification", + "multiple-choice", + "token-classification", + "question-answering", + onnx="MegatronBertOnnxConfig", + ), "mobilevit": supported_tasks_mapping( "feature-extraction", "image-classification", @@ -837,6 +896,13 @@ class TasksManager: "image-classification", onnx="MobileNetV2OnnxConfig", ), + "modernbert": supported_tasks_mapping( + "feature-extraction", + "fill-mask", + "text-classification", + "token-classification", + onnx="ModernBertOnnxConfig", + ), "mpnet": supported_tasks_mapping( "feature-extraction", "fill-mask", @@ -899,6 +965,16 @@ class TasksManager: "text-classification", onnx="OPTOnnxConfig", ), + "patchtst": supported_tasks_mapping( + "feature-extraction", + "time-series-forecasting", + onnx="PatchTSTOnnxConfig", + ), + "patchtsmixer": supported_tasks_mapping( + "feature-extraction", + "time-series-forecasting", + onnx="PatchTSMixerOnnxConfig", + ), "qwen2": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -915,6 +991,27 @@ class TasksManager: "text-classification", onnx="LlamaOnnxConfig", ), + "granite": supported_tasks_mapping( + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", + onnx="GraniteOnnxConfig", + ), + "olmo": supported_tasks_mapping( + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", + onnx="OlmoOnnxConfig", + ), + "olmo2": supported_tasks_mapping( + "feature-extraction", + "feature-extraction-with-past", + "text-generation", + "text-generation-with-past", + onnx="Olmo2OnnxConfig", + ), "pegasus": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -958,6 +1055,11 @@ class TasksManager: "image-classification", onnx="PoolFormerOnnxConfig", ), + "pvt": supported_tasks_mapping( + "feature-extraction", + "image-classification", + onnx="PvtOnnxConfig", + ), "regnet": supported_tasks_mapping( "feature-extraction", "image-classification", @@ -1017,6 +1119,23 @@ class TasksManager: "audio-classification", onnx="SEWDOnnxConfig", ), + "siglip": supported_tasks_mapping( + "feature-extraction", + "zero-shot-image-classification", + onnx="SiglipOnnxConfig", + ), + "siglip-text-model": supported_tasks_mapping( + "feature-extraction", + onnx="SiglipTextOnnxConfig", + ), + "siglip-text-with-projection": supported_tasks_mapping( + "feature-extraction", + onnx="SiglipTextWithProjectionOnnxConfig", + ), + "siglip-vision-model": supported_tasks_mapping( + "feature-extraction", + onnx="SiglipVisionModelOnnxConfig", + ), "speech-to-text": supported_tasks_mapping( "feature-extraction", "feature-extraction-with-past", @@ -1049,6 +1168,12 @@ class TasksManager: "masked-im", onnx="SwinOnnxConfig", ), + "swinv2": supported_tasks_mapping( + "feature-extraction", + "image-classification", + "masked-im", + onnx="SwinV2OnnxConfig", + ), "swin2sr": supported_tasks_mapping( "feature-extraction", "image-to-image", @@ -1095,7 +1220,19 @@ class TasksManager: onnx="VisionEncoderDecoderOnnxConfig", ), "vit": supported_tasks_mapping( - "feature-extraction", "image-classification", "masked-im", onnx="ViTOnnxConfig" + "feature-extraction", + "image-classification", + "masked-im", + onnx="ViTOnnxConfig", + ), + "vit-mae": supported_tasks_mapping( + "feature-extraction", + onnx="VitMAEOnnxConfig", + ), + "vit-msn": supported_tasks_mapping( + "feature-extraction", + "image-classification", + onnx="VitMSNOnnxConfig", ), "vits": supported_tasks_mapping( "text-to-audio", @@ -1170,12 +1307,21 @@ class TasksManager: "transformers": _SUPPORTED_MODEL_TYPE, } _UNSUPPORTED_CLI_MODEL_TYPE = { - "unet", + # diffusers model part + "clip-text", + "clip-text-with-projection", + "flux-transformer-2d", + "sd3-transformer-2d", + "t5-encoder", + "unet-2d-condition", "vae-encoder", "vae-decoder", "clip-text-model", "clip-text-with-projection", - "trocr", # supported through the vision-encoder-decoder model type + "siglip-text-model", + "siglip-text-with-projection", + # transformers model part + "trocr", # the decoder of a trocr vision-encoder-decoder } _SUPPORTED_CLI_MODEL_TYPE = ( set(_SUPPORTED_MODEL_TYPE.keys()) @@ -1637,6 +1783,7 @@ def _infer_task_from_model_name_or_path( revision: Optional[str] = None, cache_dir: str = HUGGINGFACE_HUB_CACHE, token: Optional[Union[bool, str]] = None, + library_name: Optional[str] = None, ) -> str: inferred_task_name = None @@ -1658,13 +1805,14 @@ def _infer_task_from_model_name_or_path( raise RuntimeError( f"Hugging Face Hub is not reachable and we cannot infer the task from a cached model. Make sure you are not offline, or otherwise please specify the `task` (or `--task` in command-line) argument ({', '.join(TasksManager.get_all_tasks())})." ) - library_name = cls.infer_library_from_model( - model_name_or_path, - subfolder=subfolder, - revision=revision, - cache_dir=cache_dir, - token=token, - ) + if library_name is None: + library_name = cls.infer_library_from_model( + model_name_or_path, + subfolder=subfolder, + revision=revision, + cache_dir=cache_dir, + token=token, + ) if library_name == "timm": inferred_task_name = "image-classification" @@ -1683,6 +1831,8 @@ def _infer_task_from_model_name_or_path( break if inferred_task_name is not None: break + elif library_name == "sentence_transformers": + inferred_task_name = "feature-extraction" elif library_name == "transformers": pipeline_tag = model_info.pipeline_tag transformers_info = model_info.transformersInfo @@ -1719,6 +1869,7 @@ def infer_task_from_model( revision: Optional[str] = None, cache_dir: str = HUGGINGFACE_HUB_CACHE, token: Optional[Union[bool, str]] = None, + library_name: Optional[str] = None, ) -> str: """ Infers the task from the model repo, model instance, or model class. @@ -1737,7 +1888,9 @@ def infer_task_from_model( token (`Optional[Union[bool,str]]`, defaults to `None`): The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated when running `huggingface-cli login` (stored in `huggingface_hub.constants.HF_TOKEN_PATH`). - + library_name (`Optional[str]`, defaults to `None`): + The library name of the model. Can be any of "transformers", "timm", "diffusers", "sentence_transformers". See `TasksManager.infer_library_from_model` for the priority should + none be provided. Returns: `str`: The task name automatically detected from the HF hub repo, model instance, or model class. """ @@ -1750,6 +1903,7 @@ def infer_task_from_model( revision=revision, cache_dir=cache_dir, token=token, + library_name=library_name, ) elif type(model) == type: inferred_task_name = cls._infer_task_from_model_or_model_class(model_class=model) @@ -2025,6 +2179,9 @@ def get_model_from_task( none be provided. model_kwargs (`Dict[str, Any]`, *optional*): Keyword arguments to pass to the model `.from_pretrained()` method. + library_name (`Optional[str]`, defaults to `None`): + The library name of the model. Can be any of "transformers", "timm", "diffusers", "sentence_transformers". See `TasksManager.infer_library_from_model` for the priority should + none be provided. Returns: The instance of the model. @@ -2044,7 +2201,12 @@ def get_model_from_task( original_task = task if task == "auto": task = TasksManager.infer_task_from_model( - model_name_or_path, subfolder=subfolder, revision=revision, cache_dir=cache_dir, token=token + model_name_or_path, + subfolder=subfolder, + revision=revision, + cache_dir=cache_dir, + token=token, + library_name=library_name, ) model_type = None @@ -2061,6 +2223,9 @@ def get_model_from_task( if original_task == "automatic-speech-recognition" or task == "automatic-speech-recognition": if original_task == "auto" and config.architectures is not None: model_class_name = config.architectures[0] + elif original_task == "reinforcement-learning" or task == "reinforcement-learning": + if config.architectures is not None: + model_class_name = config.architectures[0] if library_name == "diffusers": config = DiffusionPipeline.load_config(model_name_or_path, **kwargs) @@ -2080,6 +2245,7 @@ def get_model_from_task( use_auth_token = model_kwargs.pop("use_auth_token", None) token = model_kwargs.pop("token", None) trust_remote_code = model_kwargs.pop("trust_remote_code", False) + model_kwargs["torch_dtype"] = torch_dtype if use_auth_token is not None: warnings.warn( @@ -2095,7 +2261,9 @@ def get_model_from_task( device=device, cache_folder=cache_folder, token=token, + revision=revision, trust_remote_code=trust_remote_code, + model_kwargs=model_kwargs, ) else: try: diff --git a/optimum/exporters/tflite/__main__.py b/optimum/exporters/tflite/__main__.py index 0c4c7b994f..4d8d4ee7b7 100644 --- a/optimum/exporters/tflite/__main__.py +++ b/optimum/exporters/tflite/__main__.py @@ -46,7 +46,7 @@ def main(): task = args.task if task == "auto": try: - task = TasksManager.infer_task_from_model(args.model) + task = TasksManager.infer_task_from_model(args.model, library_name="transformers") except KeyError as e: raise KeyError( "The task could not be automatically inferred. Please provide the argument --task with the task " @@ -58,7 +58,12 @@ def main(): ) model = TasksManager.get_model_from_task( - task, args.model, framework="tf", cache_dir=args.cache_dir, trust_remote_code=args.trust_remote_code + task, + args.model, + framework="tf", + cache_dir=args.cache_dir, + trust_remote_code=args.trust_remote_code, + library_name="transformers", ) tflite_config_constructor = TasksManager.get_exporter_config_constructor( diff --git a/optimum/exporters/tflite/convert.py b/optimum/exporters/tflite/convert.py index c1a2010355..fb0706cacd 100644 --- a/optimum/exporters/tflite/convert.py +++ b/optimum/exporters/tflite/convert.py @@ -194,7 +194,7 @@ def prepare_converter_for_quantization( if task is None: from ...exporters import TasksManager - task = TasksManager.infer_task_from_model(model) + task = TasksManager.infer_task_from_model(model, library_name="transformers") preprocessor_kwargs = {} if isinstance(preprocessor, PreTrainedTokenizerBase): diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py index 949b54f468..58e170ba97 100644 --- a/optimum/exporters/utils.py +++ b/optimum/exporters/utils.py @@ -15,9 +15,9 @@ """Utilities for model preparation to export.""" - import copy -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union +from inspect import signature +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Tuple, Union import torch from packaging import version @@ -44,12 +44,7 @@ "Please update diffusers by running `pip install --upgrade diffusers`" ) - from diffusers import ( - DiffusionPipeline, - StableDiffusionXLImg2ImgPipeline, - StableDiffusionXLInpaintPipeline, - StableDiffusionXLPipeline, - ) + from diffusers import DiffusionPipeline from diffusers.models.attention_processor import ( Attention, AttnAddedKVProcessor, @@ -80,6 +75,20 @@ DECODER_MERGED_NAME = "decoder_model_merged" +_DIFFUSERS_CLASS_NAME_TO_SUBMODEL_TYPE = { + "CLIPTextModel": "clip-text", + "CLIPTextModelWithProjection": "clip-text-with-projection", + "FluxTransformer2DModel": "flux-transformer-2d", + "SD3Transformer2DModel": "sd3-transformer-2d", + "UNet2DConditionModel": "unet-2d-condition", + "T5EncoderModel": "t5-encoder", +} + + +def _get_diffusers_submodel_type(submodel): + return _DIFFUSERS_CLASS_NAME_TO_SUBMODEL_TYPE.get(submodel.__class__.__name__) + + def _get_submodels_for_export_diffusion( pipeline: "DiffusionPipeline", ) -> Dict[str, Union["PreTrainedModel", "ModelMixin"]]: @@ -87,56 +96,91 @@ def _get_submodels_for_export_diffusion( Returns the components of a Stable Diffusion model. """ - is_stable_diffusion_xl = isinstance( - pipeline, (StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, StableDiffusionXLInpaintPipeline) - ) - if is_stable_diffusion_xl: - projection_dim = pipeline.text_encoder_2.config.projection_dim - else: - projection_dim = pipeline.text_encoder.config.projection_dim - models_for_export = {} + is_torch_greater_or_equal_than_2_1 = version.parse(torch.__version__) >= version.parse("2.1.0") + is_sdxl = pipeline.__class__.__name__.startswith("StableDiffusionXL") + is_sd3 = pipeline.__class__.__name__.startswith("StableDiffusion3") + # Text encoder text_encoder = getattr(pipeline, "text_encoder", None) if text_encoder is not None: - if is_stable_diffusion_xl: + if is_sdxl or is_sd3: text_encoder.config.output_hidden_states = True + text_encoder.text_model.config.output_hidden_states = True + + text_encoder.config.export_model_type = _get_diffusers_submodel_type(text_encoder) models_for_export["text_encoder"] = text_encoder - # U-NET - # ONNX export of torch.nn.functional.scaled_dot_product_attention not supported for < v2.1.0 - is_torch_greater_or_equal_than_2_1 = version.parse(torch.__version__) >= version.parse("2.1.0") - if not is_torch_greater_or_equal_than_2_1: - pipeline.unet.set_attn_processor(AttnProcessor()) + # Text encoder 2 + text_encoder_2 = getattr(pipeline, "text_encoder_2", None) + if text_encoder_2 is not None: + if is_sdxl or is_sd3: + text_encoder_2.config.output_hidden_states = True + text_encoder_2.text_model.config.output_hidden_states = True - pipeline.unet.config.text_encoder_projection_dim = projection_dim - # The U-NET time_ids inputs shapes depends on the value of `requires_aesthetics_score` - # https://github.com/huggingface/diffusers/blob/v0.18.2/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py#L571 - pipeline.unet.config.requires_aesthetics_score = getattr(pipeline.config, "requires_aesthetics_score", False) - models_for_export["unet"] = pipeline.unet + text_encoder_2.config.export_model_type = _get_diffusers_submodel_type(text_encoder_2) + models_for_export["text_encoder_2"] = text_encoder_2 - # VAE Encoder https://github.com/huggingface/diffusers/blob/v0.11.1/src/diffusers/models/vae.py#L565 + # Text encoder 3 + text_encoder_3 = getattr(pipeline, "text_encoder_3", None) + if text_encoder_3 is not None: + text_encoder_3.config.export_model_type = _get_diffusers_submodel_type(text_encoder_3) + models_for_export["text_encoder_3"] = text_encoder_3 + + # U-NET + unet = getattr(pipeline, "unet", None) + if unet is not None: + # ONNX export of torch.nn.functional.scaled_dot_product_attention not supported for < v2.1.0 + if not is_torch_greater_or_equal_than_2_1: + unet.set_attn_processor(AttnProcessor()) + + # The U-NET time_ids inputs shapes depends on the value of `requires_aesthetics_score` + # https://github.com/huggingface/diffusers/blob/v0.18.2/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl_img2img.py#L571 + unet.config.requires_aesthetics_score = getattr(pipeline.config, "requires_aesthetics_score", False) + unet.config.time_cond_proj_dim = getattr(pipeline.unet.config, "time_cond_proj_dim", None) + unet.config.text_encoder_projection_dim = ( + pipeline.text_encoder.config.projection_dim + if not is_sdxl + else pipeline.text_encoder_2.config.projection_dim + ) + unet.config.export_model_type = _get_diffusers_submodel_type(unet) + models_for_export["unet"] = unet + + # Transformer + transformer = getattr(pipeline, "transformer", None) + if transformer is not None: + # ONNX export of torch.nn.functional.scaled_dot_product_attention not supported for < v2.1.0 + if not is_torch_greater_or_equal_than_2_1: + transformer.set_attn_processor(AttnProcessor()) + + transformer.config.requires_aesthetics_score = getattr(pipeline.config, "requires_aesthetics_score", False) + transformer.config.time_cond_proj_dim = getattr(pipeline.transformer.config, "time_cond_proj_dim", None) + transformer.config.text_encoder_projection_dim = pipeline.text_encoder.config.projection_dim + transformer.config.export_model_type = _get_diffusers_submodel_type(transformer) + models_for_export["transformer"] = transformer + + # VAE Encoder vae_encoder = copy.deepcopy(pipeline.vae) + + # ONNX export of torch.nn.functional.scaled_dot_product_attention not supported for < v2.1.0 if not is_torch_greater_or_equal_than_2_1: vae_encoder = override_diffusers_2_0_attn_processors(vae_encoder) + # we return the distribution parameters to be able to recreate it in the decoder vae_encoder.forward = lambda sample: {"latent_parameters": vae_encoder.encode(x=sample)["latent_dist"].parameters} models_for_export["vae_encoder"] = vae_encoder - # VAE Decoder https://github.com/huggingface/diffusers/blob/v0.11.1/src/diffusers/models/vae.py#L600 + # VAE Decoder vae_decoder = copy.deepcopy(pipeline.vae) + + # ONNX export of torch.nn.functional.scaled_dot_product_attention not supported for < v2.1.0 if not is_torch_greater_or_equal_than_2_1: vae_decoder = override_diffusers_2_0_attn_processors(vae_decoder) + vae_decoder.forward = lambda latent_sample: vae_decoder.decode(z=latent_sample) models_for_export["vae_decoder"] = vae_decoder - text_encoder_2 = getattr(pipeline, "text_encoder_2", None) - if text_encoder_2 is not None: - text_encoder_2.config.output_hidden_states = True - text_encoder_2.text_model.config.output_hidden_states = True - models_for_export["text_encoder_2"] = text_encoder_2 - return models_for_export @@ -294,33 +338,59 @@ def get_diffusion_models_for_export( `Dict[str, Tuple[Union[`PreTrainedModel`, `TFPreTrainedModel`], `ExportConfig`]: A Dict containing the model and export configs for the different components of the model. """ + models_for_export = _get_submodels_for_export_diffusion(pipeline) # Text encoder if "text_encoder" in models_for_export: + text_encoder = models_for_export["text_encoder"] text_encoder_config_constructor = TasksManager.get_exporter_config_constructor( - model=pipeline.text_encoder, - exporter=exporter, - library_name="diffusers", - task="feature-extraction", + model=text_encoder, exporter=exporter, library_name="diffusers", task="feature-extraction" ) text_encoder_export_config = text_encoder_config_constructor( - pipeline.text_encoder.config, int_dtype=int_dtype, float_dtype=float_dtype + text_encoder.config, int_dtype=int_dtype, float_dtype=float_dtype ) models_for_export["text_encoder"] = (models_for_export["text_encoder"], text_encoder_export_config) + # Text encoder 2 + if "text_encoder_2" in models_for_export: + text_encoder_2 = models_for_export["text_encoder_2"] + export_config_constructor = TasksManager.get_exporter_config_constructor( + model=text_encoder_2, exporter=exporter, library_name="diffusers", task="feature-extraction" + ) + export_config = export_config_constructor(text_encoder_2.config, int_dtype=int_dtype, float_dtype=float_dtype) + models_for_export["text_encoder_2"] = (models_for_export["text_encoder_2"], export_config) + + # Text encoder 3 + if "text_encoder_3" in models_for_export: + text_encoder_3 = models_for_export["text_encoder_3"] + export_config_constructor = TasksManager.get_exporter_config_constructor( + model=text_encoder_3, exporter=exporter, library_name="diffusers", task="feature-extraction" + ) + export_config = export_config_constructor(text_encoder_3.config, int_dtype=int_dtype, float_dtype=float_dtype) + models_for_export["text_encoder_3"] = (models_for_export["text_encoder_3"], export_config) + # U-NET - export_config_constructor = TasksManager.get_exporter_config_constructor( - model=pipeline.unet, - exporter=exporter, - library_name="diffusers", - task="semantic-segmentation", - model_type="unet", - ) - unet_export_config = export_config_constructor(pipeline.unet.config, int_dtype=int_dtype, float_dtype=float_dtype) - models_for_export["unet"] = (models_for_export["unet"], unet_export_config) + if "unet" in models_for_export: + unet = models_for_export["unet"] + export_config_constructor = TasksManager.get_exporter_config_constructor( + model=unet, exporter=exporter, library_name="diffusers", task="semantic-segmentation" + ) + unet_export_config = export_config_constructor(unet.config, int_dtype=int_dtype, float_dtype=float_dtype) + models_for_export["unet"] = (models_for_export["unet"], unet_export_config) + + # Transformer + if "transformer" in models_for_export: + transformer = models_for_export["transformer"] + export_config_constructor = TasksManager.get_exporter_config_constructor( + model=transformer, exporter=exporter, library_name="diffusers", task="semantic-segmentation" + ) + transformer_export_config = export_config_constructor( + transformer.config, int_dtype=int_dtype, float_dtype=float_dtype + ) + models_for_export["transformer"] = (models_for_export["transformer"], transformer_export_config) - # VAE Encoder https://github.com/huggingface/diffusers/blob/v0.11.1/src/diffusers/models/vae.py#L565 + # VAE Encoder vae_encoder = models_for_export["vae_encoder"] vae_config_constructor = TasksManager.get_exporter_config_constructor( model=vae_encoder, @@ -329,10 +399,12 @@ def get_diffusion_models_for_export( task="semantic-segmentation", model_type="vae-encoder", ) - vae_export_config = vae_config_constructor(vae_encoder.config, int_dtype=int_dtype, float_dtype=float_dtype) - models_for_export["vae_encoder"] = (vae_encoder, vae_export_config) + vae_encoder_export_config = vae_config_constructor( + vae_encoder.config, int_dtype=int_dtype, float_dtype=float_dtype + ) + models_for_export["vae_encoder"] = (vae_encoder, vae_encoder_export_config) - # VAE Decoder https://github.com/huggingface/diffusers/blob/v0.11.1/src/diffusers/models/vae.py#L600 + # VAE Decoder vae_decoder = models_for_export["vae_decoder"] vae_config_constructor = TasksManager.get_exporter_config_constructor( model=vae_decoder, @@ -341,21 +413,10 @@ def get_diffusion_models_for_export( task="semantic-segmentation", model_type="vae-decoder", ) - vae_export_config = vae_config_constructor(vae_decoder.config, int_dtype=int_dtype, float_dtype=float_dtype) - models_for_export["vae_decoder"] = (vae_decoder, vae_export_config) - - if "text_encoder_2" in models_for_export: - export_config_constructor = TasksManager.get_exporter_config_constructor( - model=pipeline.text_encoder_2, - exporter=exporter, - library_name="diffusers", - task="feature-extraction", - model_type="clip-text-with-projection", - ) - export_config = export_config_constructor( - pipeline.text_encoder_2.config, int_dtype=int_dtype, float_dtype=float_dtype - ) - models_for_export["text_encoder_2"] = (models_for_export["text_encoder_2"], export_config) + vae_decoder_export_config = vae_config_constructor( + vae_decoder.config, int_dtype=int_dtype, float_dtype=float_dtype + ) + models_for_export["vae_decoder"] = (vae_decoder, vae_decoder_export_config) return models_for_export @@ -619,3 +680,39 @@ def _get_submodels_and_export_configs( export_config = next(iter(models_and_export_configs.values()))[1] return export_config, models_and_export_configs + + +def check_dummy_inputs_are_allowed( + model: Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], dummy_input_names: Iterable[str] +): + """ + Checks that the dummy inputs from the ONNX config is a subset of the allowed inputs for `model`. + Args: + model (`Union[transformers.PreTrainedModel, transformers.TFPreTrainedModel`]): + The model instance. + model_inputs (`Iterable[str]`): + The model input names. + """ + + forward = model.forward if is_torch_available() and isinstance(model, torch.nn.Module) else model.call + forward_parameters = signature(forward).parameters + forward_inputs_set = set(forward_parameters.keys()) + dummy_input_names = set(dummy_input_names) + + # We are fine if config_inputs has more keys than model_inputs + if not dummy_input_names.issubset(forward_inputs_set): + raise ValueError( + f"Config dummy inputs are not a subset of the model inputs: {dummy_input_names} vs {forward_inputs_set}" + ) + + +class DisableCompileContextManager: + def __init__(self): + self._original_compile = torch.compile + + def __enter__(self): + # Turn torch.compile into a no-op + torch.compile = lambda *args, **kwargs: lambda x: x + + def __exit__(self, exc_type, exc_val, exc_tb): + torch.compile = self._original_compile diff --git a/optimum/gptq/constants.py b/optimum/gptq/constants.py index 2d3e51da7a..701868a3b8 100644 --- a/optimum/gptq/constants.py +++ b/optimum/gptq/constants.py @@ -18,6 +18,10 @@ "model.decoder.layers", "gpt_neox.layers", "model.layers", + # modules loaded by AutoModel vs AutoModelForCausalLM have different prefixes + "h", + "decoder.layers", + "layers", ] GPTQ_CONFIG = "quantize_config.json" diff --git a/optimum/gptq/data.py b/optimum/gptq/data.py index b8734da478..7e5fc0b43d 100644 --- a/optimum/gptq/data.py +++ b/optimum/gptq/data.py @@ -18,7 +18,12 @@ import numpy as np import torch -from datasets import load_dataset + +from optimum.utils.import_utils import DATASETS_IMPORT_ERROR, is_datasets_available + + +if is_datasets_available(): + from datasets import load_dataset """ @@ -113,6 +118,9 @@ def pad_block(block, pads): def get_wikitext2(tokenizer: Any, seqlen: int, nsamples: int, split: str = "train"): + if not is_datasets_available(): + raise ImportError(DATASETS_IMPORT_ERROR.format("get_wikitext2")) + if split == "train": data = load_dataset("wikitext", "wikitext-2-raw-v1", split="train") elif split == "validation": @@ -132,6 +140,9 @@ def get_wikitext2(tokenizer: Any, seqlen: int, nsamples: int, split: str = "trai def get_c4(tokenizer: Any, seqlen: int, nsamples: int, split: str = "train"): + if not is_datasets_available(): + raise ImportError(DATASETS_IMPORT_ERROR.format("get_c4")) + if split == "train": data = load_dataset("allenai/c4", split="train", data_files={"train": "en/c4-train.00000-of-01024.json.gz"}) elif split == "validation": @@ -157,6 +168,9 @@ def get_c4(tokenizer: Any, seqlen: int, nsamples: int, split: str = "train"): def get_c4_new(tokenizer: Any, seqlen: int, nsamples: int, split: str = "train"): + if not is_datasets_available(): + raise ImportError(DATASETS_IMPORT_ERROR.format("get_c4_new")) + if split == "train": data = load_dataset("allenai/c4", split="train", data_files={"train": "en/c4-train.00000-of-01024.json.gz"}) elif split == "validation": diff --git a/optimum/gptq/quantizer.py b/optimum/gptq/quantizer.py index 949d4d260d..753dbc64de 100644 --- a/optimum/gptq/quantizer.py +++ b/optimum/gptq/quantizer.py @@ -12,6 +12,7 @@ # 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 importlib import json import os from enum import Enum @@ -19,17 +20,26 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch +from packaging import version from torch import nn from tqdm.auto import tqdm from transformers import AutoTokenizer from transformers.pytorch_utils import Conv1D from transformers.utils.quantization_config import QuantizationMethod -from ..utils import is_accelerate_available, is_auto_gptq_available +from ..utils import is_accelerate_available, is_auto_gptq_available, is_gptqmodel_available from ..utils.modeling_utils import recurse_getattr +from ..version import __version__ as optimum_version from .constants import GPTQ_CONFIG from .data import get_dataset, prepare_dataset -from .utils import get_block_name_with_pattern, get_device, get_layers, get_preceding_modules, get_seqlen +from .utils import ( + get_block_name_with_pattern, + get_device, + get_layers, + get_preceding_modules, + get_seqlen, + nested_move_to, +) if is_accelerate_available(): @@ -40,14 +50,27 @@ from accelerate.hooks import remove_hook_from_module if is_auto_gptq_available(): + from auto_gptq import __version__ as autogptq_version from auto_gptq import exllama_set_max_input_length - from auto_gptq.modeling._utils import autogptq_post_init + from auto_gptq.modeling._utils import autogptq_post_init as gptq_post_init from auto_gptq.quantization import GPTQ - from auto_gptq.utils.import_utils import dynamically_import_QuantLinear + from auto_gptq.utils.import_utils import dynamically_import_QuantLinear as hf_select_quant_linear + +if is_gptqmodel_available(): + from gptqmodel import exllama_set_max_input_length + from gptqmodel.quantization import GPTQ + from gptqmodel.utils.importer import hf_select_quant_linear + from gptqmodel.utils.model import hf_convert_gptq_v1_to_v2_format, hf_convert_gptq_v2_to_v1_format + from gptqmodel.utils.model import hf_gptqmodel_post_init as gptq_post_init + from gptqmodel.version import __version__ as gptqmodel_version logger = getLogger(__name__) +def has_device_more_than_cpu(): + return torch.cuda.is_available() or (hasattr(torch, "xpu") and torch.xpu.is_available()) + + class ExllamaVersion(int, Enum): ONE = 1 TWO = 2 @@ -74,10 +97,13 @@ def __init__( batch_size: int = 1, pad_token_id: Optional[int] = None, disable_exllama: bool = False, - exllama_config: Dict[str, Any] = None, + exllama_config: Optional[Dict[str, Any]] = None, max_input_length: Optional[int] = None, cache_block_outputs: Optional[bool] = True, modules_in_block_to_quantize: Optional[List[List[str]]] = None, + checkpoint_format: str = "gptq", + meta: Optional[Dict[str, any]] = None, + backend: Optional[str] = None, *args, **kwargs, ): @@ -88,7 +114,7 @@ def __init__( dataset (`Union[List[str], str, Any]`, defaults to `None`): The dataset used for quantization. You can provide your own dataset in a list of string or in a list of tokenized data (e.g. [{ "input_ids": [ 1, 100, 15, ... ],"attention_mask": [ 1, 1, 1, ... ]},...]) - or just use the original datasets used in GPTQ paper ['wikitext2','c4','c4-new','ptb','ptb-new']. + or just use the original datasets used in GPTQ paper ['wikitext2','c4','c4-new']. group_size (int, defaults to 128): The group size to use for quantization. Recommended value is 128 and -1 uses per-column quantization. damp_percent (`float`, defaults to `0.1`): @@ -129,6 +155,13 @@ def __init__( List list of module names to quantize in the block specified. This argument is useful to exclude certain linear modules from being quantized. The block to quantize can be specified by setting `block_name_to_quantize`. We will quantize each list sequentially. If not set, we will quantize all linear layers. Example: `inside_layer_modules=[["self_attention.query_key_value"], ["mlp.dense_h_to_4h"]]` + checkpoint_format (`str`, *optional*, defaults to `gptq`): + GPTQ weight format. `gptq`(v1) is supported by both gptqmodel and auto-gptq. `gptq_v2` is gptqmodel only. + meta (`Dict[str, any]`, *optional*): + Properties, such as tooling:version, that do not directly contributes to quantization or quant inference are stored in meta. + i.e. `meta.quantizer`: ["optimum:_version_", "gptqmodel:_version_"] + backend (`str`, *optional*): + Controls which gptq kernel to be used. Valid values for gptqmodel are `auto`, `auto_trainable` and more. For auto-gptq, only valid value is None and `auto_trainable`. Ref gptqmodel backends: https://github.com/ModelCloud/GPTQModel/blob/main/gptqmodel/utils/backend.py """ self.bits = bits @@ -138,6 +171,9 @@ def __init__( self.desc_act = desc_act self.sym = sym self.true_sequential = true_sequential + self.checkpoint_format = checkpoint_format.lower() + self.meta = meta + self.backend = backend.lower() if backend is not None else None self.use_cuda_fp16 = use_cuda_fp16 self.model_seqlen = model_seqlen self.block_name_to_quantize = block_name_to_quantize @@ -161,6 +197,8 @@ def __init__( "true_sequential", "quant_method", "modules_in_block_to_quantize", + "checkpoint_format", + "meta", ] if self.bits not in [2, 3, 4, 8]: @@ -182,6 +220,29 @@ def __init__( ) self.exllama_version = self.exllama_config["version"] + def select_quant_linear(self, device_map: Union[str, dict], pack: bool = False): + if is_gptqmodel_available(): + self.quant_linear = hf_select_quant_linear( + bits=self.bits, + group_size=self.group_size, + desc_act=self.desc_act, + sym=self.sym, + checkpoint_format=self.checkpoint_format, + meta=self.meta, + device_map=device_map, + backend=self.backend, + pack=pack, + ) + else: + self.quant_linear = hf_select_quant_linear( + use_triton=False, + desc_act=self.desc_act, + group_size=self.group_size, + bits=self.bits, + disable_exllama=self.disable_exllama or self.exllama_version != ExllamaVersion.ONE, + disable_exllamav2=self.disable_exllama or self.exllama_version != ExllamaVersion.TWO, + ) + def to_dict(self): """ Returns the args in dict format. @@ -189,6 +250,20 @@ def to_dict(self): gptq_dict = {} for key in self.serialization_keys: gptq_dict[key] = getattr(self, key) + + if gptq_dict.get("meta") is None: + gptq_dict["meta"] = {} + + meta = gptq_dict["meta"] + # store both optimum:version and gptq_lib:version into quantize_config.meta.quantizer + if meta.get("quantizer") is None: + meta["quantizer"] = [f"optimum:{optimum_version}"] + + if is_gptqmodel_available(): + meta["quantizer"].append(f"gptqmodel:{gptqmodel_version}") + elif is_auto_gptq_available(): + meta["quantizer"].append(f"auto_gptq:{autogptq_version}") + return gptq_dict @classmethod @@ -205,7 +280,7 @@ def from_dict(cls, config_dict: Dict[str, Any]): """ return cls(**config_dict) - def convert_model(self, model: nn.Module): + def convert_model(self, model: nn.Module, **kwargs): """ Convert the model to a GPTQ model by getting and replacing the layers. @@ -226,7 +301,11 @@ def convert_model(self, model: nn.Module): f"Quantization disabled for {name} (only modules_in_block_to_quantize={self.modules_in_block_to_quantize} are quantized)" ) del layers_to_be_replaced[name] + + self.select_quant_linear(device_map=kwargs.get("device_map", None), pack=False) + self._replace_by_quant_layers(model, layers_to_be_replaced) + return model def get_no_split_module_classes(self, model): @@ -253,15 +332,7 @@ def _replace_by_quant_layers(self, module: nn.Module, names: List[str], name: st name (`str`, defaults to `""`): To keep track of the name of the current module """ - QuantLinear = dynamically_import_QuantLinear( - use_triton=False, - desc_act=self.desc_act, - group_size=self.group_size, - bits=self.bits, - disable_exllama=self.disable_exllama or self.exllama_version != ExllamaVersion.ONE, - disable_exllamav2=self.disable_exllama or self.exllama_version != ExllamaVersion.TWO, - ) - if isinstance(module, QuantLinear): + if isinstance(module, self.quant_linear): return for attr in dir(module): layer = getattr(module, attr) @@ -279,20 +350,37 @@ def _replace_by_quant_layers(self, module: nn.Module, names: List[str], name: st in_features = layer.weight.shape[0] out_features = layer.weight.shape[1] bias = layer.bias is not None - if not (self.desc_act) or self.group_size == -1: - new_layer = QuantLinear( + if is_gptqmodel_available(): + new_layer = self.quant_linear( self.bits, self.group_size, + self.desc_act, + self.sym, in_features, out_features, bias, - use_cuda_fp16=self.use_cuda_fp16, weight_dtype=layer.weight.dtype, ) else: - new_layer = QuantLinear( - self.bits, self.group_size, in_features, out_features, bias, weight_dtype=layer.weight.dtype - ) + if not (self.desc_act) or self.group_size == -1: + new_layer = self.quant_linear( + self.bits, + self.group_size, + in_features, + out_features, + bias, + use_cuda_fp16=self.use_cuda_fp16, + weight_dtype=layer.weight.dtype, + ) + else: + new_layer = self.quant_linear( + self.bits, + self.group_size, + in_features, + out_features, + bias, + weight_dtype=layer.weight.dtype, + ) new_layer.device = device setattr(module, attr, new_layer.to(device)) for name1, child in module.named_children(): @@ -318,13 +406,41 @@ def quantize_model(self, model: nn.Module, tokenizer: Optional[Any] = None): `nn.Module`: The quantized model """ - if not is_auto_gptq_available(): - raise RuntimeError("auto-gptq is required in order to perform quantzation : `pip install auto-gptq`") - if not torch.cuda.is_available(): - raise RuntimeError("No GPU found. A GPU is needed to quantize model.") + if not is_auto_gptq_available() and not is_gptqmodel_available(): + raise RuntimeError( + "gptqmodel or auto-gptq is required in order to perform gptq quantzation: `pip install gptqmodel` or `pip install auto-gptq`. Please notice that auto-gptq will be deprecated in the future." + ) + elif is_gptqmodel_available() and is_auto_gptq_available(): + logger.warning( + "Detected gptqmodel and auto-gptq, will use gptqmodel. The auto_gptq will be deprecated in the future." + ) + + gptq_supports_cpu = ( + is_auto_gptq_available() + and version.parse(importlib.metadata.version("auto-gptq")) > version.parse("0.4.2") + ) or is_gptqmodel_available() + + if not gptq_supports_cpu and not torch.cuda.is_available(): + raise RuntimeError( + "No cuda gpu or cpu support using Intel/IPEX found. A gpu or cpu with Intel/IPEX is required for quantization." + ) + + if not self.sym and not is_gptqmodel_available(): + raise ValueError( + "Asymmetric sym=False quantization is not supported with auto-gptq. Please use gptqmodel: `pip install gptqmodel`" + ) + + if self.checkpoint_format == "gptq_v2" and not is_gptqmodel_available(): + raise ValueError( + "gptq_v2 format only supported with gptqmodel. Please install gptqmodel: `pip install gptqmodel`" + ) model.eval() + # gptqmodel internal is gptq_v2 for asym support, gptq(v1) can only support sym=True + if is_gptqmodel_available() and self.checkpoint_format != "gptq_v2": + self.checkpoint_format = "gptq_v2" + # For Transformer model has_config = False has_device_map = False @@ -403,27 +519,32 @@ def quantize_model(self, model: nn.Module, tokenizer: Optional[Any] = None): blocks = recurse_getattr(model, self.block_name_to_quantize) + cur_layer_device = get_device(blocks[0]) + if not is_gptqmodel_available() and cur_layer_device.type == "cpu": + cur_layer_device = 0 + if not has_device_map: - # put modules from module_name_preceding_first_block on cuda + # put modules from module_name_preceding_first_block on cuda or xpu or cpu + to_device = cur_layer_device for module_name in self.module_name_preceding_first_block: module = recurse_getattr(model, module_name) if module is None: raise ValueError(f"Module {module_name} was not found in model") - module = module.to(0) - blocks[0] = blocks[0].to(0) + module = module.to(to_device) + blocks[0] = blocks[0].to(to_device) def store_input_hook(_, input, *args): kwargs = args[0] if input is None: if "hidden_states" in kwargs: - input = (kwargs["hidden_states"],) + input = (nested_move_to(kwargs["hidden_states"], cur_layer_device),) else: raise ValueError("No input value found in the foward pass") layer_inputs.append(input) other_kwargs = {} for k, v in kwargs.items(): # make sure other arguments also be captured if k not in ["hidden_states"]: - other_kwargs[k] = v + other_kwargs[k] = nested_move_to(v, cur_layer_device) layer_input_kwargs.append(other_kwargs) raise ValueError @@ -431,11 +552,7 @@ def store_input_hook(_, input, *args): handle = blocks[0].register_forward_pre_hook(store_input_hook, with_kwargs=True) for data in dataset: for k, v in data.items(): - # put the data on gpu, we won't put them back to cpu - if not has_device_map or device.type == "cpu": - data[k] = v.to(0) - else: - data[k] = v.to(device) + data[k] = nested_move_to(v, cur_layer_device) try: model(**data) except ValueError: @@ -450,6 +567,8 @@ def store_input_hook(_, input, *args): raise ValueError(f"Module {module_name} was not found in model") torch.cuda.empty_cache() + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() # Step 3: Quantize the blocks quantizers = {} @@ -460,11 +579,7 @@ def store_input_hook(_, input, *args): handle = block.register_forward_pre_hook(store_input_hook, with_kwargs=True) for data in dataset: for k, v in data.items(): - # put the data on gpu, we won't put them back to cpu - if not has_device_map or device.type == "cpu": - data[k] = v.to(0) - else: - data[k] = v.to(device) + data[k] = nested_move_to(v, cur_layer_device) try: model(**data) except ValueError: @@ -473,9 +588,12 @@ def store_input_hook(_, input, *args): # move block to cuda if needed # in case we have offload modules, we need to put them on cuda because of GPTQ object - if not has_device_map or get_device(block) == torch.device("cpu"): + if (not has_device_map or get_device(block) == torch.device("cpu")) and has_device_more_than_cpu(): block = block.to(0) layers = get_layers(block) + block_device = get_device(block) + if not is_gptqmodel_available() and block_device.type == "cpu": + block_device = 0 if isinstance(self.modules_in_block_to_quantize, list) and len(self.modules_in_block_to_quantize) > 0: if self.true_sequential: layers_name_list = self.modules_in_block_to_quantize @@ -509,15 +627,20 @@ def tmp(_, input, output): for j in range(len(dataset)): # the args are already on the gpu # don't need to store the output + layer_inputs[j] = nested_move_to(layer_inputs[j], block_device) + for k, v in layer_input_kwargs[j].items(): + layer_input_kwargs[j][k] = nested_move_to(v, block_device) + block(*layer_inputs[j], **layer_input_kwargs[j]) # remove hook for h in handles: h.remove() for name in subset_name_list: logger.info(f"Quantizing {name} in block {i + 1}/{len(blocks)}...") - scale, zero, g_idx = gptq[name].fasterquant( + quant_outputs = gptq[name].fasterquant( percdamp=self.damp_percent, group_size=self.group_size, actorder=self.desc_act ) + scale, zero, g_idx = quant_outputs[0], quant_outputs[1], quant_outputs[2] quantizers[f"{self.block_name_to_quantize}.{i}.{name}"] = ( gptq[name].quantizer, scale, @@ -543,11 +666,13 @@ def tmp(_, input, output): del layer_inputs layer_inputs = [] torch.cuda.empty_cache() + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() if self.bits == 4: # device not on gpu if device.type != "cuda" or (has_device_map and any(d in devices for d in ["cpu", "disk", "hpu"])): - if not self.disable_exllama: + if not self.disable_exllama and not is_gptqmodel_available(): logger.warning( "Found modules on cpu/disk. Using Exllama/Exllamav2 backend requires all the modules to be on GPU. Setting `disable_exllama=True`" ) @@ -578,6 +703,8 @@ def tmp(_, input, output): model = self.post_init_model(model) torch.cuda.empty_cache() + if hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.empty_cache() return model def post_init_model(self, model): @@ -601,9 +728,14 @@ def post_init_model(self, model): class StoreAttr(object): pass + if is_gptqmodel_available(): + model, _ = hf_convert_gptq_v1_to_v2_format( + model, self.bits, self.quant_linear, self.checkpoint_format, self.meta + ) + model.quantize_config = StoreAttr() model.quantize_config.desc_act = self.desc_act - model = autogptq_post_init(model, use_act_order=self.desc_act) + model = gptq_post_init(model, use_act_order=self.desc_act) if ( self.desc_act and (not self.disable_exllama and self.exllama_version == ExllamaVersion.ONE) @@ -626,19 +758,14 @@ def pack_model( quantizers (`Dict[str,Tuple]`): A mapping of the layer name and the data needed to pack the layer """ - QuantLinear = dynamically_import_QuantLinear( - use_triton=False, - desc_act=self.desc_act, - group_size=self.group_size, - bits=self.bits, - disable_exllama=self.disable_exllama or self.exllama_version != ExllamaVersion.ONE, - disable_exllamav2=self.disable_exllama or self.exllama_version != ExllamaVersion.TWO, - ) logger.info("Packing model...") layers = get_layers(model) layers = {n: layers[n] for n in quantizers} + + self.select_quant_linear(device_map=model.hf_device_map, pack=True) + self._replace_by_quant_layers(model, quantizers) - qlayers = get_layers(model, [QuantLinear]) + qlayers = get_layers(model, [self.quant_linear]) for name in qlayers: logger.info(name) quantizers[name], scale, zero, g_idx = quantizers[name] @@ -673,6 +800,15 @@ def save(self, model: nn.Module, save_dir: str, max_shard_size: str = "10GB", sa Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`). """ + + # convert gptqmodel internal gptq_v2 format to v1 for max compatibility + if is_gptqmodel_available(): + model, converted = hf_convert_gptq_v2_to_v1_format( + model, self.sym, self.bits, self.quant_linear, self.checkpoint_format, self.meta + ) + if converted: + self.checkpoint_format = "gptq" + os.makedirs(save_dir, exist_ok=True) model.save_pretrained(save_dir, max_shard_size=max_shard_size, safe_serialization=safe_serialization) with open(os.path.join(save_dir, GPTQ_CONFIG), "w", encoding="utf-8") as f: @@ -736,10 +872,12 @@ def load_quantized_model( Returns: `nn.Module`: The quantized model """ - if not torch.cuda.is_available(): - raise RuntimeError("No GPU found. A GPU is needed to run quantized model.") - if not is_auto_gptq_available(): - raise RuntimeError("auto-gptq is required in order to load quantized weights : `pip install auto-gptq`") + if not torch.cuda.is_available() and not is_gptqmodel_available(): + raise RuntimeError("No GPU found. A GPU is needed to run quantized model by auto_gptq.") + if not is_auto_gptq_available() and not is_gptqmodel_available(): + raise RuntimeError( + "gptqmodel (`pip install gptqmodel`) or auto-gptq (`pip install auto-gptq`) is required in order to load quantized weights. Please notice that auto-gptq will be deprecated in the future." + ) if not is_accelerate_available(): raise RuntimeError( "You need to install accelerate in order to load and dispatch weights to" @@ -777,7 +915,7 @@ def load_quantized_model( quantizer.exllama_version = quantizer.exllama_config["version"] quantizer.max_input_length = max_input_length - model = quantizer.convert_model(model) + model = quantizer.convert_model(model, device_map=device_map) if no_split_module_classes is None: no_split_module_classes = quantizer.get_no_split_module_classes(model) diff --git a/optimum/gptq/utils.py b/optimum/gptq/utils.py index a5f9afdaae..c32f364d2f 100644 --- a/optimum/gptq/utils.py +++ b/optimum/gptq/utils.py @@ -72,7 +72,7 @@ def get_block_name_with_pattern(model: nn.Module): modules_names = [n for n, _ in model.named_modules()] for pattern_candidate in BLOCK_PATTERNS: pattern_candidate = pattern_candidate - if any(pattern_candidate in name for name in modules_names): + if any(name.startswith(pattern_candidate) for name in modules_names): return pattern_candidate raise ValueError("Block pattern could not be match. Pass `block_name_to_quantize` argument in `quantize_model`") @@ -113,3 +113,18 @@ def get_seqlen(model: nn.Module): "We couldn't get the model sequence length. Setting it to 2048. You can overwrite this value by passing `model_seqlen` in` GPTQQuantizer`" ) return 2048 + + +def move_to(obj: torch.Tensor, device: torch.device): + if get_device(obj) != device: + obj = obj.to(device) + return obj + + +def nested_move_to(v, device): + if isinstance(v, torch.Tensor): + return move_to(v, device) + elif isinstance(v, (list, tuple)): + return type(v)([nested_move_to(e, device) for e in v]) + else: + return v diff --git a/optimum/onnxruntime/__init__.py b/optimum/onnxruntime/__init__.py index 4e25a43690..f3f1535fd4 100644 --- a/optimum/onnxruntime/__init__.py +++ b/optimum/onnxruntime/__init__.py @@ -74,33 +74,51 @@ raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: _import_structure[".utils.dummy_diffusers_objects"] = [ - "ORTStableDiffusionPipeline", + "ORTDiffusionPipeline", + "ORTPipelineForText2Image", + "ORTPipelineForImage2Image", + "ORTPipelineForInpainting", + # flux + "ORTFluxPipeline", + # lcm + "ORTLatentConsistencyModelImg2ImgPipeline", + "ORTLatentConsistencyModelPipeline", + # sd3 + "ORTStableDiffusion3Img2ImgPipeline", + "ORTStableDiffusion3InpaintPipeline", + "ORTStableDiffusion3Pipeline", + # sd "ORTStableDiffusionImg2ImgPipeline", "ORTStableDiffusionInpaintPipeline", - "ORTStableDiffusionXLPipeline", + "ORTStableDiffusionPipeline", + # xl "ORTStableDiffusionXLImg2ImgPipeline", "ORTStableDiffusionXLInpaintPipeline", - "ORTLatentConsistencyModelPipeline", - "ORTLatentConsistencyModelImg2ImgPipeline", - "ORTPipelineForImage2Image", - "ORTPipelineForInpainting", - "ORTPipelineForText2Image", - "ORTDiffusionPipeline", + "ORTStableDiffusionXLPipeline", ] else: _import_structure["modeling_diffusion"] = [ - "ORTStableDiffusionPipeline", + "ORTDiffusionPipeline", + "ORTPipelineForText2Image", + "ORTPipelineForImage2Image", + "ORTPipelineForInpainting", + # flux + "ORTFluxPipeline", + # lcm + "ORTLatentConsistencyModelImg2ImgPipeline", + "ORTLatentConsistencyModelPipeline", + # sd3 + "ORTStableDiffusion3Img2ImgPipeline", + "ORTStableDiffusion3InpaintPipeline", + "ORTStableDiffusion3Pipeline", + # sd "ORTStableDiffusionImg2ImgPipeline", "ORTStableDiffusionInpaintPipeline", - "ORTStableDiffusionXLPipeline", + "ORTStableDiffusionPipeline", + # xl "ORTStableDiffusionXLImg2ImgPipeline", "ORTStableDiffusionXLInpaintPipeline", - "ORTLatentConsistencyModelImg2ImgPipeline", - "ORTLatentConsistencyModelPipeline", - "ORTPipelineForImage2Image", - "ORTPipelineForInpainting", - "ORTPipelineForText2Image", - "ORTDiffusionPipeline", + "ORTStableDiffusionXLPipeline", ] @@ -151,30 +169,52 @@ raise OptionalDependencyNotAvailable() except OptionalDependencyNotAvailable: from ..utils.dummy_diffusers_objects import ( + # generic entrypoint ORTDiffusionPipeline, + # flux + ORTFluxPipeline, + # lcm ORTLatentConsistencyModelImg2ImgPipeline, ORTLatentConsistencyModelPipeline, + # task-specific entrypoints ORTPipelineForImage2Image, ORTPipelineForInpainting, ORTPipelineForText2Image, + # sd3 + ORTStableDiffusion3Img2ImgPipeline, + ORTStableDiffusion3InpaintPipeline, + ORTStableDiffusion3Pipeline, + # sd ORTStableDiffusionImg2ImgPipeline, ORTStableDiffusionInpaintPipeline, ORTStableDiffusionPipeline, + # xl ORTStableDiffusionXLImg2ImgPipeline, ORTStableDiffusionXLInpaintPipeline, ORTStableDiffusionXLPipeline, ) else: from .modeling_diffusion import ( + # generic entrypoint ORTDiffusionPipeline, + # flux + ORTFluxPipeline, + # lcm ORTLatentConsistencyModelImg2ImgPipeline, ORTLatentConsistencyModelPipeline, + # task-specific entrypoints ORTPipelineForImage2Image, ORTPipelineForInpainting, ORTPipelineForText2Image, + # sd3 + ORTStableDiffusion3Img2ImgPipeline, + ORTStableDiffusion3InpaintPipeline, + ORTStableDiffusion3Pipeline, + # sd ORTStableDiffusionImg2ImgPipeline, ORTStableDiffusionInpaintPipeline, ORTStableDiffusionPipeline, + # xl ORTStableDiffusionXLImg2ImgPipeline, ORTStableDiffusionXLInpaintPipeline, ORTStableDiffusionXLPipeline, diff --git a/optimum/onnxruntime/base.py b/optimum/onnxruntime/base.py index 845780cafa..4e9f23b2d1 100644 --- a/optimum/onnxruntime/base.py +++ b/optimum/onnxruntime/base.py @@ -26,7 +26,7 @@ from ..utils.logging import warn_once from .io_binding import TypeHelper from .modeling_ort import ORTModel -from .utils import get_ordered_input_names, logging +from .utils import logging logger = logging.get_logger(__name__) @@ -38,6 +38,11 @@ class ORTModelPart: It has its own `onnxruntime.InferenceSession`, and can perform a forward pass. """ + # should be in an ORTMixin + _prepare_io_binding = ORTModel._prepare_io_binding + _prepare_output_buffer = ORTModel._prepare_output_buffer + _output_shape_inference = ORTModel._output_shape_inference + _prepare_onnx_inputs = ORTModel._prepare_onnx_inputs _prepare_onnx_outputs = ORTModel._prepare_onnx_outputs @@ -48,10 +53,12 @@ def __init__(self, session: InferenceSession, parent_model: "ORTModel"): self.input_names = {input_key.name: idx for idx, input_key in enumerate(self.session.get_inputs())} self.output_names = {output_key.name: idx for idx, output_key in enumerate(self.session.get_outputs())} + self.input_dtypes = {input_key.name: input_key.type for input_key in session.get_inputs()} self.output_dtypes = {output_key.name: output_key.type for output_key in session.get_outputs()} - self._ordered_input_names = get_ordered_input_names(self.input_names.keys(), func=self.forward) + self.input_shapes = {input_key.name: input_key.shape for input_key in session.get_inputs()} + self.output_shapes = {output_key.name: output_key.shape for output_key in session.get_outputs()} @property def device(self): @@ -118,27 +125,26 @@ def forward(self, input_ids: torch.LongTensor, attention_mask: torch.LongTensor, use_torch = isinstance(input_ids, torch.Tensor) self.parent_model.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.parent_model.use_io_binding: - model_inputs = [input_ids] - if "attention_mask" in self.input_names: - model_inputs.append(attention_mask) - io_binding, output_shapes, output_buffers = self.parent_model._prepare_io_binding( - self.session, - *model_inputs, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + } - io_binding.synchronize_inputs() - self.session.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.parent_model.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.session, model_inputs) + + if self.device.type == "cpu": + self.session.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) last_hidden_state = model_outputs["last_hidden_state"] @@ -257,9 +263,7 @@ def forward( decoder_attention_mask: Optional[torch.LongTensor] = None, encoder_attention_mask: Optional[torch.LongTensor] = None, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - labels: Optional[torch.LongTensor] = None, cache_position: Optional[torch.Tensor] = None, - use_cache_branch: None = None, ) -> Seq2SeqLMOutput: # Adding use_cache_branch in the signature here is just a hack for IO Binding @@ -279,6 +283,17 @@ def forward( input_ids, past_key_values, cache_position, use_torch=use_torch ) + model_inputs = { + "input_ids": input_ids, + "encoder_hidden_states": encoder_hidden_states, + "decoder_attention_mask": decoder_attention_mask, + "encoder_attention_mask": encoder_attention_mask, + "use_cache_branch": use_cache_branch_tensor, + "cache_position": cache_position, + } + if past_key_values is not None: + model_inputs.update(zip(self.key_value_input_names, past_key_values)) + if self.parent_model.use_io_binding: known_output_shapes = self.compute_past_key_values_output_shapes( input_ids, @@ -286,53 +301,27 @@ def forward( use_cache_branch=use_cache_branch_tensor.item() if use_cache_branch_tensor is not None else None, past_key_values=past_key_values, ) - outputs_to_not_bind = self.get_outputs_not_to_bind(use_merged_cache) - # TODO: fix transformers generate to have contiguous input_ids here already - # For an unknown reason, calling `contiguous()` here is necessary to not have errors - # on CPU EP with batch size > 1, despite it being also called in _prepare_io_binding.g - model_inputs = [input_ids.contiguous()] - - if "encoder_hidden_states" in self.input_names: - model_inputs.append(encoder_hidden_states) - - if "decoder_attention_mask" in self.input_names: - model_inputs.append(decoder_attention_mask) - - if "encoder_attention_mask" in self.input_names: - model_inputs.append(encoder_attention_mask) - - if past_key_values is not None: - model_inputs += past_key_values - - if "labels" in self.input_names: - model_inputs.append(labels) - known_output_shapes.update({"loss": []}) - - if use_cache_branch_tensor is not None: - model_inputs.append(use_cache_branch_tensor) - - if "cache_position" in self.input_names: - model_inputs.append(cache_position) - - io_binding, output_shapes, output_buffers = self.parent_model._prepare_io_binding( + io_binding, output_shapes, output_buffers = self._prepare_io_binding( self.session, - *model_inputs, + model_inputs, known_output_shapes=known_output_shapes, - ordered_input_names=self._ordered_input_names, outputs_to_not_bind=outputs_to_not_bind, ) + if self.device.type == "cpu": + self.session.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() + # Set -1 for sequence_length as it could be larger than the real sequence_length for name, shape in output_shapes.items(): if name in self.key_value_output_names: output_shapes[name] = shape[:2] + (-1,) + shape[3:] - io_binding.synchronize_inputs() - self.session.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() - # Tuple of length equal to : number of layer * number of past_key_value per decoder layer (2 corresponds to the # self-attention layer and 2 to the cross-attention layer) out_past_key_values = () @@ -350,7 +339,7 @@ def forward( if not self.use_past_in_outputs: out_past_key_values = None - elif not self.use_past_in_inputs or use_merged_no_cache: + elif not self.use_past_in_inputs or use_merged_no_cache or self.no_cross_attention_cache: out_past_key_values = tuple( out_past_key_values[i : i + self.num_pkv] for i in range(0, len(out_past_key_values), self.num_pkv) ) @@ -382,21 +371,9 @@ def forward( else: raise ValueError("Unsupported num_pkv") else: - model_inputs = { - "input_ids": input_ids, - "encoder_hidden_states": encoder_hidden_states, - "decoder_attention_mask": decoder_attention_mask, - "encoder_attention_mask": encoder_attention_mask, - "use_cache_branch": use_cache_branch_tensor, - "cache_position": cache_position, - "labels": labels, - } - if past_key_values is not None: - model_inputs.update(zip(self.key_value_input_names, past_key_values)) - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.session.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) # TODO: using a new variable out_past_key_values is memory inefficient, # past_key_values is not used anymore at this point diff --git a/optimum/onnxruntime/configuration.py b/optimum/onnxruntime/configuration.py index 2e3d9f32d6..adc1984795 100644 --- a/optimum/onnxruntime/configuration.py +++ b/optimum/onnxruntime/configuration.py @@ -18,9 +18,8 @@ from dataclasses import asdict, dataclass, field from enum import Enum from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union -from datasets import Dataset from packaging.version import Version, parse from onnxruntime import __version__ as ort_version @@ -33,6 +32,10 @@ from ..utils import logging +if TYPE_CHECKING: + from datasets import Dataset + + logger = logging.get_logger(__name__) # This value is used to indicate ORT which axis it should use to quantize an operator "per-channel" @@ -117,7 +120,9 @@ def create_calibrator( class AutoCalibrationConfig: @staticmethod - def minmax(dataset: Dataset, moving_average: bool = False, averaging_constant: float = 0.01) -> CalibrationConfig: + def minmax( + dataset: "Dataset", moving_average: bool = False, averaging_constant: float = 0.01 + ) -> CalibrationConfig: """ Args: dataset (`Dataset`): @@ -151,7 +156,7 @@ def minmax(dataset: Dataset, moving_average: bool = False, averaging_constant: f @staticmethod def entropy( - dataset: Dataset, + dataset: "Dataset", num_bins: int = 128, num_quantized_bins: int = 128, ) -> CalibrationConfig: @@ -188,7 +193,7 @@ def entropy( ) @staticmethod - def percentiles(dataset: Dataset, num_bins: int = 2048, percentile: float = 99.999) -> CalibrationConfig: + def percentiles(dataset: "Dataset", num_bins: int = 2048, percentile: float = 99.999) -> CalibrationConfig: """ Args: dataset (`Dataset`): diff --git a/optimum/onnxruntime/model.py b/optimum/onnxruntime/model.py index caa662f382..4182abc925 100644 --- a/optimum/onnxruntime/model.py +++ b/optimum/onnxruntime/model.py @@ -14,10 +14,9 @@ import logging import os -from typing import Callable, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union import numpy as np -from datasets import Dataset from transformers import EvalPrediction from transformers.trainer_pt_utils import nested_concat from transformers.trainer_utils import EvalLoopOutput @@ -25,6 +24,10 @@ from onnxruntime import InferenceSession +if TYPE_CHECKING: + from datasets import Dataset + + logger = logging.getLogger(__name__) @@ -59,7 +62,7 @@ def __init__( self.session = InferenceSession(str(model_path), providers=[execution_provider]) self.onnx_input_names = {input_key.name: idx for idx, input_key in enumerate(self.session.get_inputs())} - def evaluation_loop(self, dataset: Dataset): + def evaluation_loop(self, dataset: "Dataset"): """ Run evaluation and returns metrics and predictions. diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py index 984d7f22eb..9afa1bf19a 100644 --- a/optimum/onnxruntime/modeling_decoder.py +++ b/optimum/onnxruntime/modeling_decoder.py @@ -31,7 +31,7 @@ from ..exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS, main_export from ..onnx.utils import check_model_uses_external_data -from ..utils import NormalizedConfigManager, check_if_transformers_greater +from ..utils import NormalizedConfigManager, is_transformers_version from ..utils.modeling_utils import MODEL_TO_PATCH_FOR_PAST from ..utils.save_utils import maybe_save_preprocessors from .constants import DECODER_MERGED_ONNX_FILE_PATTERN, DECODER_ONNX_FILE_PATTERN, DECODER_WITH_PAST_ONNX_FILE_PATTERN @@ -43,7 +43,7 @@ if TYPE_CHECKING: from transformers import PretrainedConfig -if check_if_transformers_greater("4.25.0"): +if is_transformers_version(">=", "4.25.0"): from transformers.generation import GenerationMixin else: from transformers.generation_utils import GenerationMixin # type: ignore # noqa: F401 @@ -149,7 +149,7 @@ def __init__( self.generation_config = generation_config - if check_if_transformers_greater("4.44.99"): + if is_transformers_version(">=", "4.44.99"): misplaced_generation_parameters = self.config._get_non_default_generation_parameters() if len(misplaced_generation_parameters) > 0: logger.warning( @@ -209,7 +209,6 @@ def forward( attention_mask: Optional[torch.FloatTensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - labels: Optional[torch.LongTensor] = None, use_cache_branch: bool = None, **kwargs, ) -> CausalLMOutputWithPast: @@ -218,8 +217,7 @@ def forward( self.raise_on_numpy_input_io_binding(use_torch) known_output_shapes = {} - use_cache_branch = None - loss = None + if self.use_cache: if past_key_values is not None: # Flatten the past_key_values (gpt_bigcode has fused key/value cache, so no need to flatten it) @@ -233,35 +231,28 @@ def forward( input_ids, past_key_values, use_torch ) - if self.use_io_binding: - # TODO: fix transformers generate to have contiguous input_ids here already - # For an unknown reason, calling `contiguous()` here is necessary to not have errors - # on CPU EP with batch size > 1, despite it being also called in _prepare_io_binding. - # I suspect the reason is the contiguous python list that messes something up? - model_inputs = [input_ids.contiguous()] - - if "attention_mask" in self.input_names: - model_inputs.append(attention_mask) - - if "position_ids" in self.input_names: - if position_ids is None: - raise ValueError("position_ids was not passed but is a required input for this ONNX model.") - model_inputs.append(position_ids.contiguous()) - - if past_key_values is not None: - model_inputs += past_key_values + # Create position_ids on the fly for batch generation + if "position_ids" in self.input_names and position_ids is None and attention_mask is not None: + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -1].unsqueeze(-1) - if use_cache_branch is not None: - model_inputs.append(use_cache_branch) + model_inputs = { + "input_ids": input_ids, + "position_ids": position_ids, + "attention_mask": attention_mask, + "use_cache_branch": use_cache_branch, + } - if "labels" in self.input_names: - model_inputs.append(labels) - known_output_shapes.update({"loss": []}) + if past_key_values is not None: + model_inputs.update( + zip(self.key_value_input_names, past_key_values), + ) - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - *model_inputs, - known_output_shapes=known_output_shapes, - ordered_input_names=self._ordered_input_names, + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding( + self.model, model_inputs, known_output_shapes=known_output_shapes ) if self.device.type == "cpu": @@ -271,32 +262,19 @@ def forward( self.model.run_with_iobinding(io_binding) io_binding.synchronize_outputs() + loss = output_buffers.get("loss", None) + logits = output_buffers["logits"].view(output_shapes["logits"]) + if self.use_cache: # Tuple of length equal to : number of layer * number of past_key_value per decoder layer(2 for the self-attention) past_key_values = tuple( output_buffers[name].view(output_shapes[name]) for name in self.key_value_output_names ) - logits = output_buffers["logits"].view(output_shapes["logits"]) - - if "loss" in self.output_names: - loss = output_buffers["loss"].view(output_shapes["loss"]) else: - model_inputs = { - "input_ids": input_ids, - "position_ids": position_ids, - "attention_mask": attention_mask, - "use_cache_branch": use_cache_branch, - "labels": labels, - } - if past_key_values is not None: - model_inputs.update( - zip(self.key_value_input_names, past_key_values), - ) - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) loss = model_outputs.get("loss", None) logits = model_outputs["logits"] @@ -340,7 +318,7 @@ def prepare_past_key_values( if self.model_type == "gemma": num_attention_heads = self.normalized_config.num_key_value_heads embed_size_per_head = self.normalized_config.head_dim - elif self.model_type in {"mistral", "llama", "qwen2"}: + elif self.model_type in {"mistral", "llama", "qwen2", "granite"}: num_attention_heads = self.normalized_config.num_key_value_heads else: num_attention_heads = self.normalized_config.num_attention_heads @@ -562,7 +540,7 @@ def _from_pretrained( ) # Since transformers 4.44, the bloom model has been updated to use the standard cache format - use_old_bloom_modeling = not check_if_transformers_greater("4.44") + use_old_bloom_modeling = not is_transformers_version(">=", "4.44") for input_name in input_dims.keys(): if input_dims[input_name][0] == "batch_size x num_heads": use_old_bloom_modeling = True diff --git a/optimum/onnxruntime/modeling_diffusion.py b/optimum/onnxruntime/modeling_diffusion.py index 3899a7b36b..193d75e0d4 100644 --- a/optimum/onnxruntime/modeling_diffusion.py +++ b/optimum/onnxruntime/modeling_diffusion.py @@ -51,13 +51,15 @@ from transformers.modeling_outputs import ModelOutput import onnxruntime as ort -from optimum.utils import check_if_diffusers_greater +from optimum.utils import is_diffusers_version from ..exporters.onnx import main_export from ..onnx.utils import _get_model_external_data_paths from ..utils import ( DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER, + DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER, DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER, + DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER, DIFFUSION_MODEL_UNET_SUBFOLDER, DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER, DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER, @@ -73,10 +75,10 @@ ) -if check_if_diffusers_greater("0.25.0"): +if is_diffusers_version(">=", "0.25.0"): from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution else: - from diffusers.models.vae import DiagonalGaussianDistribution + from diffusers.models.vae import DiagonalGaussianDistribution # type: ignore logger = logging.getLogger(__name__) @@ -92,15 +94,18 @@ class ORTDiffusionPipeline(ORTModel, DiffusionPipeline): def __init__( self, scheduler: "SchedulerMixin", - unet_session: ort.InferenceSession, vae_decoder_session: ort.InferenceSession, # optional pipeline models + unet_session: Optional[ort.InferenceSession] = None, + transformer_session: Optional[ort.InferenceSession] = None, vae_encoder_session: Optional[ort.InferenceSession] = None, text_encoder_session: Optional[ort.InferenceSession] = None, text_encoder_2_session: Optional[ort.InferenceSession] = None, + text_encoder_3_session: Optional[ort.InferenceSession] = None, # optional pipeline submodels tokenizer: Optional["CLIPTokenizer"] = None, tokenizer_2: Optional["CLIPTokenizer"] = None, + tokenizer_3: Optional["CLIPTokenizer"] = None, feature_extractor: Optional["CLIPFeatureExtractor"] = None, # stable diffusion xl specific arguments force_zeros_for_empty_prompt: bool = True, @@ -111,16 +116,20 @@ def __init__( model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None, **kwargs, ): - self.unet = ORTModelUnet(unet_session, self) - self.vae_decoder = ORTModelVaeDecoder(vae_decoder_session, self) - self.vae_encoder = ORTModelVaeEncoder(vae_encoder_session, self) if vae_encoder_session is not None else None + self.unet = ORTModelUnet(unet_session, self) if unet_session is not None else None + self.transformer = ORTModelTransformer(transformer_session, self) if transformer_session is not None else None self.text_encoder = ( ORTModelTextEncoder(text_encoder_session, self) if text_encoder_session is not None else None ) self.text_encoder_2 = ( ORTModelTextEncoder(text_encoder_2_session, self) if text_encoder_2_session is not None else None ) + self.text_encoder_3 = ( + ORTModelTextEncoder(text_encoder_3_session, self) if text_encoder_3_session is not None else None + ) # We wrap the VAE Decoder & Encoder in a single object to simulate diffusers API + self.vae_encoder = ORTModelVaeEncoder(vae_encoder_session, self) if vae_encoder_session is not None else None + self.vae_decoder = ORTModelVaeDecoder(vae_decoder_session, self) if vae_decoder_session is not None else None self.vae = ORTWrapperVae(self.vae_encoder, self.vae_decoder) # we allow passing these as torch models for now @@ -130,18 +139,22 @@ def __init__( self.scheduler = scheduler self.tokenizer = tokenizer self.tokenizer_2 = tokenizer_2 + self.tokenizer_3 = tokenizer_3 self.feature_extractor = feature_extractor all_pipeline_init_args = { "vae": self.vae, "unet": self.unet, + "transformer": self.transformer, "text_encoder": self.text_encoder, "text_encoder_2": self.text_encoder_2, + "text_encoder_3": self.text_encoder_3, "safety_checker": self.safety_checker, "image_encoder": self.image_encoder, "scheduler": self.scheduler, "tokenizer": self.tokenizer, "tokenizer_2": self.tokenizer_2, + "tokenizer_3": self.tokenizer_3, "feature_extractor": self.feature_extractor, "requires_aesthetics_score": requires_aesthetics_score, "force_zeros_for_empty_prompt": force_zeros_for_empty_prompt, @@ -157,7 +170,10 @@ def __init__( # inits ort specific attributes self.shared_attributes_init( - model=unet_session, use_io_binding=use_io_binding, model_save_dir=model_save_dir, **kwargs + model=unet_session if unet_session is not None else transformer_session, + use_io_binding=use_io_binding, + model_save_dir=model_save_dir, + **kwargs, ) def _save_pretrained(self, save_directory: Union[str, Path]): @@ -165,10 +181,12 @@ def _save_pretrained(self, save_directory: Union[str, Path]): models_to_save_paths = { (self.unet, save_directory / DIFFUSION_MODEL_UNET_SUBFOLDER), + (self.transformer, save_directory / DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER), (self.vae_decoder, save_directory / DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER), (self.vae_encoder, save_directory / DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER), (self.text_encoder, save_directory / DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER), (self.text_encoder_2, save_directory / DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER), + (self.text_encoder_3, save_directory / DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER), } for model, save_path in models_to_save_paths: if model is not None: @@ -192,6 +210,8 @@ def _save_pretrained(self, save_directory: Union[str, Path]): self.tokenizer.save_pretrained(save_directory / "tokenizer") if self.tokenizer_2 is not None: self.tokenizer_2.save_pretrained(save_directory / "tokenizer_2") + if self.tokenizer_3 is not None: + self.tokenizer_3.save_pretrained(save_directory / "tokenizer_3") if self.feature_extractor is not None: self.feature_extractor.save_pretrained(save_directory / "feature_extractor") @@ -208,10 +228,12 @@ def _from_pretrained( cache_dir: str = HUGGINGFACE_HUB_CACHE, token: Optional[Union[bool, str]] = None, unet_file_name: str = ONNX_WEIGHTS_NAME, + transformer_file_name: str = ONNX_WEIGHTS_NAME, vae_decoder_file_name: str = ONNX_WEIGHTS_NAME, vae_encoder_file_name: str = ONNX_WEIGHTS_NAME, text_encoder_file_name: str = ONNX_WEIGHTS_NAME, text_encoder_2_file_name: str = ONNX_WEIGHTS_NAME, + text_encoder_3_file_name: str = ONNX_WEIGHTS_NAME, use_io_binding: Optional[bool] = None, provider: str = "CPUExecutionProvider", provider_options: Optional[Dict[str, Any]] = None, @@ -230,10 +252,12 @@ def _from_pretrained( allow_patterns.update( { unet_file_name, + transformer_file_name, vae_decoder_file_name, vae_encoder_file_name, text_encoder_file_name, text_encoder_2_file_name, + text_encoder_3_file_name, SCHEDULER_CONFIG_NAME, cls.config_name, CONFIG_NAME, @@ -259,10 +283,12 @@ def _from_pretrained( model_paths = { "unet": model_save_path / DIFFUSION_MODEL_UNET_SUBFOLDER / unet_file_name, + "transformer": model_save_path / DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER / transformer_file_name, "vae_decoder": model_save_path / DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER / vae_decoder_file_name, "vae_encoder": model_save_path / DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER / vae_encoder_file_name, "text_encoder": model_save_path / DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER / text_encoder_file_name, "text_encoder_2": model_save_path / DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER / text_encoder_2_file_name, + "text_encoder_3": model_save_path / DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER / text_encoder_3_file_name, } sessions = {} @@ -276,7 +302,7 @@ def _from_pretrained( ) submodels = {} - for submodel in {"scheduler", "tokenizer", "tokenizer_2", "feature_extractor"}: + for submodel in {"scheduler", "tokenizer", "tokenizer_2", "tokenizer_3", "feature_extractor"}: if kwargs.get(submodel, None) is not None: submodels[submodel] = kwargs.pop(submodel) elif config.get(submodel, (None, None))[0] is not None: @@ -385,17 +411,24 @@ def to(self, device: Union[torch.device, str, int]): if device.type == "cuda" and self.providers[0] == "TensorrtExecutionProvider": return self - self.unet.session.set_providers([provider], provider_options=[provider_options]) self.vae_decoder.session.set_providers([provider], provider_options=[provider_options]) + if self.unet is not None: + self.unet.session.set_providers([provider], provider_options=[provider_options]) + if self.transformer is not None: + self.transformer.session.set_providers([provider], provider_options=[provider_options]) if self.vae_encoder is not None: self.vae_encoder.session.set_providers([provider], provider_options=[provider_options]) if self.text_encoder is not None: self.text_encoder.session.set_providers([provider], provider_options=[provider_options]) if self.text_encoder_2 is not None: self.text_encoder_2.session.set_providers([provider], provider_options=[provider_options]) + if self.text_encoder_3 is not None: + self.text_encoder_3.session.set_providers([provider], provider_options=[provider_options]) - self.providers = self.unet.session.get_providers() + self.providers = ( + self.unet.session.get_providers() if self.unet is not None else self.transformer.session.get_providers() + ) self._device = device return self @@ -404,16 +437,31 @@ def to(self, device: Union[torch.device, str, int]): def _load_config(cls, config_name_or_path: Union[str, os.PathLike], **kwargs): return cls.load_config(config_name_or_path, **kwargs) - def _save_config(self, save_directory): - self.save_config(save_directory) + def _save_config(self, save_directory: Union[str, Path]): + model_dir = ( + self.model_save_dir + if not isinstance(self.model_save_dir, TemporaryDirectory) + else self.model_save_dir.name + ) + save_dir = Path(save_directory) + original_config = Path(model_dir) / self.config_name + if original_config.exists(): + if not save_dir.exists(): + save_dir.mkdir(parents=True) + + shutil.copy(original_config, save_dir) + else: + self.save_config(save_directory) @property def components(self) -> Dict[str, Any]: components = { "vae": self.vae, "unet": self.unet, + "transformer": self.transformer, "text_encoder": self.text_encoder, "text_encoder_2": self.text_encoder_2, + "text_encoder_3": self.text_encoder_3, "safety_checker": self.safety_checker, "image_encoder": self.image_encoder, } @@ -443,9 +491,13 @@ def __init__(self, session: ort.InferenceSession, parent_pipeline: ORTDiffusionP self.input_names = {input_key.name: idx for idx, input_key in enumerate(self.session.get_inputs())} self.output_names = {output_key.name: idx for idx, output_key in enumerate(self.session.get_outputs())} + self.input_dtypes = {input_key.name: input_key.type for input_key in self.session.get_inputs()} self.output_dtypes = {output_key.name: output_key.type for output_key in self.session.get_outputs()} + self.input_shapes = {input_key.name: input_key.shape for input_key in self.session.get_inputs()} + self.output_shapes = {output_key.name: output_key.shape for output_key in self.session.get_outputs()} + config_file_path = Path(session._model_path).parent / self.config_name if not config_file_path.is_file(): # config is mandatory for the model part to be used for inference @@ -543,13 +595,18 @@ def __init__(self, *args, **kwargs): ) self.register_to_config(time_cond_proj_dim=None) + if len(self.input_shapes["timestep"]) > 0: + logger.warning( + "The exported unet onnx model expects a non scalar timestep input. " + "We will have to unsqueeze the timestep input at each iteration which might be inefficient. " + "Please re-export the pipeline with newer version of optimum and diffusers to avoid this warning." + ) + def forward( self, sample: Union[np.ndarray, torch.Tensor], timestep: Union[np.ndarray, torch.Tensor], encoder_hidden_states: Union[np.ndarray, torch.Tensor], - text_embeds: Optional[Union[np.ndarray, torch.Tensor]] = None, - time_ids: Optional[Union[np.ndarray, torch.Tensor]] = None, timestep_cond: Optional[Union[np.ndarray, torch.Tensor]] = None, cross_attention_kwargs: Optional[Dict[str, Any]] = None, added_cond_kwargs: Optional[Dict[str, Any]] = None, @@ -557,15 +614,13 @@ def forward( ): use_torch = isinstance(sample, torch.Tensor) - if len(timestep.shape) == 0: + if len(self.input_shapes["timestep"]) > 0: timestep = timestep.unsqueeze(0) model_inputs = { "sample": sample, "timestep": timestep, "encoder_hidden_states": encoder_hidden_states, - "text_embeds": text_embeds, - "time_ids": time_ids, "timestep_cond": timestep_cond, **(cross_attention_kwargs or {}), **(added_cond_kwargs or {}), @@ -581,6 +636,42 @@ def forward( return ModelOutput(**model_outputs) +class ORTModelTransformer(ORTPipelinePart): + def forward( + self, + hidden_states: Union[np.ndarray, torch.Tensor], + encoder_hidden_states: Union[np.ndarray, torch.Tensor], + pooled_projections: Union[np.ndarray, torch.Tensor], + timestep: Union[np.ndarray, torch.Tensor], + guidance: Optional[Union[np.ndarray, torch.Tensor]] = None, + txt_ids: Optional[Union[np.ndarray, torch.Tensor]] = None, + img_ids: Optional[Union[np.ndarray, torch.Tensor]] = None, + joint_attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + ): + use_torch = isinstance(hidden_states, torch.Tensor) + + model_inputs = { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "pooled_projections": pooled_projections, + "timestep": timestep, + "guidance": guidance, + "txt_ids": txt_ids, + "img_ids": img_ids, + **(joint_attention_kwargs or {}), + } + + onnx_inputs = self.prepare_onnx_inputs(use_torch, **model_inputs) + onnx_outputs = self.session.run(None, onnx_inputs) + model_outputs = self.prepare_onnx_outputs(use_torch, *onnx_outputs) + + if return_dict: + return model_outputs + + return ModelOutput(**model_outputs) + + class ORTModelTextEncoder(ORTPipelinePart): def forward( self, @@ -599,11 +690,13 @@ def forward( if output_hidden_states: model_outputs["hidden_states"] = [] - for i in range(self.config.num_hidden_layers): + num_layers = self.num_hidden_layers if hasattr(self, "num_hidden_layers") else self.num_decoder_layers + for i in range(num_layers): model_outputs["hidden_states"].append(model_outputs.pop(f"hidden_states.{i}")) model_outputs["hidden_states"].append(model_outputs.get("last_hidden_state")) else: - for i in range(self.config.num_hidden_layers): + num_layers = self.num_hidden_layers if hasattr(self, "num_hidden_layers") else self.num_decoder_layers + for i in range(num_layers): model_outputs.pop(f"hidden_states.{i}", None) if return_dict: @@ -620,7 +713,7 @@ def __init__(self, *args, **kwargs): if not hasattr(self.config, "scaling_factor"): logger.warning( "The `scaling_factor` attribute is missing from the VAE encoder configuration. " - "Please re-export the model with newer version of optimum and diffusers." + "Please re-export the model with newer version of optimum and diffusers to avoid this warning." ) self.register_to_config(scaling_factor=2 ** (len(self.config.block_out_channels) - 1)) @@ -660,7 +753,7 @@ def __init__(self, *args, **kwargs): if not hasattr(self.config, "scaling_factor"): logger.warning( "The `scaling_factor` attribute is missing from the VAE decoder configuration. " - "Please re-export the model with newer version of optimum and diffusers." + "Please re-export the model with newer version of optimum and diffusers to avoid this warning." ) self.register_to_config(scaling_factor=2 ** (len(self.config.block_out_channels) - 1)) @@ -871,6 +964,80 @@ class ORTLatentConsistencyModelImg2ImgPipeline(ORTDiffusionPipeline, LatentConsi auto_model_class = LatentConsistencyModelImg2ImgPipeline +class ORTUnavailablePipeline: + MIN_VERSION = None + + def __init__(self, *args, **kwargs): + raise NotImplementedError( + f"The pipeline {self.__class__.__name__} is not available in the current version of `diffusers`. " + f"Please upgrade `diffusers` to {self.MIN_VERSION} or later." + ) + + +if is_diffusers_version(">=", "0.29.0"): + from diffusers import StableDiffusion3Img2ImgPipeline, StableDiffusion3Pipeline + + @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) + class ORTStableDiffusion3Pipeline(ORTDiffusionPipeline, StableDiffusion3Pipeline): + """ + ONNX Runtime-powered stable diffusion pipeline corresponding to [diffusers.StableDiffusion3Pipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img#diffusers.StableDiffusion3Pipeline). + """ + + main_input_name = "prompt" + export_feature = "text-to-image" + auto_model_class = StableDiffusion3Pipeline + + @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) + class ORTStableDiffusion3Img2ImgPipeline(ORTDiffusionPipeline, StableDiffusion3Img2ImgPipeline): + """ + ONNX Runtime-powered stable diffusion pipeline corresponding to [diffusers.StableDiffusion3Img2ImgPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/img2img#diffusers.StableDiffusion3Img2ImgPipeline). + """ + + main_input_name = "image" + export_feature = "image-to-image" + auto_model_class = StableDiffusion3Img2ImgPipeline + +else: + + class ORTStableDiffusion3Pipeline(ORTUnavailablePipeline): + MIN_VERSION = "0.29.0" + + class ORTStableDiffusion3Img2ImgPipeline(ORTUnavailablePipeline): + MIN_VERSION = "0.29.0" + + +if is_diffusers_version(">=", "0.30.0"): + from diffusers import FluxPipeline, StableDiffusion3InpaintPipeline + + @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) + class ORTStableDiffusion3InpaintPipeline(ORTDiffusionPipeline, StableDiffusion3InpaintPipeline): + """ + ONNX Runtime-powered stable diffusion pipeline corresponding to [diffusers.StableDiffusion3InpaintPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/inpaint#diffusers.StableDiffusion3InpaintPipeline). + """ + + main_input_name = "prompt" + export_feature = "inpainting" + auto_model_class = StableDiffusion3InpaintPipeline + + @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) + class ORTFluxPipeline(ORTDiffusionPipeline, FluxPipeline): + """ + ONNX Runtime-powered stable diffusion pipeline corresponding to [diffusers.FluxPipeline](https://huggingface.co/docs/diffusers/api/pipelines/flux/text2img#diffusers.FluxPipeline). + """ + + main_input_name = "prompt" + export_feature = "text-to-image" + auto_model_class = FluxPipeline + +else: + + class ORTStableDiffusion3InpaintPipeline(ORTUnavailablePipeline): + MIN_VERSION = "0.30.0" + + class ORTFluxPipeline(ORTUnavailablePipeline): + MIN_VERSION = "0.30.0" + + SUPPORTED_ORT_PIPELINES = [ ORTStableDiffusionPipeline, ORTStableDiffusionImg2ImgPipeline, @@ -880,6 +1047,10 @@ class ORTLatentConsistencyModelImg2ImgPipeline(ORTDiffusionPipeline, LatentConsi ORTStableDiffusionXLInpaintPipeline, ORTLatentConsistencyModelPipeline, ORTLatentConsistencyModelImg2ImgPipeline, + ORTStableDiffusion3Pipeline, + ORTStableDiffusion3Img2ImgPipeline, + ORTStableDiffusion3InpaintPipeline, + ORTFluxPipeline, ] @@ -897,23 +1068,27 @@ def _get_ort_class(pipeline_class_name: str, throw_error_if_not_exist: bool = Tr ORT_TEXT2IMAGE_PIPELINES_MAPPING = OrderedDict( [ + ("flux", ORTFluxPipeline), + ("latent-consistency", ORTLatentConsistencyModelPipeline), ("stable-diffusion", ORTStableDiffusionPipeline), + ("stable-diffusion-3", ORTStableDiffusion3Pipeline), ("stable-diffusion-xl", ORTStableDiffusionXLPipeline), - ("latent-consistency", ORTLatentConsistencyModelPipeline), ] ) ORT_IMAGE2IMAGE_PIPELINES_MAPPING = OrderedDict( [ + ("latent-consistency", ORTLatentConsistencyModelImg2ImgPipeline), ("stable-diffusion", ORTStableDiffusionImg2ImgPipeline), + ("stable-diffusion-3", ORTStableDiffusion3Img2ImgPipeline), ("stable-diffusion-xl", ORTStableDiffusionXLImg2ImgPipeline), - ("latent-consistency", ORTLatentConsistencyModelImg2ImgPipeline), ] ) ORT_INPAINT_PIPELINES_MAPPING = OrderedDict( [ ("stable-diffusion", ORTStableDiffusionInpaintPipeline), + ("stable-diffusion-3", ORTStableDiffusion3InpaintPipeline), ("stable-diffusion-xl", ORTStableDiffusionXLInpaintPipeline), ] ) diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py index 8e5a814b68..e9633343c7 100644 --- a/optimum/onnxruntime/modeling_ort.py +++ b/optimum/onnxruntime/modeling_ort.py @@ -72,7 +72,6 @@ ONNX_WEIGHTS_NAME, check_io_binding, get_device_for_provider, - get_ordered_input_names, get_provider_for_device, parse_device, validate_provider_availability, @@ -276,8 +275,6 @@ def __init__( self.output_names = {output_key.name: idx for idx, output_key in enumerate(model.get_outputs())} self.output_dtypes = {output_key.name: output_key.type for output_key in model.get_outputs()} - self._ordered_input_names = get_ordered_input_names(self.input_names.keys(), func=self.forward) - @property def dtype(self) -> torch.dtype: """ @@ -773,43 +770,23 @@ def _output_shape_inference(self, axis_name: Union[str, int], dimensions: Dict[s """ if isinstance(axis_name, int): return axis_name - # It is actually covered below, but this is to make things faster. + elif axis_name in dimensions: return dimensions[axis_name] - # Tokens is going to be populated by iterating over every match for the self.output_shape_inference_pattern. - # This pattern matches 4 things: axis names, integer values, operators (+, -, *, /) and parenthesis. - tokens = [] - for idx, match_ in enumerate(re.finditer(self.output_shape_inference_pattern, axis_name)): - groups = match_.groups() - matched_group = None - for idx, group in enumerate(groups): - if group is not None: - matched_group = idx - break - - # For every match except an axis name, we simply append the content of the match to the tokens list. - # For an axis name, we check if it is specified in the `dimensions` dictionary. If for some reason it is - # not there, or its value not an integer, the shape inference process stops and we return the axis name as - # is. - if matched_group == 0: - dim = dimensions.get(groups[0], None) - if dim is None or not isinstance(dim, int): - return axis_name - tokens.append(str(dim)) - else: - tokens.append(groups[matched_group]) + # faster way to do the same thing, assuming the axis names are well defined (by us in the exporter config) + tokens = axis_name.split(" ") + for idx, token in enumerate(tokens): + if token in dimensions: + tokens[idx] = str(dimensions[token]) - # Here it should not be problematic to use eval since anything not matching the pattern would trigger an - # exception. return int(eval(" ".join(tokens))) # TODO: this method is bloated with state arguments (that are accesible using self) why ? def _prepare_io_binding( self, model: ort.InferenceSession, - *model_inputs: torch.Tensor, - ordered_input_names: List[str], + model_inputs: Dict[str, torch.Tensor], known_output_shapes: Optional[Dict[str, Tuple[int]]] = None, outputs_to_not_bind: Optional[Union[Set[str], str]] = None, ) -> Tuple[ort.IOBinding, Dict[str, Tuple[int]], Dict[str, torch.Tensor]]: @@ -819,10 +796,8 @@ def _prepare_io_binding( Args: model (`ort.InferenceSession`): The model for which we want to bind the inputs and outputs. - *model_inputs: - The inputs of the model. - ordered_input_names (`List[str]`): - Names of the inputs, that must match with the order of model_inputs. + model_inputs (`Dict[str, torch.Tensor]`): + The inputs to bind to the model. known_output_shapes (`Optional[Dict[str, Tuple[int]]]`, defaults to `None`): It can be hard to infer all the output shapes from the inputs only. For instance for the past key / values. It is possible to explicitely pass the shape via this argument. @@ -835,36 +810,39 @@ def _prepare_io_binding( """ io_binding = model.io_binding() - name_to_np_type = TypeHelper.get_io_numpy_type_map(model) + input_shapes = {} + for input_name in self.input_names.keys(): + input_shapes[input_name] = model_inputs[input_name].shape - input_name_to_shape = {} - for idx, tensor in enumerate(model_inputs): - if tensor is None: - continue - name = ordered_input_names[idx] - tensor = tensor.contiguous() - input_name_to_shape[name] = tensor.shape + if not model_inputs[input_name].is_contiguous(): + model_inputs[input_name] = model_inputs[input_name].contiguous() + + tensor_dtype = model_inputs[input_name].dtype + expected_dtype = TypeHelper.ort_type_to_torch_type(self.input_dtypes[input_name]) + if tensor_dtype != expected_dtype: + model_inputs[input_name] = model_inputs[input_name].to(expected_dtype) - data_ptr = tensor.data_ptr() - if "past" in name and data_ptr == 0: + data_ptr = model_inputs[input_name].data_ptr() + if data_ptr == 0: # During first generation, sequence_length can be 0 when use_cache=True, which results in data_ptr to also be 0. # To keep compatibility with IO binding, we pass the data pointer of input_ids instead. This will have no impact because past_key_values will not be used during the first generation. - data_ptr = model_inputs[0].data_ptr() + data_ptr = model_inputs["input_ids"].data_ptr() io_binding.bind_input( - name, - tensor.device.type, + input_name, + self.device.type, IOBindingHelper.get_device_index(self.device), - name_to_np_type[name], - tuple(tensor.shape), + TypeHelper.ort_type_to_numpy_type(self.input_dtypes[input_name]), + model_inputs[input_name].shape, data_ptr, ) + dimensions = {} for input_ in model.get_inputs(): shape = input_.shape for idx, axis in enumerate(shape): if isinstance(axis, str): - dimensions[axis] = input_name_to_shape[input_.name][idx] + dimensions[axis] = input_shapes[input_.name][idx] output_shapes = {} output_buffers = {} @@ -887,32 +865,25 @@ def _prepare_io_binding( output_shape = [] for axis_name in output_node.shape: output_shape.append(self._output_shape_inference(axis_name, dimensions)) + output_buffer = self._prepare_output_buffer(model, output_shape, output_name) + data_ptr = output_buffer.data_ptr() + io_binding.bind_output( output_name, - output_buffer.device.type, + self.device.type, IOBindingHelper.get_device_index(self.device), - name_to_np_type[output_name], + TypeHelper.ort_type_to_numpy_type(output_node.type), output_shape, - output_buffer.data_ptr(), + data_ptr, ) - output_shapes[output_name] = output_shape + output_buffers[output_name] = output_buffer + output_shapes[output_name] = output_shape return io_binding, output_shapes, output_buffers - def prepare_io_binding( - self, *model_inputs, ordered_input_names, outputs_to_not_bind=None, known_output_shapes=None - ): - return self._prepare_io_binding( - self.model, - *model_inputs, - ordered_input_names=ordered_input_names, - known_output_shapes=known_output_shapes, - outputs_to_not_bind=outputs_to_not_bind, - ) - def raise_on_numpy_input_io_binding(self, use_torch: bool): """ Raises an error if IO Binding is requested although the tensor used are numpy arrays. @@ -928,29 +899,57 @@ def raise_on_numpy_input_io_binding(self, use_torch: bool): ) def _prepare_onnx_inputs( - self, use_torch: bool, **inputs: Union[torch.Tensor, np.ndarray] + self, use_torch: bool, model_inputs: Dict[str, Union[torch.Tensor, np.ndarray]] ) -> Dict[str, np.ndarray]: + """ + Prepares the inputs for ONNX Runtime by converting them to numpy arrays with the expected dtype. + + Args: + use_torch (`bool`): + Whether the inputs are torch.Tensor or not. + inputs (`Dict[str, Union[torch.Tensor, np.ndarray]]`): + The inputs to prepare for ONNX Runtime. + + Returns: + `Dict[str, np.ndarray]`: The inputs prepared for ONNX Runtime. + """ + onnx_inputs = {} - # converts pytorch inputs into numpy inputs for onnx + for input_name in self.input_names.keys(): - onnx_inputs[input_name] = inputs.pop(input_name) + if model_inputs.get(input_name, None) is None: + raise ValueError(f"Input {input_name} is required by model but not provided.") if use_torch: - onnx_inputs[input_name] = onnx_inputs[input_name].numpy(force=True) + onnx_inputs[input_name] = model_inputs[input_name].numpy(force=True) + else: + onnx_inputs[input_name] = model_inputs[input_name] - if onnx_inputs[input_name].dtype != self.input_dtypes[input_name]: - onnx_inputs[input_name] = onnx_inputs[input_name].astype( - TypeHelper.ort_type_to_numpy_type(self.input_dtypes[input_name]) - ) + expected_dtype = TypeHelper.ort_type_to_numpy_type(self.input_dtypes[input_name]) + + if onnx_inputs[input_name].dtype != expected_dtype: + onnx_inputs[input_name] = onnx_inputs[input_name].astype(expected_dtype) return onnx_inputs def _prepare_onnx_outputs( - self, use_torch: bool, *onnx_outputs: np.ndarray + self, use_torch: bool, onnx_outputs: List[np.ndarray] ) -> Dict[str, Union[torch.Tensor, np.ndarray]]: + """ + Prepares the outputs from ONNX Runtime by converting them to torch.Tensor if requested. + + Args: + use_torch (`bool`): + Whether the outputs should be torch.Tensor or not. + onnx_outputs (`List[np.ndarray]`): + The outputs from ONNX Runtime. + + Returns: + `Dict[str, Union[torch.Tensor, np.ndarray]]`: The outputs prepared for the user. + """ + model_outputs = {} - # converts onnxruntime outputs into tensor for standard outputs for output_name, idx in self.output_names.items(): model_outputs[output_name] = onnx_outputs[idx] @@ -1088,26 +1087,28 @@ def forward( if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_ids, - attention_mask, - token_type_ids, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) if "last_hidden_state" in self.output_names: last_hidden_state = model_outputs["last_hidden_state"] @@ -1243,29 +1244,31 @@ def forward( use_torch = isinstance(input_ids, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if token_type_ids is None and "token_type_ids" in self.input_names: + if "token_type_ids" in self.input_names and token_type_ids is None: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_ids, - attention_mask, - token_type_ids, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -1338,28 +1341,25 @@ def forward( if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_ids, - attention_mask, - token_type_ids, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() - # TODO: this is the same routine in all io binding branches, should we refactor it into a prepare_io_binding_outputs method? start_logits = output_buffers["start_logits"].view(output_shapes["start_logits"]) end_logits = output_buffers["end_logits"].view(output_shapes["end_logits"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) start_logits = model_outputs["start_logits"] end_logits = model_outputs["end_logits"] @@ -1448,26 +1448,28 @@ def forward( if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_ids, - attention_mask, - token_type_ids, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -1541,26 +1543,28 @@ def forward( if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_ids, - attention_mask, - token_type_ids, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -1627,26 +1631,28 @@ def forward( if token_type_ids is None and "token_type_ids" in self.input_names: token_type_ids = torch.zeros_like(input_ids) if use_torch else np.zeros_like(input_ids) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_ids, - attention_mask, - token_type_ids, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -1696,7 +1702,7 @@ def forward( @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForImageClassification(ORTModel): """ - ONNX Model for image-classification tasks. This class officially supports beit, convnext, convnextv2, data2vec_vision, deit, levit, mobilenet_v1, mobilenet_v2, mobilevit, poolformer, resnet, segformer, swin, vit. + ONNX Model for image-classification tasks. This class officially supports beit, convnext, convnextv2, data2vec_vision, deit, dinov2, levit, mobilenet_v1, mobilenet_v2, mobilevit, poolformer, resnet, segformer, swin, swinv2, vit. """ auto_model_class = AutoModelForImageClassification @@ -1717,24 +1723,26 @@ def forward( use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - pixel_values, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "pixel_values": pixel_values, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"pixel_values": pixel_values} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -1784,7 +1792,7 @@ def forward( @add_end_docstrings(ONNX_MODEL_END_DOCSTRING) class ORTModelForSemanticSegmentation(ORTModel): """ - ONNX Model for semantic-segmentation, with an all-MLP decode head on top e.g. for ADE20k, CityScapes. This class officially supports segformer. + ONNX Model for semantic-segmentation, with an all-MLP decode head on top e.g. for ADE20k, CityScapes. This class officially supports maskformer, segformer. """ auto_model_class = AutoModelForSemanticSegmentation @@ -1805,24 +1813,26 @@ def forward( use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - pixel_values, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "pixel_values": pixel_values, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"pixel_values": pixel_values} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -1932,25 +1942,27 @@ def forward( use_torch = isinstance(model_input, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - model_input, - attention_mask, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + self.input_name: model_input, + "attention_mask": attention_mask, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {self.input_name: model_input, "attention_mask": attention_mask} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -2009,35 +2021,36 @@ def forward( use_torch = isinstance(input_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: - input_size = input_values.shape[1] - output_sizes = [] + model_inputs = { + "input_values": input_values, + } - def _conv_output_size(input_size, kernel_size, stride): - return (input_size - kernel_size) // stride + 1 + if self.use_io_binding: + batch_size = input_values.shape[0] + final_input_size = input_values.shape[-1] for kernel_size, stride in zip(self.config.conv_kernel, self.config.conv_stride): - input_size = _conv_output_size(input_size, kernel_size, stride) - output_sizes.append(input_size) + final_input_size = (final_input_size - kernel_size) // stride + 1 - known_output_shapes = {"logits": [input_values.shape[0], output_sizes[-1], self.config.vocab_size]} + known_output_shapes = {"logits": [batch_size, final_input_size, self.config.vocab_size]} - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_values, ordered_input_names=self._ordered_input_names, known_output_shapes=known_output_shapes + io_binding, output_shapes, output_buffers = self._prepare_io_binding( + self.model, model_inputs, known_output_shapes=known_output_shapes ) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) else: - model_inputs = {"input_values": input_values} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -2104,25 +2117,28 @@ def forward( use_torch = isinstance(input_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - input_values, ordered_input_names=self._ordered_input_names - ) + model_inputs = { + "input_values": input_values, + } + + if self.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.model, model_inputs) # run inference with binding & synchronize in case of multiple CUDA streams - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() logits = output_buffers["logits"].view(output_shapes["logits"]) embeddings = output_buffers["embeddings"].view(output_shapes["embeddings"]) else: - model_inputs = {"input_values": input_values} - - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] embeddings = model_outputs["embeddings"] @@ -2182,14 +2198,14 @@ def forward( use_torch = isinstance(input_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: + if self.use_io_binding: raise NotImplementedError() else: model_inputs = {"input_values": input_values} - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) logits = model_outputs["logits"] @@ -2241,29 +2257,33 @@ def forward( ): use_torch = isinstance(pixel_values, torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: - input_shapes = pixel_values.shape - io_binding, output_shapes, output_buffers = self.prepare_io_binding( - pixel_values, - ordered_input_names=self._ordered_input_names, - known_output_shapes={ - "reconstruction": [ - input_shapes[0], - input_shapes[1], - input_shapes[2] * self.config.upscale, - input_shapes[3] * self.config.upscale, - ] - }, + + model_inputs = { + "pixel_values": pixel_values, + } + + if self.use_io_binding: + batch_size, num_channels, height, width = pixel_values.shape + known_output_shapes = { + "reconstruction": [batch_size, num_channels, height * self.config.upscale, width * self.config.upscale] + } + + io_binding, output_shapes, output_buffers = self._prepare_io_binding( + self.model, model_inputs, known_output_shapes=known_output_shapes ) - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() + reconstruction = output_buffers["reconstruction"].view(output_shapes["reconstruction"]) else: - model_inputs = {"pixel_values": pixel_values} - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) reconstruction = model_outputs["reconstruction"] return ImageSuperResolutionOutput(reconstruction=reconstruction) @@ -2318,23 +2338,26 @@ def forward(self, **model_inputs: Union[torch.Tensor, np.ndarray]): use_torch = isinstance(next(iter(model_inputs.values())), torch.Tensor) self.raise_on_numpy_input_io_binding(use_torch) - if self.device.type == "cuda" and self.use_io_binding: + if self.use_io_binding: # TODO: should this be used in favor of `model.prepare_io_binding`? io_binding = IOBindingHelper.prepare_io_binding(self, **model_inputs) # run inference with binding - io_binding.synchronize_inputs() - self.model.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.model.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.model.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() model_outputs = {} for name, output in zip(self.output_names.keys(), io_binding._iobinding.get_outputs()): model_outputs[name] = IOBindingHelper.to_pytorch(output) else: - onnx_inputs = self._prepare_onnx_inputs(use_torch, **model_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) onnx_outputs = self.model.run(None, onnx_inputs) - model_outputs = self._prepare_onnx_outputs(use_torch, *onnx_outputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) # converts output to namedtuple for pipelines post-processing return ModelOutput(**model_outputs) diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py index 27e0dc01b4..a3063826be 100644 --- a/optimum/onnxruntime/modeling_seq2seq.py +++ b/optimum/onnxruntime/modeling_seq2seq.py @@ -23,7 +23,6 @@ from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union -import numpy as np import torch from huggingface_hub import hf_hub_download from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE @@ -43,7 +42,7 @@ from ..exporters.onnx import main_export from ..onnx.utils import _get_external_data_paths -from ..utils import check_if_transformers_greater +from ..utils import is_transformers_version from ..utils.file_utils import validate_file_exists from ..utils.save_utils import maybe_load_preprocessors, maybe_save_preprocessors from .base import ORTDecoderForSeq2Seq, ORTEncoder @@ -64,13 +63,13 @@ ) -if check_if_transformers_greater("4.25.0"): +if is_transformers_version(">=", "4.25.0"): from transformers.generation import GenerationMixin else: from transformers.generation_utils import GenerationMixin # type: ignore -if check_if_transformers_greater("4.43.0"): +if is_transformers_version(">=", "4.43.0"): from transformers.cache_utils import EncoderDecoderCache else: EncoderDecoderCache = dict @@ -363,43 +362,28 @@ def forward( use_torch = isinstance(input_features, torch.Tensor) self.parent_model.raise_on_numpy_input_io_binding(use_torch) - if self.parent_model.device.type == "cuda" and self.parent_model.use_io_binding: - model_inputs = ( - [input_features, attention_mask] if "attention_mask" in self.input_names else [input_features] - ) - io_binding, output_shapes, output_buffers = self.parent_model._prepare_io_binding( - self.session, - *model_inputs, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "input_features": input_features, + "attention_mask": attention_mask, + } - io_binding.synchronize_inputs() - self.session.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.parent_model.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.session, model_inputs) - last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) - else: - if use_torch: - onnx_inputs = {"input_features": input_features.cpu().detach().numpy()} - if "attention_mask" in self.input_names: - onnx_inputs["attention_mask"] = attention_mask.cpu().detach().numpy() + if self.device.type == "cpu": + self.session.run_with_iobinding(io_binding) else: - onnx_inputs = {"input_features": input_features} - if "attention_mask" in self.input_names: - onnx_inputs["attention_mask"] = attention_mask + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() - # TODO: Replace with a better solution - # attention_mask is exported with int64 datatype and tokenizer produces int32 input - # for speech2text model. Hence, the input is type casted for inference. - if "attention_mask" in self.input_names: - if self.session.get_inputs()[1].type == "tensor(int64)": - onnx_inputs["attention_mask"] = onnx_inputs["attention_mask"].astype(np.int64) - - outputs = self.session.run(None, onnx_inputs) + last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) + else: + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) + onnx_outputs = self.session.run(None, onnx_inputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) - last_hidden_state = outputs[self.output_names["last_hidden_state"]] - if use_torch: - last_hidden_state = torch.from_numpy(last_hidden_state).to(self.device) + last_hidden_state = model_outputs["last_hidden_state"] return BaseModelOutput(last_hidden_state=last_hidden_state) @@ -422,60 +406,30 @@ def forward( use_torch = isinstance(pixel_values, torch.Tensor) self.parent_model.raise_on_numpy_input_io_binding(use_torch) - if self.parent_model.device.type == "cuda" and self.parent_model.use_io_binding: - known_output_shapes = self.compute_encoder_known_output_shapes(pixel_values) + model_inputs = { + "pixel_values": pixel_values, + } - io_binding, output_shapes, output_buffers = self.parent_model._prepare_io_binding( - self.session, - pixel_values, - known_output_shapes=known_output_shapes, - ordered_input_names=self._ordered_input_names, - ) + if self.parent_model.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.session, model_inputs) - io_binding.synchronize_inputs() - self.session.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.device.type == "cpu": + self.session.run_with_iobinding(io_binding) + else: + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) else: - if use_torch: - onnx_inputs = {"pixel_values": pixel_values.cpu().detach().numpy()} - else: - onnx_inputs = {"pixel_values": pixel_values} - - outputs = self.session.run(None, onnx_inputs) + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) + onnx_outputs = self.session.run(None, onnx_inputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) - last_hidden_state = outputs[self.output_names["last_hidden_state"]] - if use_torch: - last_hidden_state = torch.from_numpy(last_hidden_state).to(self.device) + last_hidden_state = model_outputs["last_hidden_state"] return BaseModelOutput(last_hidden_state=last_hidden_state) - def compute_encoder_known_output_shapes(self, pixel_values: torch.FloatTensor) -> Dict[str, List[int]]: - if self.normalized_config.config.model_type == "donut-swin": - # TODO: kind of weird to export to ONNX with dynamic output shape if it is in fact static... - encoder_sequence_length = ( - self.normalized_config.config.image_size[0] - * self.normalized_config.config.image_size[1] - // self.normalized_config.config.hidden_size - ) - elif self.normalized_config.config.model_type in ["vit", "deit"]: - return None - else: - raise ValueError( - f"Unsupported encoder model type {self.normalized_config.config.model_type} for ORTForVisionSeq2Seq with IOBinding." - "Currently supported models are vit, donut-swin and deit." - "Please submit a PR to add support for this model type." - ) - - return { - "last_hidden_state": [ - pixel_values.shape[0], # batch size - encoder_sequence_length, - self.normalized_config.config.hidden_size, - ] - } - class ORTEncoderForPix2Struct(ORTEncoder): """ @@ -496,41 +450,28 @@ def forward( use_torch = isinstance(flattened_patches, torch.Tensor) self.parent_model.raise_on_numpy_input_io_binding(use_torch) - if self.parent_model.device.type == "cuda" and self.parent_model.use_io_binding: - model_inputs = ( - [flattened_patches, attention_mask] if "attention_mask" in self.input_names else [flattened_patches] - ) - io_binding, output_shapes, output_buffers = self.parent_model._prepare_io_binding( - self.session, - *model_inputs, - ordered_input_names=self._ordered_input_names, - ) + model_inputs = { + "flattened_patches": flattened_patches, + "attention_mask": attention_mask, + } - io_binding.synchronize_inputs() - self.session.run_with_iobinding(io_binding) - io_binding.synchronize_outputs() + if self.parent_model.use_io_binding: + io_binding, output_shapes, output_buffers = self._prepare_io_binding(self.session, model_inputs) - last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) - else: - if use_torch: - onnx_inputs = {"flattened_patches": flattened_patches.cpu().detach().numpy()} - if "attention_mask" in self.input_names: - onnx_inputs["attention_mask"] = attention_mask.cpu().detach().numpy() + if self.device.type == "cpu": + self.session.run_with_iobinding(io_binding) else: - onnx_inputs = {"flattened_patches": flattened_patches} - if "attention_mask" in self.input_names: - onnx_inputs["attention_mask"] = attention_mask - - if "attention_mask" in self.input_names: - if self.session.get_inputs()[1].type == "tensor(int64)": - onnx_inputs["attention_mask"] = onnx_inputs["attention_mask"].astype(np.int64) + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() - outputs = self.session.run(None, onnx_inputs) - - last_hidden_state = outputs[self.output_names["last_hidden_state"]] + last_hidden_state = output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) + else: + onnx_inputs = self._prepare_onnx_inputs(use_torch, model_inputs) + onnx_outputs = self.session.run(None, onnx_inputs) + model_outputs = self._prepare_onnx_outputs(use_torch, onnx_outputs) - if use_torch: - last_hidden_state = torch.from_numpy(last_hidden_state).to(self.device) + last_hidden_state = model_outputs["last_hidden_state"] return BaseModelOutput(last_hidden_state=last_hidden_state) @@ -705,7 +646,7 @@ def show_deprecated_argument(arg_name): generation_config = GenerationConfig.from_model_config(config) self.generation_config = generation_config - if check_if_transformers_greater("4.44.99"): + if is_transformers_version(">=", "4.44.99"): misplaced_generation_parameters = self.config._get_non_default_generation_parameters() if len(misplaced_generation_parameters) > 0: logger.warning( @@ -1164,7 +1105,6 @@ def forward( decoder_input_ids: Optional[torch.LongTensor] = None, encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - labels: Optional[torch.LongTensor] = None, **kwargs, ) -> Seq2SeqLMOutput: # Encode if needed : first prediction pass @@ -1181,7 +1121,6 @@ def forward( past_key_values=past_key_values, encoder_hidden_states=encoder_outputs.last_hidden_state, encoder_attention_mask=attention_mask, - labels=labels, ) return Seq2SeqLMOutput( @@ -1297,7 +1236,6 @@ def forward( decoder_input_ids: Optional[torch.LongTensor] = None, encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - labels: Optional[torch.LongTensor] = None, cache_position: Optional[torch.Tensor] = None, **kwargs, ) -> Seq2SeqLMOutput: @@ -1316,7 +1254,6 @@ def forward( encoder_hidden_states=encoder_outputs.last_hidden_state, encoder_attention_mask=attention_mask, cache_position=cache_position, - labels=labels, ) return Seq2SeqLMOutput( @@ -1477,10 +1414,8 @@ def forward( decoder_input_ids: Optional[torch.LongTensor] = None, encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - labels: Optional[torch.LongTensor] = None, **kwargs, ) -> Seq2SeqLMOutput: - # Encode if needed : first prediction pass if encoder_outputs is None: encoder_outputs = self.encoder(pixel_values=pixel_values) @@ -1489,17 +1424,18 @@ def forward( if past_key_values is None or not self.use_cache or self.use_merged else self.decoder_with_past ) + decoder_outputs = model( input_ids=decoder_input_ids, past_key_values=past_key_values, encoder_hidden_states=encoder_outputs.last_hidden_state, - labels=labels, ) return Seq2SeqLMOutput( - loss=decoder_outputs.get("loss", None), + loss=decoder_outputs.loss, logits=decoder_outputs.logits, past_key_values=decoder_outputs.past_key_values, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, ) def prepare_inputs_for_generation( @@ -1577,42 +1513,33 @@ def forward( decoder_attention_mask: Optional[torch.BoolTensor] = None, encoder_outputs: Optional[Tuple[Tuple[torch.Tensor]]] = None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None, - labels: Optional[torch.LongTensor] = None, **kwargs, ) -> Seq2SeqLMOutput: - # Encode if needed : first prediction pass - # Encode if needed (training, first prediction pass) if encoder_outputs is None: encoder_outputs = self.encoder( flattened_patches=flattened_patches, attention_mask=attention_mask, ) - # TODO: for some reason the attention_mask for pix2struct is a float in transformers and not an int64. This messes up with the exporter - # hardcodes int64 input dtype for the attention mask. This workaround is quite ugly, it should be fixed rather in the ONNX exporter. - if isinstance(attention_mask, torch.Tensor): - attention_mask = attention_mask.to(torch.int64) - else: - attention_mask = attention_mask.astype(np.int64) - model = ( self.decoder - if past_key_values is None or not self.use_cache or self.use_merged + if self.use_merged or not self.use_cache or past_key_values is None else self.decoder_with_past ) + decoder_outputs = model( input_ids=decoder_input_ids, decoder_attention_mask=decoder_attention_mask, past_key_values=past_key_values, encoder_hidden_states=encoder_outputs.last_hidden_state, encoder_attention_mask=attention_mask, - labels=labels, ) return Seq2SeqLMOutput( - loss=decoder_outputs.get("loss", None), + loss=decoder_outputs.loss, logits=decoder_outputs.logits, past_key_values=decoder_outputs.past_key_values, + encoder_last_hidden_state=encoder_outputs.last_hidden_state, ) def prepare_inputs_for_generation( diff --git a/optimum/onnxruntime/quantization.py b/optimum/onnxruntime/quantization.py index 056123f8d8..054a2310a6 100644 --- a/optimum/onnxruntime/quantization.py +++ b/optimum/onnxruntime/quantization.py @@ -21,7 +21,6 @@ from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union import onnx -from datasets import Dataset, load_dataset from packaging.version import Version, parse from transformers import AutoConfig @@ -29,6 +28,7 @@ from onnxruntime.quantization import CalibrationDataReader, QuantFormat, QuantizationMode, QuantType from onnxruntime.quantization.onnx_quantizer import ONNXQuantizer from onnxruntime.quantization.qdq_quantizer import QDQQuantizer +from optimum.utils.import_utils import requires_backends from ..quantization_base import OptimumQuantizer from ..utils.save_utils import maybe_save_preprocessors @@ -40,6 +40,7 @@ if TYPE_CHECKING: + from datasets import Dataset from transformers import PretrainedConfig LOGGER = logging.getLogger(__name__) @@ -48,7 +49,7 @@ class ORTCalibrationDataReader(CalibrationDataReader): __slots__ = ["batch_size", "dataset", "_dataset_iter"] - def __init__(self, dataset: Dataset, batch_size: int = 1): + def __init__(self, dataset: "Dataset", batch_size: int = 1): if dataset is None: raise ValueError("Provided dataset is None.") @@ -100,7 +101,7 @@ def __init__(self, onnx_model_path: Path, config: Optional["PretrainedConfig"] = if self.config is None: try: self.config = AutoConfig.from_pretrained(self.onnx_model_path.parent) - except OSError: + except (OSError, ValueError): LOGGER.warning( f"Could not load the config for {self.onnx_model_path} automatically, this might make " "the quantized model harder to use because it will not be able to be loaded by an ORTModel without " @@ -134,6 +135,7 @@ def from_pretrained( model_or_path = Path(model_or_path) path = None + config = None if isinstance(model_or_path, ORTModelForConditionalGeneration): raise NotImplementedError(ort_quantizer_error_message) elif isinstance(model_or_path, Path) and file_name is None: @@ -147,17 +149,17 @@ def from_pretrained( file_name = onnx_files[0].name if isinstance(model_or_path, ORTModel): - if path is None: - path = Path(model_or_path.model._model_path) + path = Path(model_or_path.model._model_path) + config = model_or_path.config elif os.path.isdir(model_or_path): path = Path(model_or_path) / file_name else: raise ValueError(f"Unable to load model from {model_or_path}.") - return cls(path) + return cls(path, config=config) def fit( self, - dataset: Dataset, + dataset: "Dataset", calibration_config: CalibrationConfig, onnx_augmented_model_name: Union[str, Path] = "augmented_model.onnx", operators_to_quantize: Optional[List[str]] = None, @@ -211,7 +213,7 @@ def fit( def partial_fit( self, - dataset: Dataset, + dataset: "Dataset", calibration_config: CalibrationConfig, onnx_augmented_model_name: Union[str, Path] = "augmented_model.onnx", operators_to_quantize: Optional[List[str]] = None, @@ -427,7 +429,7 @@ def get_calibration_dataset( seed: int = 2016, use_auth_token: Optional[Union[bool, str]] = None, token: Optional[Union[bool, str]] = None, - ) -> Dataset: + ) -> "Dataset": """ Creates the calibration `datasets.Dataset` to use for the post-training static quantization calibration step. @@ -473,6 +475,10 @@ def get_calibration_dataset( "provided." ) + requires_backends(self, ["datasets"]) + + from datasets import load_dataset + calib_dataset = load_dataset( dataset_name, name=dataset_config_name, @@ -491,7 +497,7 @@ def get_calibration_dataset( return self.clean_calibration_dataset(processed_calib_dataset) - def clean_calibration_dataset(self, dataset: Dataset) -> Dataset: + def clean_calibration_dataset(self, dataset: "Dataset") -> "Dataset": model = onnx.load(self.onnx_model_path) model_inputs = {input.name for input in model.graph.input} ignored_columns = list(set(dataset.column_names) - model_inputs) diff --git a/optimum/onnxruntime/runs/calibrator.py b/optimum/onnxruntime/runs/calibrator.py index c493a94374..bfdcd64d92 100644 --- a/optimum/onnxruntime/runs/calibrator.py +++ b/optimum/onnxruntime/runs/calibrator.py @@ -1,6 +1,4 @@ -from typing import Dict, List - -from datasets import Dataset +from typing import TYPE_CHECKING, Dict, List from ...runs_base import Calibrator from .. import ORTQuantizer @@ -9,10 +7,14 @@ from ..preprocessors.passes import ExcludeGeLUNodes, ExcludeLayerNormNodes, ExcludeNodeAfter, ExcludeNodeFollowedBy +if TYPE_CHECKING: + from datasets import Dataset + + class OnnxRuntimeCalibrator(Calibrator): def __init__( self, - calibration_dataset: Dataset, + calibration_dataset: "Dataset", quantizer: ORTQuantizer, model_path: str, qconfig: QuantizationConfig, diff --git a/optimum/onnxruntime/trainer.py b/optimum/onnxruntime/trainer.py index 66273cbcf9..47a98e19c8 100644 --- a/optimum/onnxruntime/trainer.py +++ b/optimum/onnxruntime/trainer.py @@ -14,6 +14,7 @@ """ The ORTTrainer class, to easily train a 🤗 Transformers from scratch or finetune it on a new task with ONNX Runtime. """ + import functools import math import os @@ -27,8 +28,8 @@ # Integrations must be imported before ML frameworks: # isort: off +import safetensors from transformers.integrations import hp_params - from transformers.utils import is_accelerate_available from packaging import version @@ -58,7 +59,7 @@ from transformers.modeling_utils import PreTrainedModel, unwrap_model from transformers.tokenization_utils_base import PreTrainedTokenizerBase from transformers.trainer import Trainer -from transformers.trainer_callback import TrainerCallback, TrainerState +from transformers.trainer_callback import ExportableState, TrainerCallback, TrainerState from transformers.trainer_pt_utils import ( get_model_param_count, get_module_class_from_name, @@ -77,13 +78,15 @@ ) from transformers.training_args import ParallelMode from transformers.utils import ( + SAFE_WEIGHTS_NAME, + WEIGHTS_NAME, is_apex_available, is_sagemaker_dp_enabled, is_sagemaker_mp_enabled, ) from ..utils import logging -from ..utils.import_utils import check_if_transformers_greater +from ..utils.import_utils import is_transformers_version from .training_args import ORTOptimizerNames, ORTTrainingArguments from .utils import ( is_onnxruntime_training_available, @@ -93,7 +96,7 @@ if is_apex_available(): from apex import amp -if check_if_transformers_greater("4.33"): +if is_transformers_version(">=", "4.33"): from transformers.integrations.deepspeed import ( deepspeed_init, deepspeed_load_checkpoint, @@ -102,7 +105,7 @@ else: from transformers.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_zero3_enabled -if check_if_transformers_greater("4.39"): +if is_transformers_version(">=", "4.39"): from transformers.utils import is_torch_xla_available as is_torch_tpu_xla_available if is_torch_tpu_xla_available(): @@ -119,11 +122,12 @@ # Name of the files used for checkpointing TRAINER_STATE_NAME = "trainer_state.json" +TRAINING_ARGS_NAME = "training_args.bin" logger = logging.get_logger(__name__) -class ModuleWithLoss(nn.Module): +class ModuleWithLoss(PreTrainedModel): def __init__(self, model, args, label_smoother): super().__init__() self._original_model = model @@ -131,11 +135,11 @@ def __init__(self, model, args, label_smoother): # Label smoothing self.label_smoother = label_smoother - def forward(self, inputs: Dict[str, Union[torch.Tensor, Any]], return_outputs): + def forward(self, inputs: Dict[str, Union[torch.Tensor, Any]], return_outputs, num_items_in_batch): # The compute_model_plus_loss_internal is assigned once the class is instantiated. # It should have same signature as Trainer.compute_loss(). # We do this to avoid potential un-synced states if we duplicated compute loss codes . - return self.compute_model_plus_loss_internal(self._original_model, inputs, return_outputs) + return self.compute_model_plus_loss_internal(self._original_model, inputs, return_outputs, num_items_in_batch) @property def module(self): @@ -291,14 +295,14 @@ def _set_signature_columns_if_needed(self): # Labels may be named label or label_ids, the default data collator handles that. self._signature_columns += list(set(["label", "label_ids"] + self.label_names)) - def compute_loss(self, model_with_loss, inputs, return_outputs=False): + def compute_loss(self, model_with_loss, inputs, return_outputs=False, num_items_in_batch=None): # Run model forward + loss compute. if isinstance(self.model, ModuleWithLoss): # ORTModule Does not support the BatchEncoding Type so we have to convert to a dict. dict_inputs = dict(inputs.items()) - return model_with_loss(dict_inputs, return_outputs) + return model_with_loss(dict_inputs, return_outputs, num_items_in_batch) else: - return super().compute_loss(model_with_loss, inputs, return_outputs) + return super().compute_loss(model_with_loss, inputs, return_outputs, num_items_in_batch) def train( self, @@ -508,8 +512,13 @@ def _inner_training_loop( if not delay_optimizer_creation: self.create_optimizer_and_scheduler(num_training_steps=max_steps) - self.state = TrainerState() + self.state = TrainerState( + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ] + ) self.state.is_hyper_param_search = trial is not None + self.state.train_batch_size = self._train_batch_size # Compute absolute values for logging, eval, and save if given as ratio if args.logging_steps is not None: @@ -798,12 +807,16 @@ def get_dataloader_sampler(dataloader): self.lr_scheduler.step() model.zero_grad() - grad_norm: Optional[float] = None self.state.global_step += 1 self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch self.control = self.callback_handler.on_step_end(args, self.state, self.control) - self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval) + if is_transformers_version(">=", "4.47.0"): + self._maybe_log_save_evaluate( + tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time + ) + else: + self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval) else: self.control = self.callback_handler.on_substep_end(args, self.state, self.control) @@ -818,8 +831,13 @@ def get_dataloader_sampler(dataloader): self.control.should_training_stop = True self.control = self.callback_handler.on_epoch_end(args, self.state, self.control) - self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval) + if is_transformers_version(">=", "4.47.0"): + self._maybe_log_save_evaluate( + tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time + ) + else: + self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval) if DebugOption.TPU_METRICS_DEBUG in self.args.debug: logger.warning( "You enabled PyTorch/XLA debug metrics which is not supported by ONNX " @@ -1072,3 +1090,39 @@ def get_ort_optimizer_cls_and_kwargs(args: ORTTrainingArguments) -> Tuple[Any, A else: raise ValueError(f"ORTTrainer cannot instantiate unsupported optimizer: {args.optim}") return optimizer_cls, optimizer_kwargs + + def _save(self, output_dir: Optional[str] = None, state_dict=None): + # If we are executing this function, we are the process zero, so we don't check for that. + output_dir = output_dir if output_dir is not None else self.args.output_dir + os.makedirs(output_dir, exist_ok=True) + logger.info(f"Saving model checkpoint to {output_dir}") + + supported_classes = (PreTrainedModel,) + # Save a trained model and configuration using `save_pretrained()`. + # They can then be reloaded using `from_pretrained()` + if not isinstance(self.model, supported_classes): + if state_dict is None: + state_dict = self.model.state_dict() + + if isinstance(self.accelerator.unwrap_model(self.model), supported_classes): + self.accelerator.unwrap_model(self.model).save_pretrained( + output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors + ) + else: + logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") + if self.args.save_safetensors: + safetensors.torch.save_model( + self.model, os.path.join(output_dir, SAFE_WEIGHTS_NAME), metadata={"format": "pt"} + ) + else: + torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) + else: + self.model.save_pretrained( + output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors + ) + + if self.processing_class is not None: + self.processing_class.save_pretrained(output_dir) + + # Good practice: save your training arguments together with the trained model + torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) diff --git a/optimum/onnxruntime/trainer_seq2seq.py b/optimum/onnxruntime/trainer_seq2seq.py index 1565ffa6ac..a76374a5ec 100644 --- a/optimum/onnxruntime/trainer_seq2seq.py +++ b/optimum/onnxruntime/trainer_seq2seq.py @@ -22,7 +22,7 @@ from transformers.trainer_utils import PredictionOutput from transformers.utils import is_accelerate_available, logging -from ..utils.import_utils import check_if_transformers_greater +from ..utils.import_utils import is_transformers_version from .trainer import ORTTrainer @@ -33,7 +33,7 @@ "The package `accelerate` is required to use the ORTTrainer. Please install it following https://huggingface.co/docs/accelerate/basic_tutorials/install." ) -if check_if_transformers_greater("4.33"): +if is_transformers_version(">=", "4.33"): from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled else: from transformers.deepspeed import is_deepspeed_zero3_enabled diff --git a/optimum/onnxruntime/training_args.py b/optimum/onnxruntime/training_args.py index 6135abc137..6eb2bb4904 100644 --- a/optimum/onnxruntime/training_args.py +++ b/optimum/onnxruntime/training_args.py @@ -44,13 +44,13 @@ ) from transformers.utils.generic import strtobool -from ..utils.import_utils import check_if_transformers_greater +from ..utils.import_utils import is_transformers_version if is_torch_available(): import torch -if is_accelerate_available() and check_if_transformers_greater("4.38.0"): +if is_accelerate_available() and is_transformers_version(">=", "4.38.0"): from transformers.trainer_pt_utils import AcceleratorConfig @@ -481,7 +481,7 @@ def __post_init__(self): os.environ[f"{prefix}SYNC_MODULE_STATES"] = self.fsdp_config.get("sync_module_states", "true") os.environ[f"{prefix}USE_ORIG_PARAMS"] = self.fsdp_config.get("use_orig_params", "false") - if is_accelerate_available() and check_if_transformers_greater("4.38.0"): + if is_accelerate_available() and is_transformers_version(">=", "4.38.0"): if not isinstance(self.accelerator_config, (AcceleratorConfig)): if self.accelerator_config is None: self.accelerator_config = AcceleratorConfig() diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py index 128e2406f1..be395927cc 100644 --- a/optimum/onnxruntime/utils.py +++ b/optimum/onnxruntime/utils.py @@ -114,7 +114,7 @@ class ORTConfigManager: "bart": "bart", "bert": "bert", "big-bird": "bert", - # "bigbird-pegasus": None, # bug in `fusion_skiplayernorm.py` + "bigbird-pegasus": "bart", "blenderbot": "bert", "bloom": "gpt2", "camembert": "bert", @@ -128,7 +128,7 @@ class ORTConfigManager: "gpt-neo": "gpt2", "gpt-neox": "gpt2", "gptj": "gpt2", - # longt5 with O4 results in segmentation fault + "granite": "gpt2", "longt5": "bert", "llama": "gpt2", "marian": "bart", @@ -177,6 +177,7 @@ def check_optimization_supported_model(cls, model_type: str, optimization_config "clip", "vit", "swin", + "swinv2", ] model_type = model_type.replace("_", "-") if (model_type not in cls._conf) or (cls._conf[model_type] not in supported_model_types_for_optimization): diff --git a/optimum/pipelines/pipelines_base.py b/optimum/pipelines/pipelines_base.py index 7690143f13..0016c73ff0 100644 --- a/optimum/pipelines/pipelines_base.py +++ b/optimum/pipelines/pipelines_base.py @@ -46,7 +46,7 @@ from transformers.pipelines import infer_framework_load_model from ..bettertransformer import BetterTransformer -from ..utils import check_if_transformers_greater, is_onnxruntime_available +from ..utils import is_onnxruntime_available, is_transformers_version from ..utils.file_utils import find_files_matching_pattern @@ -189,7 +189,7 @@ def load_bettertransformer( if model_kwargs is None: # the argument was first introduced in 4.36.0 but most models didn't have an sdpa implementation then # see https://github.com/huggingface/transformers/blob/v4.36.0/src/transformers/modeling_utils.py#L1258 - if check_if_transformers_greater("4.36.0"): + if is_transformers_version(">=", "4.36.0"): model_kwargs = {"attn_implementation": "eager"} else: model_kwargs = {} diff --git a/optimum/runs_base.py b/optimum/runs_base.py index 3a1d164c60..dadd445818 100644 --- a/optimum/runs_base.py +++ b/optimum/runs_base.py @@ -2,13 +2,12 @@ import subprocess from contextlib import contextmanager from time import perf_counter_ns -from typing import Set +from typing import TYPE_CHECKING, Set import numpy as np import optuna import torch import transformers -from datasets import Dataset from tqdm import trange from . import version as optimum_version @@ -21,6 +20,9 @@ from .utils.runs import RunConfig, cpu_info_command +if TYPE_CHECKING: + from datasets import Dataset + os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -34,7 +36,7 @@ def get_autoclass_name(task): class Calibrator: def __init__( - self, calibration_dataset: Dataset, quantizer, model_path, qconfig, calibration_params, node_exclusion + self, calibration_dataset: "Dataset", quantizer, model_path, qconfig, calibration_params, node_exclusion ): self.calibration_dataset = calibration_dataset self.quantizer = quantizer diff --git a/optimum/subpackages.py b/optimum/subpackages.py index 8729581521..24eca2a139 100644 --- a/optimum/subpackages.py +++ b/optimum/subpackages.py @@ -48,6 +48,8 @@ def load_namespace_modules(namespace: str, module: str): dist_name = dist.metadata["Name"] if not dist_name.startswith(f"{namespace}-"): continue + if dist_name == f"{namespace}-benchmark": + continue package_import_name = dist_name.replace("-", ".") module_import_name = f"{package_import_name}.{module}" if module_import_name in sys.modules: diff --git a/optimum/utils/__init__.py b/optimum/utils/__init__.py index db7d1f6975..b4097f0f80 100644 --- a/optimum/utils/__init__.py +++ b/optimum/utils/__init__.py @@ -16,7 +16,9 @@ from .constant import ( CONFIG_NAME, DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER, + DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER, DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER, + DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER, DIFFUSION_MODEL_UNET_SUBFOLDER, DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER, DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER, @@ -33,15 +35,22 @@ check_if_transformers_greater, is_accelerate_available, is_auto_gptq_available, + is_datasets_available, is_diffusers_available, + is_diffusers_version, + is_gptqmodel_available, is_onnx_available, is_onnxruntime_available, is_pydantic_available, is_sentence_transformers_available, + is_tf_available, is_timm_available, + is_torch_available, is_torch_onnx_support_available, + is_torch_version, + is_transformers_available, + is_transformers_version, require_numpy_strictly_lower, - torch_version, ) from .input_generators import ( DEFAULT_DUMMY_SHAPES, @@ -50,12 +59,16 @@ DummyAudioInputGenerator, DummyBboxInputGenerator, DummyCodegenDecoderTextInputGenerator, + DummyDecisionTransformerInputGenerator, DummyDecoderTextInputGenerator, DummyEncodecInputGenerator, + DummyFluxTransformerTextInputGenerator, + DummyFluxTransformerVisionInputGenerator, DummyInputGenerator, DummyIntGenerator, DummyLabelsGenerator, DummyPastKeyValuesGenerator, + DummyPatchTSTInputGenerator, DummyPix2StructInputGenerator, DummyPointsGenerator, DummySeq2SeqDecoderTextInputGenerator, @@ -63,6 +76,9 @@ DummySpeechT5InputGenerator, DummyTextInputGenerator, DummyTimestepInputGenerator, + DummyTransformerTextInputGenerator, + DummyTransformerTimestepInputGenerator, + DummyTransformerVisionInputGenerator, DummyVisionEmbeddingsGenerator, DummyVisionEncoderDecoderPastKeyValuesGenerator, DummyVisionInputGenerator, @@ -70,6 +86,8 @@ FalconDummyPastKeyValuesGenerator, GemmaDummyPastKeyValuesGenerator, GPTBigCodeDummyPastKeyValuesGenerator, + LongformerDummyTextInputGenerator, + MCTCTDummyAudioInputGenerator, MistralDummyPastKeyValuesGenerator, MultiQueryPastKeyValuesGenerator, ) @@ -82,5 +100,6 @@ NormalizedTextAndVisionConfig, NormalizedTextConfig, NormalizedTextConfigWithGQA, + NormalizedTimeSeriesForecastingConfig, NormalizedVisionConfig, ) diff --git a/optimum/utils/constant.py b/optimum/utils/constant.py index 4497b5246d..eb7a67e9ec 100644 --- a/optimum/utils/constant.py +++ b/optimum/utils/constant.py @@ -15,8 +15,10 @@ CONFIG_NAME = "config.json" DIFFUSION_MODEL_UNET_SUBFOLDER = "unet" -DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER = "text_encoder" +DIFFUSION_MODEL_TRANSFORMER_SUBFOLDER = "transformer" DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER = "vae_decoder" DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER = "vae_encoder" +DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER = "text_encoder" DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER = "text_encoder_2" +DIFFUSION_MODEL_TEXT_ENCODER_3_SUBFOLDER = "text_encoder_3" ONNX_WEIGHTS_NAME = "model.onnx" diff --git a/optimum/utils/dummy_diffusers_objects.py b/optimum/utils/dummy_diffusers_objects.py index 35d1ffe9fc..ff8b587e19 100644 --- a/optimum/utils/dummy_diffusers_objects.py +++ b/optimum/utils/dummy_diffusers_objects.py @@ -15,6 +15,50 @@ from .import_utils import DummyObject, requires_backends +class ORTDiffusionPipeline(metaclass=DummyObject): + _backends = ["diffusers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["diffusers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["diffusers"]) + + +class ORTPipelineForText2Image(metaclass=DummyObject): + _backends = ["diffusers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["diffusers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["diffusers"]) + + +class ORTPipelineForImage2Image(metaclass=DummyObject): + _backends = ["diffusers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["diffusers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["diffusers"]) + + +class ORTPipelineForInpainting(metaclass=DummyObject): + _backends = ["diffusers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["diffusers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["diffusers"]) + + class ORTStableDiffusionPipeline(metaclass=DummyObject): _backends = ["diffusers"] @@ -70,6 +114,17 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["diffusers"]) +class ORTStableDiffusionXLInpaintPipeline(metaclass=DummyObject): + _backends = ["diffusers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["diffusers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["diffusers"]) + + class ORTLatentConsistencyModelPipeline(metaclass=DummyObject): _backends = ["diffusers"] @@ -81,7 +136,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["diffusers"]) -class ORTDiffusionPipeline(metaclass=DummyObject): +class ORTLatentConsistencyModelImg2ImgPipeline(metaclass=DummyObject): _backends = ["diffusers"] def __init__(self, *args, **kwargs): @@ -92,7 +147,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["diffusers"]) -class ORTPipelineForText2Image(metaclass=DummyObject): +class ORTStableDiffusion3Pipeline(metaclass=DummyObject): _backends = ["diffusers"] def __init__(self, *args, **kwargs): @@ -103,7 +158,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["diffusers"]) -class ORTPipelineForImage2Image(metaclass=DummyObject): +class ORTStableDiffusion3Img2ImgPipeline(metaclass=DummyObject): _backends = ["diffusers"] def __init__(self, *args, **kwargs): @@ -114,7 +169,18 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["diffusers"]) -class ORTPipelineForInpainting(metaclass=DummyObject): +class ORTStableDiffusion3InpaintPipeline(metaclass=DummyObject): + _backends = ["diffusers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["diffusers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["diffusers"]) + + +class ORTFluxPipeline(metaclass=DummyObject): _backends = ["diffusers"] def __init__(self, *args, **kwargs): diff --git a/optimum/utils/import_utils.py b/optimum/utils/import_utils.py index 35a6294ab5..8da1df5fac 100644 --- a/optimum/utils/import_utils.py +++ b/optimum/utils/import_utils.py @@ -13,82 +13,188 @@ # limitations under the License. """Import utilities.""" +import importlib.metadata import importlib.util -import inspect -import sys +import operator as op from collections import OrderedDict from contextlib import contextmanager -from typing import Tuple, Union +from logging import getLogger +from typing import List, Optional, Tuple, Union import numpy as np from packaging import version -from transformers.utils import is_torch_available -def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[Tuple[bool, str], bool]: - # Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version - package_exists = importlib.util.find_spec(pkg_name) is not None - package_version = "N/A" - if package_exists: - try: - package_version = importlib.metadata.version(pkg_name) - package_exists = True - except importlib.metadata.PackageNotFoundError: - package_exists = False - if return_version: - return package_exists, package_version - else: - return package_exists - - -# The package importlib_metadata is in a different place, depending on the python version. -if sys.version_info < (3, 8): - import importlib_metadata -else: - import importlib.metadata as importlib_metadata - +logger = getLogger(__name__) TORCH_MINIMUM_VERSION = version.parse("1.11.0") TRANSFORMERS_MINIMUM_VERSION = version.parse("4.25.0") DIFFUSERS_MINIMUM_VERSION = version.parse("0.22.0") AUTOGPTQ_MINIMUM_VERSION = version.parse("0.4.99") # Allows 0.5.0.dev0 - +GPTQMODEL_MINIMUM_VERSION = version.parse("1.6.0") # This is the minimal required version to support some ONNX Runtime features ORT_QUANTIZE_MINIMUM_VERSION = version.parse("1.4.0") +STR_OPERATION_TO_FUNC = {">": op.gt, ">=": op.ge, "==": op.eq, "!=": op.ne, "<=": op.le, "<": op.lt} -_onnx_available = _is_package_available("onnx") -# importlib.metadata.version seem to not be robust with the ONNX Runtime extensions (`onnxruntime-gpu`, etc.) -_onnxruntime_available = importlib.util.find_spec("onnxruntime") is not None +def _is_package_available( + pkg_name: str, + return_version: bool = False, + pkg_distributions: Optional[List[str]] = None, +) -> Union[Tuple[bool, str], bool]: + """ + Check if a package is available in the current environment and not just an importable module by checking its version. + Optionally return the version of the package. + + Args: + pkg_name (str): The name of the package to check. + return_version (bool): Whether to return the version of the package. + pkg_distributions (Optional[List[str]]): A list of package distributions (e.g. "package-name", "package-name-gpu", etc.) to check for the package. + + Returns: + Union[Tuple[bool, str], bool]: A tuple of the package availability and the version of the package if `return_version` is `True`. + """ + package_exists = importlib.util.find_spec(pkg_name) is not None + package_version = "N/A" + + if pkg_distributions is None: + pkg_distributions = [pkg_name] + else: + pkg_distributions.append(pkg_name) + + if package_exists: + for pkg in pkg_distributions: + try: + package_version = importlib.metadata.version(pkg) + package_exists = True + break + except importlib.metadata.PackageNotFoundError: + package_exists = False + pass + + if return_version: + return package_exists, package_version + else: + return package_exists + + +_onnx_available = _is_package_available("onnx") _pydantic_available = _is_package_available("pydantic") _accelerate_available = _is_package_available("accelerate") -_diffusers_available = _is_package_available("diffusers") _auto_gptq_available = _is_package_available("auto_gptq") +_gptqmodel_available = _is_package_available("gptqmodel") _timm_available = _is_package_available("timm") _sentence_transformers_available = _is_package_available("sentence_transformers") +_datasets_available = _is_package_available("datasets") +_diffusers_available, _diffusers_version = _is_package_available("diffusers", return_version=True) +_transformers_available, _transformers_version = _is_package_available("transformers", return_version=True) +_torch_available, _torch_version = _is_package_available("torch", return_version=True) +_onnxruntime_available, _onnxruntime_version = _is_package_available( + "onnxruntime", + return_version=True, + pkg_distributions=[ + "onnxruntime-gpu", + "onnxruntime-rocm", + "onnxruntime-training", + # list in https://github.com/microsoft/onnxruntime/blob/main/setup.py#L56C1-L98C91 + "onnxruntime-training-rocm", + "onnxruntime-training-cpu", + "onnxruntime-openvino", + "onnxruntime-vitisai", + "onnxruntime-armnn", + "onnxruntime-cann", + "onnxruntime-dnnl", + "onnxruntime-acl", + "onnxruntime-tvm", + "onnxruntime-qnn", + "onnxruntime-migraphx", + "ort-migraphx-nightly", + "ort-rocm-nightly", + ], +) +_tf_available, _tf_version = _is_package_available( + "tensorflow", + return_version=True, + pkg_distributions=[ + "tensorflow", + "tensorflow-cpu", + "tensorflow-gpu", + "tensorflow-rocm", + "tensorflow-macos", + "tensorflow-aarch64", + "tf-nightly", + "tf-nightly-cpu", + "tf-nightly-gpu", + "tf-nightly-rocm", + "tf-nightly-macos", + "intel-tensorflow", + "intel-tensorflow-avx512", + ], +) -torch_version = None -if is_torch_available(): - torch_version = version.parse(importlib_metadata.version("torch")) +if _tf_available and version.parse(_tf_version) < version.parse("2"): + logger.warning( + "TensorFlow 2.0 or higher is required to use the TensorFlow backend. " + "Please install the latest version of TensorFlow, or switch to another backend." + ) + _tf_available = False -_is_torch_onnx_support_available = is_torch_available() and ( - TORCH_MINIMUM_VERSION.major, - TORCH_MINIMUM_VERSION.minor, -) <= ( - torch_version.major, - torch_version.minor, -) +# This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319 +def compare_versions(library_or_version: Union[str, version.Version], operation: str, requirement_version: str): + """ + Compare a library version to some requirement using a given operation. + + Arguments: + library_or_version (`str` or `packaging.version.Version`): + A library name or a version to check. + operation (`str`): + A string representation of an operator, such as `">"` or `"<="`. + requirement_version (`str`): + The version to compare the library version against + """ + if operation not in STR_OPERATION_TO_FUNC.keys(): + raise ValueError(f"`operation` must be one of {list(STR_OPERATION_TO_FUNC.keys())}, received {operation}") + operation = STR_OPERATION_TO_FUNC[operation] + if isinstance(library_or_version, str): + library_or_version = version.parse(importlib.metadata.version(library_or_version)) + return operation(library_or_version, version.parse(requirement_version)) -_diffusers_version = None -if _diffusers_available: - try: - _diffusers_version = importlib_metadata.version("diffusers") - except importlib_metadata.PackageNotFoundError: - _diffusers_available = False + +def is_transformers_version(operation: str, reference_version: str): + """ + Compare the current Transformers version to a given reference with an operation. + """ + if not _transformers_available: + return False + return compare_versions(version.parse(_transformers_version), operation, reference_version) + + +def is_diffusers_version(operation: str, reference_version: str): + """ + Compare the current diffusers version to a given reference with an operation. + """ + if not _diffusers_available: + return False + return compare_versions(version.parse(_diffusers_version), operation, reference_version) + + +def is_torch_version(operation: str, reference_version: str): + """ + Compare the current torch version to a given reference with an operation. + """ + if not _torch_available: + return False + + import torch + + return compare_versions(version.parse(version.parse(torch.__version__).base_version), operation, reference_version) + + +_is_torch_onnx_support_available = _torch_available and is_torch_version(">=", TORCH_MINIMUM_VERSION.base_version) def is_torch_onnx_support_available(): @@ -100,14 +206,6 @@ def is_onnx_available(): def is_onnxruntime_available(): - try: - # Try to import the source file of onnxruntime - if you run the tests from `tests` the function gets - # confused since there a folder named `onnxruntime` in `tests`. Therefore, `_onnxruntime_available` - # will be set to `True` even if not installed. - mod = importlib.import_module("onnxruntime") - inspect.getsourcefile(mod) - except Exception: - return False return _onnxruntime_available @@ -131,14 +229,41 @@ def is_sentence_transformers_available(): return _sentence_transformers_available +def is_datasets_available(): + return _datasets_available + + +def is_transformers_available(): + return _transformers_available + + +def is_torch_available(): + return _torch_available + + +def is_tf_available(): + return _tf_available + + def is_auto_gptq_available(): if _auto_gptq_available: - version_autogptq = version.parse(importlib_metadata.version("auto_gptq")) - if AUTOGPTQ_MINIMUM_VERSION < version_autogptq: + v = version.parse(importlib.metadata.version("auto_gptq")) + if v >= AUTOGPTQ_MINIMUM_VERSION: return True else: raise ImportError( - f"Found an incompatible version of auto-gptq. Found version {version_autogptq}, but only version above {AUTOGPTQ_MINIMUM_VERSION} are supported" + f"Found an incompatible version of auto-gptq. Found version {v}, but only version >= {AUTOGPTQ_MINIMUM_VERSION} are supported" + ) + + +def is_gptqmodel_available(): + if _gptqmodel_available: + v = version.parse(importlib.metadata.version("gptqmodel")) + if v >= GPTQMODEL_MINIMUM_VERSION: + return True + else: + raise ImportError( + f"Found an incompatible version of gptqmodel. Found version {v}, but only version >= {GPTQMODEL_MINIMUM_VERSION} are supported" ) @@ -159,6 +284,7 @@ def check_if_pytorch_greater(target_version: str, message: str): pass +# TODO : Remove check_if_transformers_greater, check_if_diffusers_greater, check_if_torch_greater def check_if_transformers_greater(target_version: Union[str, version.Version]) -> bool: """ Checks whether the current install of transformers is greater than or equal to the target version. @@ -203,10 +329,10 @@ def check_if_torch_greater(target_version: str) -> bool: Returns: bool: whether the check is True or not. """ - if not is_torch_available(): + if not _torch_available: return False - return torch_version >= version.parse(target_version) + return version.parse(_torch_version) >= version.parse(target_version) @contextmanager @@ -230,21 +356,28 @@ def require_numpy_strictly_lower(package_version: str, message: str): -U transformers`. Please note that you may need to restart your runtime after installation. """ +DATASETS_IMPORT_ERROR = """ +{0} requires the datasets library but it was not found in your environment. You can install it with pip: +`pip install datasets`. Please note that you may need to restart your runtime after installation. +""" + + BACKENDS_MAPPING = OrderedDict( [ ("diffusers", (is_diffusers_available, DIFFUSERS_IMPORT_ERROR)), ( "transformers_431", - (lambda: check_if_transformers_greater("4.31"), "{0} " + TRANSFORMERS_IMPORT_ERROR.format("4.31")), + (lambda: is_transformers_version(">=", "4.31"), "{0} " + TRANSFORMERS_IMPORT_ERROR.format("4.31")), ), ( "transformers_432", - (lambda: check_if_transformers_greater("4.32"), "{0} " + TRANSFORMERS_IMPORT_ERROR.format("4.32")), + (lambda: is_transformers_version(">=", "4.32"), "{0} " + TRANSFORMERS_IMPORT_ERROR.format("4.32")), ), ( "transformers_434", - (lambda: check_if_transformers_greater("4.34"), "{0} " + TRANSFORMERS_IMPORT_ERROR.format("4.34")), + (lambda: is_transformers_version(">=", "4.34"), "{0} " + TRANSFORMERS_IMPORT_ERROR.format("4.34")), ), + ("datasets", (is_datasets_available, DATASETS_IMPORT_ERROR)), ] ) diff --git a/optimum/utils/input_generators.py b/optimum/utils/input_generators.py index dac14a3811..6a265061fd 100644 --- a/optimum/utils/input_generators.py +++ b/optimum/utils/input_generators.py @@ -20,9 +20,8 @@ from typing import Any, List, Optional, Tuple, Union import numpy as np -from transformers.utils import is_tf_available, is_torch_available -from ..utils import check_if_transformers_greater +from ..utils import is_diffusers_version, is_tf_available, is_torch_available, is_transformers_version from .normalized_config import ( NormalizedConfig, NormalizedEncoderDecoderConfig, @@ -36,7 +35,7 @@ import torch if is_tf_available(): - import tensorflow as tf + import tensorflow as tf # type: ignore def check_framework_is_available(func): @@ -384,6 +383,7 @@ class DummyTextInputGenerator(DummyInputGenerator): "input_ids", "attention_mask", "encoder_attention_mask", + "global_attention_mask", "token_type_ids", "position_ids", ) @@ -426,24 +426,47 @@ def __init__( self.padding_side = padding_side self.normalized_config = normalized_config - def generate( - self, - input_name: str, - framework: str = "pt", - int_dtype: str = "int64", - float_dtype: str = "fp32", - ): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): min_value = 0 max_value = 2 if input_name != "input_ids" else self.vocab_size - shape = [self.batch_size, self.sequence_length] + if self.task == "multiple-choice": shape = [self.batch_size, self.num_choices, self.sequence_length] - if "mask" in input_name: + else: + shape = [self.batch_size, self.sequence_length] + + if input_name in ["attention_mask", "encoder_attention_mask"]: return self.random_mask_tensor(shape, padding_side=self.padding_side, framework=framework, dtype=int_dtype) else: return self.random_int_tensor(shape, max_value, min_value=min_value, framework=framework, dtype=int_dtype) +class LongformerDummyTextInputGenerator(DummyTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "input_ids", + "attention_mask", + "token_type_ids", + "global_attention_mask", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "global_attention_mask": + attention_mask = super().generate( + "attention_mask", framework=framework, int_dtype=int_dtype, float_dtype=float_dtype + ) + + if framework == "pt": + global_attention_mask = torch.zeros_like(attention_mask) + elif framework == "tf": + global_attention_mask = tf.zeros_like(attention_mask) + else: + global_attention_mask = np.zeros_like(attention_mask) + + return global_attention_mask + + return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) + + class DummyXPathSeqInputGenerator(DummyTextInputGenerator): """ Generates dummy xpath sequences. @@ -507,6 +530,43 @@ class DummyDecoderTextInputGenerator(DummyTextInputGenerator): ) +class DummyDecisionTransformerInputGenerator(DummyTextInputGenerator): + """ + Generates dummy decision transformer inputs. + """ + + SUPPORTED_INPUT_NAMES = ( + "states", + "actions", + "timesteps", + "returns_to_go", + "attention_mask", + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.act_dim = self.normalized_config.config.act_dim + self.state_dim = self.normalized_config.config.state_dim + self.max_ep_len = self.normalized_config.config.max_ep_len + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "states": + shape = [self.batch_size, self.sequence_length, self.state_dim] + elif input_name == "actions": + shape = [self.batch_size, self.sequence_length, self.act_dim] + elif input_name == "rewards": + shape = [self.batch_size, self.sequence_length, 1] + elif input_name == "returns_to_go": + shape = [self.batch_size, self.sequence_length, 1] + elif input_name == "attention_mask": + shape = [self.batch_size, self.sequence_length] + elif input_name == "timesteps": + shape = [self.batch_size, self.sequence_length] + return self.random_int_tensor(shape=shape, max_value=self.max_ep_len, framework=framework, dtype=int_dtype) + + return self.random_float_tensor(shape, min_value=-2.0, max_value=2.0, framework=framework, dtype=float_dtype) + + class DummySeq2SeqDecoderTextInputGenerator(DummyDecoderTextInputGenerator): SUPPORTED_INPUT_NAMES = ( "decoder_input_ids", @@ -860,23 +920,31 @@ def __init__( ): self.task = task self.vocab_size = normalized_config.vocab_size - self.text_encoder_projection_dim = normalized_config.text_encoder_projection_dim - self.time_ids = 5 if normalized_config.requires_aesthetics_score else 6 + self.text_encoder_projection_dim = getattr(normalized_config, "text_encoder_projection_dim", None) + self.time_ids = 5 if getattr(normalized_config, "requires_aesthetics_score", False) else 6 if random_batch_size_range: low, high = random_batch_size_range self.batch_size = random.randint(low, high) else: self.batch_size = batch_size - self.time_cond_proj_dim = normalized_config.config.time_cond_proj_dim + self.time_cond_proj_dim = getattr(normalized_config.config, "time_cond_proj_dim", None) def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): if input_name == "timestep": - shape = [self.batch_size] - return self.random_int_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=int_dtype) + shape = [] # a scalar with no dimension (it can be int or float depending on the sd architecture) + return self.random_float_tensor(shape, max_value=self.vocab_size, framework=framework, dtype=float_dtype) if input_name == "text_embeds": + if self.text_encoder_projection_dim is None: + raise ValueError( + "Unable to infer the value of `text_encoder_projection_dim` for generating `text_embeds`, please double check the config of your model." + ) dim = self.text_encoder_projection_dim elif input_name == "timestep_cond": + if self.time_cond_proj_dim is None: + raise ValueError( + "Unable to infer the value of `time_cond_proj_dim` for generating `timestep_cond`, please double check the config of your model." + ) dim = self.time_cond_proj_dim else: dim = self.time_ids @@ -1027,7 +1095,7 @@ def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int class BloomDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator): def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): - if check_if_transformers_greater("4.44"): + if is_transformers_version(">=", "4.44"): return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) else: past_key_shape = ( @@ -1411,3 +1479,116 @@ def generate( float_dtype: str = "fp32", ): return self.random_int_tensor(shape=(1,), min_value=20, max_value=22, framework=framework, dtype=int_dtype) + + +class DummyTransformerTimestepInputGenerator(DummyTimestepInputGenerator): + SUPPORTED_INPUT_NAMES = ("timestep",) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "timestep": + shape = [self.batch_size] # With transformer diffusers, timestep is a 1D tensor + 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 DummyTransformerVisionInputGenerator(DummyVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ("hidden_states",) + + +class DummyTransformerTextInputGenerator(DummySeq2SeqDecoderTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "encoder_hidden_states", + "pooled_projection", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "encoder_hidden_states": + return super().generate(input_name, framework, int_dtype, float_dtype)[0] + + elif input_name == "pooled_projections": + return self.random_float_tensor( + [self.batch_size, self.normalized_config.projection_size], framework=framework, dtype=float_dtype + ) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyFluxTransformerVisionInputGenerator(DummyTransformerVisionInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "hidden_states", + "img_ids", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "hidden_states": + shape = [self.batch_size, (self.height // 2) * (self.width // 2), self.num_channels] + return self.random_float_tensor(shape, framework=framework, dtype=float_dtype) + elif input_name == "img_ids": + shape = ( + [(self.height // 2) * (self.width // 2), 3] + if is_diffusers_version(">=", "0.31.0") + else [self.batch_size, (self.height // 2) * (self.width // 2), 3] + ) + return self.random_int_tensor(shape, max_value=1, framework=framework, dtype=int_dtype) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyFluxTransformerTextInputGenerator(DummyTransformerTextInputGenerator): + SUPPORTED_INPUT_NAMES = ( + "encoder_hidden_states", + "pooled_projections", + "guidance", + "txt_ids", + ) + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "txt_ids": + shape = ( + [self.sequence_length, 3] + if is_diffusers_version(">=", "0.31.0") + else [self.batch_size, self.sequence_length, 3] + ) + return self.random_int_tensor(shape, max_value=1, framework=framework, dtype=int_dtype) + elif input_name == "guidance": + shape = [self.batch_size] + return self.random_float_tensor(shape, min_value=0, max_value=1, framework=framework, dtype=float_dtype) + + return super().generate(input_name, framework, int_dtype, float_dtype) + + +class DummyPatchTSTInputGenerator(DummyInputGenerator): + SUPPORTED_INPUT_NAMES = ("past_values",) + + def __init__( + self, + task: str, + normalized_config: NormalizedConfig, + batch_size: int = DEFAULT_DUMMY_SHAPES["batch_size"], + **kwargs, + ): + self.task = task + self.normalized_config = normalized_config + + self.batch_size = batch_size + self.context_length = normalized_config.context_length + self.num_input_channels = normalized_config.num_input_channels + + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + return self.random_float_tensor( + shape=[self.batch_size, self.context_length, self.num_input_channels], + min_value=-1, + max_value=1, + framework=framework, + dtype=float_dtype, + ) + + +class MCTCTDummyAudioInputGenerator(DummyAudioInputGenerator): + def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"): + if input_name == "input_features": + shape = [self.batch_size, self.sequence_length, self.normalized_config.input_features_per_channel] + return self.random_float_tensor(shape, min_value=-1, max_value=1, framework=framework, dtype=float_dtype) + + return super().generate(input_name, framework=framework, int_dtype=int_dtype, float_dtype=float_dtype) diff --git a/optimum/utils/normalized_config.py b/optimum/utils/normalized_config.py index 81207b7649..3f497b5920 100644 --- a/optimum/utils/normalized_config.py +++ b/optimum/utils/normalized_config.py @@ -77,6 +77,11 @@ def has_attribute(self, attr_name): return True +class NormalizedTimeSeriesForecastingConfig(NormalizedConfig): + NUM_INPUT_CHANNELS = "num_input_channels" + CONTEXT_LENGTH = "context_length" + + class NormalizedTextConfig(NormalizedConfig): VOCAB_SIZE = "vocab_size" HIDDEN_SIZE = "hidden_size" @@ -204,8 +209,10 @@ class NormalizedConfigManager: 'data2vec-text', 'data2vec-vision', 'detr', + 'dinov2', 'flaubert', 'groupvit', + 'hiera', 'ibert', 'layoutlm', 'layoutlmv3', @@ -216,6 +223,8 @@ class NormalizedConfigManager: 'owlvit', 'perceiver', 'roformer', + 'segformer', + 'siglip', 'squeezebert', 'table-transformer', """ @@ -225,8 +234,8 @@ class NormalizedConfigManager: "albert": NormalizedTextConfig, "bart": BartLikeNormalizedTextConfig, "bert": NormalizedTextConfig, - # "big_bird": NormalizedTextConfig, - # "bigbird_pegasus": BartLikeNormalizedTextConfig, + "big-bird": NormalizedTextConfig, + "bigbird-pegasus": BartLikeNormalizedTextConfig, "blenderbot": BartLikeNormalizedTextConfig, "blenderbot-small": BartLikeNormalizedTextConfig, "bloom": NormalizedTextConfig.with_args(num_layers="n_layer"), @@ -281,6 +290,7 @@ class NormalizedConfigManager: "xlm-roberta": NormalizedTextConfig, "yolos": NormalizedVisionConfig, "qwen2": NormalizedTextConfig, + "granite": NormalizedTextConfigWithGQA, } @classmethod diff --git a/optimum/utils/preprocessing/base.py b/optimum/utils/preprocessing/base.py index dc995ccc50..7cfda13ba7 100644 --- a/optimum/utils/preprocessing/base.py +++ b/optimum/utils/preprocessing/base.py @@ -20,15 +20,16 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union -from datasets import Dataset, DatasetDict -from datasets import load_dataset as datasets_load_dataset from transformers import PreTrainedTokenizerBase from transformers.image_processing_utils import BaseImageProcessor +from optimum.utils.import_utils import requires_backends + from .. import logging if TYPE_CHECKING: + from datasets import Dataset, DatasetDict from transformers import PretrainedConfig @@ -102,11 +103,14 @@ def create_dataset_processing_func( def prepare_dataset( self, - dataset: Union[DatasetDict, Dataset], + dataset: Union["DatasetDict", "Dataset"], data_keys: Dict[str, str], ref_keys: Optional[List[str]] = None, split: Optional[str] = None, - ) -> Union[DatasetDict, Dataset]: + ) -> Union["DatasetDict", "Dataset"]: + requires_backends(self, ["datasets"]) + from datasets import Dataset + if isinstance(dataset, Dataset) and split is not None: raise ValueError("A Dataset and a split name were provided, but splits are for DatasetDict.") elif split is not None: @@ -131,7 +135,12 @@ def load_dataset( num_samples: Optional[int] = None, shuffle: bool = False, **load_dataset_kwargs, - ) -> Union[DatasetDict, Dataset]: + ) -> Union["DatasetDict", "Dataset"]: + requires_backends(self, ["datasets"]) + + from datasets import Dataset, DatasetDict + from datasets import load_dataset as datasets_load_dataset + dataset = datasets_load_dataset(path, **load_dataset_kwargs) if isinstance(dataset, DatasetDict) and load_smallest_split: diff --git a/optimum/utils/testing_utils.py b/optimum/utils/testing_utils.py index 76fe9a05b1..88b1acdb78 100644 --- a/optimum/utils/testing_utils.py +++ b/optimum/utils/testing_utils.py @@ -28,6 +28,7 @@ from . import ( is_accelerate_available, is_auto_gptq_available, + is_datasets_available, is_diffusers_available, is_sentence_transformers_available, is_timm_available, @@ -146,6 +147,10 @@ def require_sentence_transformers(test_case): return unittest.skipUnless(is_sentence_transformers_available(), "test requires sentence-transformers")(test_case) +def require_datasets(test_case): + return unittest.skipUnless(is_datasets_available(), "test requires datasets")(test_case) + + def grid_parameters( parameters: Dict[str, Iterable[Any]], yield_dict: bool = False, diff --git a/optimum/version.py b/optimum/version.py index 4fff28e5c9..a1a07cef8f 100644 --- a/optimum/version.py +++ b/optimum/version.py @@ -12,4 +12,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.24.0.dev0" +__version__ = "1.24.0" diff --git a/pyproject.toml b/pyproject.toml index 99a0f1c85f..17bcd90e06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ markers = [ "rocm_ep_test", "tensorflow_test", "timm_test", + "datasets_test", "run_in_series", "run_slow", "accelerate_test", diff --git a/setup.py b/setup.py index 82892bfcc8..d132975aa4 100644 --- a/setup.py +++ b/setup.py @@ -13,14 +13,11 @@ REQUIRED_PKGS = [ - "coloredlogs", - "sympy", "transformers>=4.29", "torch>=1.11", "packaging", "numpy", "huggingface_hub>=0.8.0", - "datasets", ] # TODO: unpin pytest once https://github.com/huggingface/transformers/pull/29154 is merged & released @@ -34,7 +31,6 @@ "Pillow", "sacremoses", "torchvision", - "diffusers>=0.17.0", "torchaudio", "einops", "timm", @@ -54,7 +50,7 @@ "datasets>=1.2.1", "evaluate", "protobuf>=3.20.1", - "transformers<4.47.0", + "transformers>=4.36,<4.49.0", ], "onnxruntime-gpu": [ "onnx", @@ -62,20 +58,28 @@ "datasets>=1.2.1", "evaluate", "protobuf>=3.20.1", - "accelerate", # ORTTrainer requires it. - "transformers<4.47.0", + "transformers>=4.36,<4.49.0", + ], + "onnxruntime-training": [ + "torch-ort", + "onnxruntime-training>=1.11.0", + "datasets>=1.2.1", + "accelerate", + "evaluate", + "protobuf>=3.20.1", + "transformers>=4.36,<4.49.0", ], "exporters": [ "onnx", "onnxruntime", "timm", - "transformers<4.47.0", + "transformers>=4.36,<4.49.0", ], "exporters-gpu": [ "onnx", "onnxruntime-gpu", "timm", - "transformers<4.47.0", + "transformers>=4.36,<4.49.0", ], "exporters-tf": [ "tensorflow>=2.4,<=2.12.1", @@ -86,7 +90,7 @@ "h5py", "numpy<1.24.0", "datasets<=2.16", - "transformers>=4.26,<4.38", + "transformers>=4.36,<4.38", ], "diffusers": ["diffusers"], "intel": "optimum-intel>=1.18.0", @@ -123,9 +127,10 @@ "Intended Audience :: Education", "Intended Audience :: Science/Research", "Operating System :: OS Independent", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", "Topic :: Scientific/Engineering :: Artificial Intelligence", ], keywords="transformers, quantization, pruning, optimization, training, inference, onnx, onnx runtime, intel, " @@ -137,7 +142,7 @@ packages=find_namespace_packages(include=["optimum*"]), install_requires=REQUIRED_PKGS, extras_require=EXTRAS_REQUIRE, - python_requires=">=3.7.0", + python_requires=">=3.9.0", include_package_data=True, zip_safe=False, entry_points={"console_scripts": ["optimum-cli=optimum.commands.optimum_cli:main"]}, diff --git a/tests/exporters/Dockerfile_exporters_gpu b/tests/exporters/Dockerfile_exporters_gpu index 1f2bd35066..9e9bce2a8b 100644 --- a/tests/exporters/Dockerfile_exporters_gpu +++ b/tests/exporters/Dockerfile_exporters_gpu @@ -1,6 +1,6 @@ # use version with cudnn 8.5 to match torch==1.13.1 that uses 8.5.0.96 # has Python 3.8.10 -FROM nvcr.io/nvidia/tensorrt:22.08-py3 +FROM nvcr.io/nvidia/tensorrt:24.02-py3 CMD nvidia-smi # Ignore interactive questions during `docker build` diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index c8a33b0be3..8705765bb2 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. + VALIDATE_EXPORT_ON_SHAPES_SLOW = { "batch_size": [1, 3, 5], "sequence_length": [8, 33, 96, 154], @@ -36,6 +37,7 @@ PYTORCH_EXPORT_MODELS_TINY = { "albert": "hf-internal-testing/tiny-random-AlbertModel", + "audio-spectrogram-transformer": "hf-internal-testing/tiny-random-ASTModel", "beit": "hf-internal-testing/tiny-random-BeitForImageClassification", "bert": { "hf-internal-testing/tiny-random-BertModel": [ @@ -49,8 +51,8 @@ "nreimers/BERT-Tiny_L-2_H-128_A-2": ["feature-extraction"], }, "bart": "hf-internal-testing/tiny-random-bart", - # "big-bird": "hf-internal-testing/tiny-random-BigBirdModel", - # "bigbird-pegasus": "hf-internal-testing/tiny-random-bigbird_pegasus", + "big-bird": "hf-internal-testing/tiny-random-BigBirdModel", + "bigbird-pegasus": "hf-internal-testing/tiny-random-bigbird_pegasus", "blenderbot-small": "hf-internal-testing/tiny-random-BlenderbotModel", "blenderbot": "hf-internal-testing/tiny-random-BlenderbotModel", "bloom": "hf-internal-testing/tiny-random-BloomModel", @@ -67,10 +69,12 @@ "data2vec-audio": "hf-internal-testing/tiny-random-Data2VecAudioModel", "deberta": "hf-internal-testing/tiny-random-DebertaModel", "deberta-v2": "hf-internal-testing/tiny-random-DebertaV2Model", + "decision-transformer": "edbeeching/decision-transformer-gym-hopper-medium", "deit": "hf-internal-testing/tiny-random-DeiTModel", + "dinov2": "hf-internal-testing/tiny-random-Dinov2Model", "donut": "fxmarty/tiny-doc-qa-vision-encoder-decoder", "donut-swin": "hf-internal-testing/tiny-random-DonutSwinModel", - "detr": "hf-internal-testing/tiny-random-DetrModel", # hf-internal-testing/tiny-random-detr is larger + "detr": "hf-internal-testing/tiny-random-DetrModel", "distilbert": "hf-internal-testing/tiny-random-DistilBertModel", "dpt": "hf-internal-testing/tiny-random-DPTModel", "electra": "hf-internal-testing/tiny-random-ElectraModel", @@ -100,7 +104,9 @@ "gpt-neo": "hf-internal-testing/tiny-random-GPTNeoModel", "gpt-neox": "hf-internal-testing/tiny-random-GPTNeoXForCausalLM", "gptj": "hf-internal-testing/tiny-random-GPTJModel", + "granite": "hf-internal-testing/tiny-random-GraniteForCausalLM", "groupvit": "hf-internal-testing/tiny-random-groupvit", + "hiera": "hf-internal-testing/tiny-random-HieraForImageClassification", "ibert": "hf-internal-testing/tiny-random-IBertModel", "imagegpt": "hf-internal-testing/tiny-random-ImageGPTModel", "levit": "hf-internal-testing/tiny-random-LevitModel", @@ -109,24 +115,33 @@ "lilt": "hf-internal-testing/tiny-random-LiltModel", "llama": "fxmarty/tiny-llama-fast-tokenizer", "longt5": "fxmarty/tiny-random-working-LongT5Model", - # "longformer": "allenai/longformer-base-4096", + "longformer": "hf-internal-testing/tiny-random-LongformerModel", "m2m-100": "hf-internal-testing/tiny-random-m2m_100", "marian": "sshleifer/tiny-marian-en-de", # hf-internal-testing ones are broken "markuplm": "hf-internal-testing/tiny-random-MarkupLMModel", + "maskformer": "hf-internal-testing/tiny-random-MaskFormerForInstanceSegmentation", "mbart": "hf-internal-testing/tiny-random-mbart", + "mctct": "hf-internal-testing/tiny-random-MCTCTModel", + "megatron-bert": "hf-internal-testing/tiny-random-MegatronBertModel", + "mgp-str": "hf-internal-testing/tiny-random-MgpstrForSceneTextRecognition", "mistral": "echarlaix/tiny-random-mistral", "mobilebert": "hf-internal-testing/tiny-random-MobileBertModel", "mobilenet-v2": "hf-internal-testing/tiny-random-MobileNetV2Model", - "mobilenet-v1": "google/mobilenet_v1_0.75_192", + "mobilenet-v1": "hf-internal-testing/tiny-random-MobileNetV1Model", "mobilevit": "hf-internal-testing/tiny-random-mobilevit", + "modernbert": "hf-internal-testing/tiny-random-ModernBertForMaskedLM", "mpnet": "hf-internal-testing/tiny-random-MPNetModel", "mpt": "hf-internal-testing/tiny-random-MptForCausalLM", "mt5": "lewtun/tiny-random-mt5", "musicgen": "hf-internal-testing/tiny-random-MusicgenForConditionalGeneration", "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel", + "olmo": "hf-internal-testing/tiny-random-OlmoForCausalLM", + "olmo2": "hf-internal-testing/tiny-random-Olmo2ForCausalLM", "opt": "hf-internal-testing/tiny-random-OPTModel", "owlv2": "hf-internal-testing/tiny-random-Owlv2Model", "owlvit": "hf-tiny-model-private/tiny-random-OwlViTModel", + "patchtst": "ibm/test-patchtst", + "patchtsmixer": "ibm/test-patchtsmixer", "pegasus": "hf-internal-testing/tiny-random-PegasusModel", "perceiver": { "hf-internal-testing/tiny-random-language_perceiver": ["fill-mask", "text-classification"], @@ -135,8 +150,9 @@ "phi": "echarlaix/tiny-random-PhiForCausalLM", "phi3": "Xenova/tiny-random-Phi3ForCausalLM", "pix2struct": "fxmarty/pix2struct-tiny-random", - # "rembert": "google/rembert", + "rembert": "hf-internal-testing/tiny-random-RemBertModel", "poolformer": "hf-internal-testing/tiny-random-PoolFormerModel", + "pvt": "hf-internal-testing/tiny-random-PvtForImageClassification", "qwen2": "fxmarty/tiny-dummy-qwen2", "regnet": "hf-internal-testing/tiny-random-RegNetModel", "resnet": "hf-internal-testing/tiny-random-resnet", @@ -144,16 +160,21 @@ "roformer": "hf-internal-testing/tiny-random-RoFormerModel", "sam": "fxmarty/sam-vit-tiny-random", "segformer": "hf-internal-testing/tiny-random-SegformerModel", + "siglip": "hf-internal-testing/tiny-random-SiglipModel", + "siglip-vision-model": "hf-internal-testing/tiny-random-SiglipVisionModel", "splinter": "hf-internal-testing/tiny-random-SplinterModel", "squeezebert": "hf-internal-testing/tiny-random-SqueezeBertModel", "swin": "hf-internal-testing/tiny-random-SwinModel", + "swinv2": "hf-internal-testing/tiny-random-Swinv2Model", "swin2sr": "hf-internal-testing/tiny-random-Swin2SRModel", "t5": "hf-internal-testing/tiny-random-t5", "table-transformer": "hf-internal-testing/tiny-random-TableTransformerModel", "vit": "hf-internal-testing/tiny-random-vit", + "vit-mae": "hf-internal-testing/tiny-random-ViTMAEModel", + "vit-msn": "hf-internal-testing/tiny-random-ViTMSNForImageClassification", "vits": "echarlaix/tiny-random-vits", "yolos": "hf-internal-testing/tiny-random-YolosModel", - "whisper": "openai/whisper-tiny.en", # hf-internal-testing ones are broken + "whisper": "optimum-internal-testing/tiny-random-whisper", "hubert": "hf-internal-testing/tiny-random-HubertModel", "wav2vec2": "hf-internal-testing/tiny-random-Wav2Vec2Model", "wav2vec2-conformer": "hf-internal-testing/tiny-random-wav2vec2-conformer", @@ -178,9 +199,6 @@ "hf-internal-testing/tiny-random-UniSpeechSatForPreTraining": ["audio-frame-classification"], "hf-internal-testing/tiny-random-UniSpeechSatForXVector": ["audio-xvector"], }, - "audio-spectrogram-transformer": "Ericwang/tiny-random-ast", - # Disabled for now because some operator seems to not be supported by ONNX. - # "mctct": "hf-internal-testing/tiny-random-MCTCTModel", "speech-to-text": "hf-internal-testing/tiny-random-Speech2TextModel", "speecht5": "hf-internal-testing/tiny-random-SpeechT5ForTextToSpeech", "xlm": "hf-internal-testing/tiny-random-XLMModel", @@ -198,39 +216,41 @@ }, } - +# TODO: enable export slow tests PYTORCH_EXPORT_MODELS_LARGE = { "albert": "albert-base-v2", + "audio-spectrogram-transformer": "nielsr/audio-spectogram-transformer-finetuned-audioset-10-10-0.4593", "beit": "microsoft/beit-base-patch16-224", "bert": "bert-base-cased", "bart": "facebook/bart-base", - # "big-bird": "google/bigbird-roberta-base", - # "bigbird-pegasus": "hf-internal-testing/tiny-random-bigbird_pegasus", + "big-bird": "google/bigbird-roberta-base", + "bigbird-pegasus": "hf-internal-testing/tiny-random-bigbird_pegasus", "blenderbot-small": "facebook/blenderbot_small-90M", "blenderbot": "facebook/blenderbot-90M", - "bloom": "hf-internal-testing/tiny-random-BloomModel", # Not using bigscience/bloom-560m because it goes OOM. + "bloom": "bigscience/bloom-560m", "camembert": "camembert-base", "clip": "openai/clip-vit-base-patch32", "convbert": "YituTech/conv-bert-base", "convnext": "facebook/convnext-tiny-224", - "codegen": "hf-internal-testing/tiny-random-CodeGenModel", # Not using Salesforce/codegen-350M-multi because it takes too much time for testing. + "codegen": "Salesforce/codegen-350M-multi", "data2vec-text": "facebook/data2vec-text-base", "data2vec-vision": "facebook/data2vec-vision-base", "data2vec-audio": "facebook/data2vec-audio-base", - "deberta": "hf-internal-testing/tiny-random-DebertaModel", # Not using microsoft/deberta-base because it takes too much time for testing. - "deberta-v2": "hf-internal-testing/tiny-random-DebertaV2Model", # Not using microsoft/deberta-v2-xlarge because it takes too much time for testing. + "deberta": "microsoft/deberta-base", + "deberta-v2": "microsoft/deberta-v2-xlarge", "deit": "facebook/deit-small-patch16-224", - "detr": "hf-internal-testing/tiny-random-detr", # Not using facebook/detr-resnet-50 because it takes too much time for testing. + "detr": "facebook/detr-resnet-50", "distilbert": "distilbert-base-cased", "electra": "google/electra-base-generator", "encoder-decoder": "patrickvonplaten/bert2bert_cnn_daily_mail", - "flaubert": "hf-internal-testing/tiny-random-flaubert", # TODO + "flaubert": "flaubert/flaubert_small_cased", "gemma": "google/gemma-2b", "gpt2": "gpt2", "gpt-neo": "EleutherAI/gpt-neo-125M", "gpt-neox": "EleutherAI/gpt-neox-20b", - "gptj": "anton-l/gpt-j-tiny-random", # TODO + "gptj": "architext/gptj-162M", "groupvit": "nvidia/groupvit-gcc-yfcc", + "hiera": "facebook/hiera-tiny-224-in1k-hf", "ibert": "kssteven/ibert-roberta-base", "imagegpt": "openai/imagegpt-small", "levit": "facebook/levit-128S", @@ -238,37 +258,45 @@ "layoutlmv3": "microsoft/layoutlmv3-base", "lilt": "SCUT-DLVCLab/lilt-roberta-en-base", "llama": "decapoda-research/llama-65b-hf", - "longt5": "fxmarty/tiny-random-working-LongT5Model", # Not using google/long-t5-local-base because it takes too much time for testing. - # "longformer": "allenai/longformer-base-4096", - "m2m-100": "hf-internal-testing/tiny-random-m2m_100", # Not using facebook/m2m100_418M because it takes too much time for testing. + "longt5": "google/long-t5-local-base", + "longformer": "allenai/longformer-base-4096", + "m2m-100": "facebook/m2m100_418M", "marian": "Helsinki-NLP/opus-mt-en-de", "markuplm": "hf-internal-testing/tiny-random-MarkupLMModel", + "maskformer": "facebook/maskformer-swin-tiny-coco", "mbart": "sshleifer/tiny-mbart", + "mgp-str": "alibaba-damo/mgp-str-base", "mobilebert": "google/mobilebert-uncased", - # "mobilenet_v1": "google/mobilenet_v1_0.75_192", - # "mobilenet_v2": "google/mobilenet_v2_0.35_96", + "mobilenet_v1": "google/mobilenet_v1_0.75_192", + "mobilenet_v2": "google/mobilenet_v2_0.35_96", "mobilevit": "apple/mobilevit-small", + "modernbert": "answerdotai/ModernBERT-base", "mpt": "mosaicml/mpt-7b", - "mt5": "lewtun/tiny-random-mt5", # Not using google/mt5-small because it takes too much time for testing. + "mt5": "google/mt5-small", "musicgen": "facebook/musicgen-small", "nystromformer": "hf-internal-testing/tiny-random-NystromformerModel", "owlv2": "google/owlv2-base-patch16", "owlvit": "google/owlvit-base-patch32", - "perceiver": "hf-internal-testing/tiny-random-PerceiverModel", # Not using deepmind/language-perceiver because it takes too much time for testing. - # "rembert": "google/rembert", + "perceiver": "deepmind/language-perceiver", + "rembert": "google/rembert", "poolformer": "hf-internal-testing/tiny-random-PoolFormerModel", + "pvt": "hf-internal-testing/tiny-random-PvtForImageClassification", "regnet": "facebook/regnet-y-040", "resnet": "microsoft/resnet-50", "roberta": "roberta-base", "roformer": "junnyu/roformer_chinese_base", "sam": "facebook/sam-vit-base", "segformer": "nvidia/segformer-b0-finetuned-ade-512-512", + "siglip": "google/siglip-base-patch16-224", "splinter": "hf-internal-testing/tiny-random-SplinterModel", "squeezebert": "squeezebert/squeezebert-uncased", "swin": "microsoft/swin-tiny-patch4-window7-224", + "swinv2": "microsoft/swinv2-tiny-patch4-window16-256", "t5": "t5-small", "table-transformer": "microsoft/table-transformer-detection", "vit": "google/vit-base-patch16-224", + "vit-mae": "facebook/vit-mae-base", + "vit-msn": "facebook/vit-msn-small", "yolos": "hustvl/yolos-tiny", "whisper": "openai/whisper-tiny.en", "hubert": "facebook/hubert-base-ls960", @@ -279,9 +307,7 @@ "sew-d": "asapp/sew-d-tiny-100k-ft-ls100h", "unispeech": "microsoft/unispeech-1350-en-353-fr-ft-1h", "unispeech-sat": "microsoft/unispeech-sat-base", - "audio-spectrogram-transformer": "nielsr/audio-spectogram-transformer-finetuned-audioset-10-10-0.4593", - # Disabled for now because some operator seems to not be supported by ONNX. - # "mctct": "speechbrain/m-ctc-t-large", + "mctct": "speechbrain/m-ctc-t-large", "speech-to-text": "codenamewei/speech-to-text", "xlm": "xlm-clm-ende-1024", "xlm-roberta": "Unbabel/xlm-roberta-comet-small", @@ -296,9 +322,11 @@ } PYTORCH_DIFFUSION_MODEL = { + "flux": "optimum-internal-testing/tiny-random-flux", + "latent-consistency": "echarlaix/tiny-random-latent-consistency", "stable-diffusion": "hf-internal-testing/tiny-stable-diffusion-torch", + "stable-diffusion-3": "yujiepan/stable-diffusion-3-tiny-random", "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl", - "latent-consistency": "echarlaix/tiny-random-latent-consistency", } PYTORCH_TIMM_MODEL = { diff --git a/tests/exporters/onnx/test_onnx_export.py b/tests/exporters/onnx/test_onnx_export.py index 7671d6cd2e..aef3bc8dfc 100644 --- a/tests/exporters/onnx/test_onnx_export.py +++ b/tests/exporters/onnx/test_onnx_export.py @@ -207,20 +207,18 @@ def _onnx_export( model.config.pad_token_id = 0 if is_torch_available(): - from optimum.utils import torch_version + from optimum.utils.import_utils import _torch_version, _transformers_version if not onnx_config.is_transformers_support_available: - import transformers - pytest.skip( "Skipping due to incompatible Transformers version. Minimum required is" - f" {onnx_config.MIN_TRANSFORMERS_VERSION}, got: {transformers.__version__}" + f" {onnx_config.MIN_TRANSFORMERS_VERSION}, got: {_transformers_version}" ) if not onnx_config.is_torch_support_available: pytest.skip( "Skipping due to incompatible PyTorch version. Minimum required is" - f" {onnx_config.MIN_TORCH_VERSION}, got: {torch_version}" + f" {onnx_config.MIN_TORCH_VERSION}, got: {_torch_version}" ) atol = onnx_config.ATOL_FOR_VALIDATION @@ -299,7 +297,6 @@ def _onnx_export_diffusion_models(self, model_type: str, model_name: str, device with TemporaryDirectory() as tmpdirname: _, onnx_outputs = export_models( models_and_onnx_configs=models_and_onnx_configs, - opset=14, output_dir=Path(tmpdirname), device=device, ) @@ -307,7 +304,6 @@ def _onnx_export_diffusion_models(self, model_type: str, model_name: str, device models_and_onnx_configs=models_and_onnx_configs, onnx_named_outputs=onnx_outputs, output_dir=Path(tmpdirname), - atol=1e-4, use_subprocess=False, ) diff --git a/tests/onnx/test_onnx_export_custom_module.py b/tests/onnx/test_onnx_export_custom_module.py index 4398c14f01..9416093c84 100644 --- a/tests/onnx/test_onnx_export_custom_module.py +++ b/tests/onnx/test_onnx_export_custom_module.py @@ -22,7 +22,7 @@ if is_torch_available(): import torch - from transformers.models.deberta import modeling_deberta + from transformers.models.sew_d import modeling_sew_d from optimum.utils import check_if_torch_greater @@ -36,7 +36,7 @@ def test_training(self): """Tests export of StableDropout in training mode.""" devnull = open(os.devnull, "wb") # drop_prob must be > 0 for the test to be meaningful - sd = modeling_deberta.StableDropout(0.1) + sd = modeling_sew_d.StableDropout(0.1) # Avoid warnings in training mode do_constant_folding = False # Dropout is a no-op in inference mode diff --git a/tests/onnxruntime/ds_configs/ds_config_zero_stage_1.json b/tests/onnxruntime-training/ds_configs/ds_config_zero_stage_1.json similarity index 100% rename from tests/onnxruntime/ds_configs/ds_config_zero_stage_1.json rename to tests/onnxruntime-training/ds_configs/ds_config_zero_stage_1.json diff --git a/tests/onnxruntime/ds_configs/ds_config_zero_stage_2.json b/tests/onnxruntime-training/ds_configs/ds_config_zero_stage_2.json similarity index 100% rename from tests/onnxruntime/ds_configs/ds_config_zero_stage_2.json rename to tests/onnxruntime-training/ds_configs/ds_config_zero_stage_2.json diff --git a/tests/onnxruntime/ds_configs/ds_config_zero_stage_3.json b/tests/onnxruntime-training/ds_configs/ds_config_zero_stage_3.json similarity index 100% rename from tests/onnxruntime/ds_configs/ds_config_zero_stage_3.json rename to tests/onnxruntime-training/ds_configs/ds_config_zero_stage_3.json diff --git a/tests/onnxruntime/ds_configs/ds_config_zero_stage_inifinity.json b/tests/onnxruntime-training/ds_configs/ds_config_zero_stage_inifinity.json similarity index 100% rename from tests/onnxruntime/ds_configs/ds_config_zero_stage_inifinity.json rename to tests/onnxruntime-training/ds_configs/ds_config_zero_stage_inifinity.json diff --git a/tests/onnxruntime/training/nightly_test_examples.py b/tests/onnxruntime-training/test_examples.py similarity index 91% rename from tests/onnxruntime/training/nightly_test_examples.py rename to tests/onnxruntime-training/test_examples.py index a16913a097..5873f238af 100644 --- a/tests/onnxruntime/training/nightly_test_examples.py +++ b/tests/onnxruntime-training/test_examples.py @@ -25,7 +25,7 @@ class ORTTrainerExampleTest(unittest.TestCase): def test_text_classification(self): subprocess.run( - "cp ../examples/onnxruntime/training/text-classification/run_glue.py ./", + "cp examples/onnxruntime/training/text-classification/run_glue.py ./", shell=True, ) @@ -51,7 +51,7 @@ def test_text_classification(self): def test_token_classification(self): subprocess.run( - "cp ../examples/onnxruntime/training/token-classification/run_ner.py ./", + "cp examples/onnxruntime/training/token-classification/run_ner.py ./", shell=True, ) @@ -75,7 +75,7 @@ def test_token_classification(self): def test_translation(self): subprocess.run( - "cp ../examples/onnxruntime/training/translation/run_translation.py ./", + "cp examples/onnxruntime/training/translation/run_translation.py ./", shell=True, ) @@ -105,7 +105,7 @@ def test_translation(self): @pytest.mark.skip(reason="skip for now") def test_summarization(self): subprocess.run( - "cp ../examples/onnxruntime/training/summarization/run_summarization.py ./", + "cp examples/onnxruntime/training/summarization/run_summarization.py ./", shell=True, ) @@ -139,7 +139,7 @@ def test_stable_diffusion_txt2img(self): @pytest.mark.skip(reason="skip for now") def test_question_answering(self): subprocess.run( - "cp ../examples/onnxruntime/training/question-answering/run_qa.py ./", + "cp examples/onnxruntime/training/question-answering/run_qa.py ./", shell=True, ) @@ -166,7 +166,7 @@ def test_question_answering(self): @pytest.mark.skip(reason="skip for now") def test_language_modeling(self): subprocess.run( - "cp ../examples/onnxruntime/training/question-answering/run_qa.py ./", + "cp examples/onnxruntime/training/question-answering/run_qa.py ./", shell=True, ) @@ -194,7 +194,7 @@ def test_language_modeling(self): @pytest.mark.skip(reason="skip for now") def test_image_classification(self): subprocess.run( - "cp ../examples/onnxruntime/training/image-classification/run_image_classification.py ./", + "cp examples/onnxruntime/training/image-classification/run_image_classification.py ./", shell=True, ) diff --git a/tests/onnxruntime/training/nightly_test_trainer.py b/tests/onnxruntime-training/test_trainer.py similarity index 97% rename from tests/onnxruntime/training/nightly_test_trainer.py rename to tests/onnxruntime-training/test_trainer.py index e24ee30617..ac4413c639 100644 --- a/tests/onnxruntime/training/nightly_test_trainer.py +++ b/tests/onnxruntime-training/test_trainer.py @@ -60,11 +60,11 @@ nltk.download("punkt") _ENCODERS_TO_TEST = { - ("distilbert", "distilbert-base-cased"), + ("distilbert", "distilbert-base-uncased"), } _DECODERS_TO_TEST = { - ("gpt2", "gpt2"), + ("gpt2", "distilgpt2"), } _SEQ2SEQ_MODELS_TO_TEST = { @@ -78,11 +78,6 @@ "data_collator": default_data_collator, "data_collator_class": DataCollatorWithPadding, }, - # "token-classification": { - # "dataset": ["conll2003"], - # "metric": ["seqeval"], - # "data_collator_class": DataCollatorForTokenClassification, - # }, } _DECODER_TASKS_DATASETS_CONFIGS = { @@ -235,7 +230,7 @@ def load_and_prepare(task): def load_and_prepare_glue(model_name, data_metric_config, max_seq_length, padding="max_length", **kwargs): # Prepare model - model = AutoModelForSequenceClassification.from_pretrained(model_name) + model = AutoModelForSequenceClassification.from_pretrained(model_name, attn_implementation="eager") tokenizer = AutoTokenizer.from_pretrained(model_name) # Prepare dataset @@ -295,7 +290,9 @@ def load_and_prepare_ner(model_name, data_metric_config, max_seq_length, padding label_list = dataset["train"].features[f"{task}_tags"].feature.names # Prepare model - model = AutoModelForTokenClassification.from_pretrained(model_name, num_labels=len(label_list)) + model = AutoModelForTokenClassification.from_pretrained( + model_name, num_labels=len(label_list), attn_implementation="eager" + ) if model_name.split("-")[0] in {"gpt2", "roberta"}: tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, add_prefix_space=True) else: @@ -387,7 +384,7 @@ def load_and_prepare_clm(model_name, data_metric_config, max_seq_length, padding metric = load(*data_metric_config["metric"]) # Prepare model - model = AutoModelForCausalLM.from_pretrained(model_name) + model = AutoModelForCausalLM.from_pretrained(model_name, attn_implementation="eager") tokenizer = AutoTokenizer.from_pretrained(model_name) # Prepare dataset @@ -462,7 +459,7 @@ def compute_metrics(eval_pred): def load_and_prepare_xsum(model_name, data_metric_config, _, **kwargs): # Prepare model - model = AutoModelForSeq2SeqLM.from_pretrained(model_name) + model = AutoModelForSeq2SeqLM.from_pretrained(model_name, attn_implementation="eager") tokenizer = AutoTokenizer.from_pretrained(model_name) # Load dataset and metric @@ -600,7 +597,7 @@ def test_trainer_fp32(self, test_name, model_name, task, data_metric_config): trainer.train() trainer.save_model() trainer.evaluate() - trainer.predict(test_dataset) + # trainer.predict(test_dataset) gc.collect() @slow @@ -639,7 +636,7 @@ def test_trainer_fp32_with_label_smoothing(self, test_name, model_name, task, da trainer.train() trainer.save_model() trainer.evaluate() - trainer.predict(test_dataset) + # trainer.predict(test_dataset) gc.collect() @slow @@ -678,7 +675,7 @@ def test_trainer_fp16(self, test_name, model_name, task, data_metric_config): trainer.train() trainer.save_model() trainer.evaluate() - trainer.predict(test_dataset) + # trainer.predict(test_dataset) gc.collect() @@ -730,7 +727,7 @@ def test_trainer_fp16_ds_stage1(self, test_name, model_name, task, data_metric_c weight_decay=self.weight_decay, logging_dir=tmp_dir, fp16=True, - deepspeed="onnxruntime/ds_configs/ds_config_zero_stage_1.json", + deepspeed="tests/onnxruntime-training/ds_configs/ds_config_zero_stage_1.json", ) trainer, _ = get_ort_trainer( @@ -769,7 +766,7 @@ def test_trainer_fp16_ds_stage2(self, test_name, model_name, task, data_metric_c weight_decay=self.weight_decay, logging_dir=tmp_dir, fp16=True, - deepspeed="onnxruntime/ds_configs/ds_config_zero_stage_2.json", + deepspeed="tests/onnxruntime-training/ds_configs/ds_config_zero_stage_2.json", ) trainer, _ = get_ort_trainer( diff --git a/tests/onnxruntime/docker/Dockerfile_onnxruntime_gpu b/tests/onnxruntime/docker/Dockerfile_onnxruntime_gpu deleted file mode 100644 index 9013697e04..0000000000 --- a/tests/onnxruntime/docker/Dockerfile_onnxruntime_gpu +++ /dev/null @@ -1,26 +0,0 @@ -# use version with CUDA 11.8 and TensorRT 8.5.1.7 to match ORT 1.14 requirements -FROM nvcr.io/nvidia/tensorrt:22.12-py3 -CMD nvidia-smi - -# Ignore interactive questions during `docker build` -ENV DEBIAN_FRONTEND noninteractive - -# Install and update tools to minimize security vulnerabilities -RUN apt-get update -RUN apt-get install -y software-properties-common wget apt-utils patchelf git libprotobuf-dev protobuf-compiler cmake \ - bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 mercurial subversion libopenmpi-dev ffmpeg && \ - apt-get clean -RUN unattended-upgrade -RUN apt-get autoremove -y - -RUN python -m pip install -U pip - -RUN pip install transformers torch onnxruntime-gpu -RUN pip install datasets evaluate diffusers scipy - -# Install Optimum -COPY . /workspace/optimum -RUN pip install /workspace/optimum[onnxruntime-gpu,tests] - -ENV TEST_LEVEL=1 -CMD pytest onnxruntime/test_*.py --durations=0 -s -vvvvv -m cuda_ep_test -m trt_ep_test diff --git a/tests/onnxruntime/docker/Dockerfile_onnxruntime_trainer b/tests/onnxruntime/docker/Dockerfile_onnxruntime_trainer deleted file mode 100644 index 74add3f07e..0000000000 --- a/tests/onnxruntime/docker/Dockerfile_onnxruntime_trainer +++ /dev/null @@ -1,83 +0,0 @@ -#!/usr/bin/env python -# coding=utf-8 -# Copyright 2023 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. - -# Use nvidia/cuda image -FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04 -CMD nvidia-smi - -# Ignore interactive questions during `docker build` -ENV DEBIAN_FRONTEND noninteractive - -# Bash shell -RUN chsh -s /bin/bash -SHELL ["/bin/bash", "-c"] - -# Versions -ARG PYTHON_VERSION=3.9 -ARG TORCH_CUDA_VERSION=cu118 -ARG TORCH_VERSION=2.0.0 -ARG TORCHVISION_VERSION=0.15.1 - -# Install and update tools to minimize security vulnerabilities -RUN apt-get update -RUN apt-get install -y software-properties-common wget apt-utils patchelf git libprotobuf-dev protobuf-compiler cmake \ - bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 mercurial subversion libopenmpi-dev ffmpeg && \ - apt-get clean -RUN unattended-upgrade -RUN apt-get autoremove -y - -# Install miniconda (comes with python 3.9 default) -ARG BUILD_USER=onnxruntimedev -ARG MINICONDA_PREFIX=/home/$BUILD_USER/miniconda3 -RUN apt-get install curl - -ARG CONDA_URL=https://repo.anaconda.com/miniconda/Miniconda3-py37_4.9.2-Linux-x86_64.sh -RUN curl -fSsL --insecure ${CONDA_URL} -o install-conda.sh && \ - /bin/bash ./install-conda.sh -b -p $MINICONDA_PREFIX && \ - $MINICONDA_PREFIX/bin/conda clean -ya && \ - $MINICONDA_PREFIX/bin/conda install -y python=${PYTHON_VERSION} - -ENV PATH=$MINICONDA_PREFIX/bin:${PATH} - -ARG PYTHON_EXE=$MINICONDA_PREFIX/bin/python - -# (Optional) Intall test dependencies -RUN $PYTHON_EXE -m pip install git+https://github.com/huggingface/transformers -RUN $PYTHON_EXE -m pip install datasets accelerate evaluate coloredlogs absl-py rouge_score seqeval scipy sacrebleu nltk scikit-learn parameterized sentencepiece -RUN $PYTHON_EXE -m pip install deepspeed mpi4py -# RUN $PYTHON_EXE -m pip install optuna ray sigopt wandb - -# PyTorch -RUN $PYTHON_EXE -m pip install onnx ninja -RUN $PYTHON_EXE -m pip install torch==${TORCH_VERSION} torchvision==${TORCHVISION_VERSION} -f https://download.pytorch.org/whl/${TORCH_CUDA_VERSION} - -# ORT Module -RUN $PYTHON_EXE -m pip install onnxruntime-training==1.16.3 -f https://download.onnxruntime.ai/onnxruntime_stable_cu118.html -RUN $PYTHON_EXE -m pip install torch-ort -ENV TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX" -RUN $PYTHON_EXE -m pip install --upgrade protobuf==3.20.2 -RUN $PYTHON_EXE -m torch_ort.configure - -# https://github.com/vllm-project/vllm/issues/1726 -RUN pip uninstall nvidia-nccl-cu12 -y - -# Install Optimum -COPY . /workspace/optimum -RUN pip install /workspace/optimum[tests] - -ENV TEST_LEVEL=1 -CMD RUN_SLOW=1 pytest -v -rs onnxruntime/training/nightly_test_trainer.py --durations=0 -CMD RUN_SLOW=1 pytest -v -rs onnxruntime/training/nightly_test_examples.py --durations=0 \ No newline at end of file diff --git a/tests/onnxruntime/test_diffusion.py b/tests/onnxruntime/test_diffusion.py index 956566f0e1..749e078456 100644 --- a/tests/onnxruntime/test_diffusion.py +++ b/tests/onnxruntime/test_diffusion.py @@ -25,8 +25,8 @@ from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker from diffusers.utils import load_image from parameterized import parameterized +from testing_utils import MODEL_NAMES, SEED, ORTModelTestMixin from transformers.testing_utils import require_torch_gpu -from utils_onnxruntime_tests import MODEL_NAMES, SEED, ORTModelTestMixin from optimum.onnxruntime import ( ORTDiffusionPipeline, @@ -34,6 +34,7 @@ ORTPipelineForInpainting, ORTPipelineForText2Image, ) +from optimum.utils import is_transformers_version from optimum.utils.testing_utils import grid_parameters, require_diffusers @@ -53,6 +54,7 @@ def _generate_prompts(batch_size=1): "guidance_scale": 7.5, "output_type": "np", } + return inputs @@ -71,7 +73,30 @@ def _generate_images(height=128, width=128, batch_size=1, channel=3, input_type= class ORTPipelineForText2ImageTest(ORTModelTestMixin): - SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"] + SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + "latent-consistency", + ] + if is_transformers_version(">=", "4.45"): + SUPPORTED_ARCHITECTURES += ["stable-diffusion-3", "flux"] + + NEGATIVE_PROMPT_SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + "latent-consistency", + ] + + if is_transformers_version(">=", "4.45"): + NEGATIVE_PROMPT_SUPPORTED_ARCHITECTURES += ["stable-diffusion-3"] + + CALLBACK_SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + "latent-consistency", + ] + if is_transformers_version(">=", "4.45"): + CALLBACK_SUPPORTED_ARCHITECTURES += ["flux"] ORTMODEL_CLASS = ORTPipelineForText2Image AUTOMODEL_CLASS = AutoPipelineForText2Image @@ -81,8 +106,7 @@ class ORTPipelineForText2ImageTest(ORTModelTestMixin): def generate_inputs(self, height=128, width=128, batch_size=1): inputs = _generate_prompts(batch_size=batch_size) - inputs["height"] = height - inputs["width"] = width + inputs["height"], inputs["width"] = height, width return inputs @@ -120,8 +144,8 @@ def test_num_images_per_prompt(self, model_arch: str): pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) for batch_size in [1, 3]: - for height in [64, 128]: - for width in [64, 128]: + for height in [16, 32]: + for width in [16, 32]: for num_images_per_prompt in [1, 3]: inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images @@ -142,12 +166,12 @@ def test_compare_to_diffusers_pipeline(self, model_arch: str): for output_type in ["latent", "np", "pt"]: inputs["output_type"] = output_type - ort_output = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images - diffusers_output = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images + ort_images = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images + diffusers_images = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images - np.testing.assert_allclose(ort_output, diffusers_output, atol=1e-4, rtol=1e-2) + np.testing.assert_allclose(ort_images, diffusers_images, atol=1e-4, rtol=1e-2) - @parameterized.expand(SUPPORTED_ARCHITECTURES) + @parameterized.expand(CALLBACK_SUPPORTED_ARCHITECTURES) @require_diffusers def test_callback(self, model_arch: str): model_args = {"test_name": model_arch, "model_arch": model_arch} @@ -164,6 +188,7 @@ def __init__(self): def __call__(self, *args, **kwargs) -> None: self.has_been_called = True self.number_of_steps += 1 + return kwargs ort_callback = Callback() auto_callback = Callback() @@ -171,9 +196,8 @@ def __call__(self, *args, **kwargs) -> None: ort_pipe = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) auto_pipe = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) - # callback_steps=1 to trigger callback every step - ort_pipe(**inputs, callback=ort_callback, callback_steps=1) - auto_pipe(**inputs, callback=auto_callback, callback_steps=1) + ort_pipe(**inputs, callback_on_step_end=ort_callback) + auto_pipe(**inputs, callback_on_step_end=auto_callback) self.assertTrue(ort_callback.has_been_called) self.assertTrue(auto_callback.has_been_called) @@ -200,10 +224,22 @@ def test_shape(self, model_arch: str): elif output_type == "pt": self.assertEqual(outputs.shape, (batch_size, 3, height, width)) else: - self.assertEqual( - outputs.shape, - (batch_size, 4, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor), - ) + if model_arch == "flux": + expected_height = height // (pipeline.vae_scale_factor * 2) + expected_width = width // (pipeline.vae_scale_factor * 2) + channels = pipeline.transformer.config.in_channels + expected_shape = (batch_size, expected_height * expected_width, channels) + else: + expected_height = height // pipeline.vae_scale_factor + expected_width = width // pipeline.vae_scale_factor + out_channels = ( + pipeline.unet.config.out_channels + if getattr(pipeline, "unet", None) is not None + else pipeline.transformer.config.out_channels + ) + expected_shape = (batch_size, out_channels, expected_height, expected_width) + + self.assertEqual(outputs.shape, expected_shape) @parameterized.expand(SUPPORTED_ARCHITECTURES) @require_diffusers @@ -224,60 +260,39 @@ def test_image_reproducibility(self, model_arch: str): self.assertFalse(np.array_equal(ort_outputs_1.images[0], ort_outputs_3.images[0])) np.testing.assert_allclose(ort_outputs_1.images[0], ort_outputs_2.images[0], atol=1e-4, rtol=1e-2) - @parameterized.expand(SUPPORTED_ARCHITECTURES) + @parameterized.expand(NEGATIVE_PROMPT_SUPPORTED_ARCHITECTURES) def test_negative_prompt(self, model_arch: str): model_args = {"test_name": model_arch, "model_arch": model_arch} self._setup(model_args) height, width, batch_size = 64, 64, 1 inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) + inputs["negative_prompt"] = ["This is a negative prompt"] * batch_size - negative_prompt = ["This is a negative prompt"] - pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) + ort_pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) + diffusers_pipeline = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) - images_1 = pipeline(**inputs, negative_prompt=negative_prompt, generator=get_generator("pt", SEED)).images - prompt = inputs.pop("prompt") - - if model_arch == "stable-diffusion-xl": - ( - inputs["prompt_embeds"], - inputs["negative_prompt_embeds"], - inputs["pooled_prompt_embeds"], - inputs["negative_pooled_prompt_embeds"], - ) = pipeline.encode_prompt( - prompt=prompt, - num_images_per_prompt=1, - device=torch.device("cpu"), - do_classifier_free_guidance=True, - negative_prompt=negative_prompt, - ) - else: - inputs["prompt_embeds"], inputs["negative_prompt_embeds"] = pipeline.encode_prompt( - prompt=prompt, - num_images_per_prompt=1, - device=torch.device("cpu"), - do_classifier_free_guidance=True, - negative_prompt=negative_prompt, - ) - - images_2 = pipeline(**inputs, generator=get_generator("pt", SEED)).images - - np.testing.assert_allclose(images_1, images_2, atol=1e-4, rtol=1e-2) + ort_images = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images + diffusers_images = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images + + np.testing.assert_allclose(ort_images, diffusers_images, atol=1e-4, rtol=1e-2) @parameterized.expand( grid_parameters( { "model_arch": SUPPORTED_ARCHITECTURES, - "provider": ["CUDAExecutionProvider", "ROCMExecutionProvider", "TensorrtExecutionProvider"], + "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"], } ) ) - @pytest.mark.rocm_ep_test @pytest.mark.cuda_ep_test @pytest.mark.trt_ep_test @require_torch_gpu @require_diffusers def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str): + if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]: + self.skipTest("Testing a single arch for TensorrtExecutionProvider") + model_args = {"test_name": test_name, "model_arch": model_arch} self._setup(model_args) @@ -285,9 +300,9 @@ def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str): inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[test_name], provider=provider) + self.assertEqual(pipeline.device.type, "cuda") outputs = pipeline(**inputs).images - self.assertIsInstance(outputs, np.ndarray) self.assertEqual(outputs.shape, (batch_size, height, width, 3)) @@ -326,7 +341,19 @@ def test_safety_checker(self, model_arch: str): class ORTPipelineForImage2ImageTest(ORTModelTestMixin): - SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl", "latent-consistency"] + SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + "latent-consistency", + ] + if is_transformers_version(">=", "4.45"): + SUPPORTED_ARCHITECTURES += ["stable-diffusion-3"] + + CALLBACK_SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + "latent-consistency", + ] AUTOMODEL_CLASS = AutoPipelineForImage2Image ORTMODEL_CLASS = ORTPipelineForImage2Image @@ -340,6 +367,7 @@ def generate_inputs(self, height=128, width=128, batch_size=1, channel=3, input_ height=height, width=width, batch_size=batch_size, channel=channel, input_type=input_type ) + inputs["height"], inputs["width"] = height, width inputs["strength"] = 0.75 return inputs @@ -373,14 +401,14 @@ def test_num_images_per_prompt(self, model_arch: str): pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) for batch_size in [1, 3]: - for height in [64, 128]: - for width in [64, 128]: + for height in [16, 32]: + for width in [16, 32]: for num_images_per_prompt in [1, 3]: inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images self.assertEqual(outputs.shape, (batch_size * num_images_per_prompt, height, width, 3)) - @parameterized.expand(SUPPORTED_ARCHITECTURES) + @parameterized.expand(CALLBACK_SUPPORTED_ARCHITECTURES) @require_diffusers def test_callback(self, model_arch: str): model_args = {"test_name": model_arch, "model_arch": model_arch} @@ -398,15 +426,16 @@ def __init__(self): def __call__(self, *args, **kwargs) -> None: self.has_been_called = True self.number_of_steps += 1 + return kwargs ort_pipe = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) auto_pipe = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) ort_callback = Callback() auto_callback = Callback() - # callback_steps=1 to trigger callback every step - ort_pipe(**inputs, callback=ort_callback, callback_steps=1) - auto_pipe(**inputs, callback=auto_callback, callback_steps=1) + + ort_pipe(**inputs, callback_on_step_end=ort_callback) + auto_pipe(**inputs, callback_on_step_end=auto_callback) self.assertTrue(ort_callback.has_been_called) self.assertEqual(ort_callback.number_of_steps, auto_callback.number_of_steps) @@ -434,9 +463,19 @@ def test_shape(self, model_arch: str): elif output_type == "pt": self.assertEqual(outputs.shape, (batch_size, 3, height, width)) else: + out_channels = ( + pipeline.unet.config.out_channels + if pipeline.unet is not None + else pipeline.transformer.config.out_channels + ) self.assertEqual( outputs.shape, - (batch_size, 4, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor), + ( + batch_size, + out_channels, + height // pipeline.vae_scale_factor, + width // pipeline.vae_scale_factor, + ), ) @parameterized.expand(SUPPORTED_ARCHITECTURES) @@ -454,10 +493,10 @@ def test_compare_to_diffusers_pipeline(self, model_arch: str): for output_type in ["latent", "np", "pt"]: inputs["output_type"] = output_type - ort_output = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images - diffusers_output = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images + ort_images = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images + diffusers_images = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images - np.testing.assert_allclose(ort_output, diffusers_output, atol=1e-4, rtol=1e-2) + np.testing.assert_allclose(ort_images, diffusers_images, atol=1e-4, rtol=1e-2) @parameterized.expand(SUPPORTED_ARCHITECTURES) @require_diffusers @@ -482,16 +521,18 @@ def test_image_reproducibility(self, model_arch: str): grid_parameters( { "model_arch": SUPPORTED_ARCHITECTURES, - "provider": ["CUDAExecutionProvider", "ROCMExecutionProvider", "TensorrtExecutionProvider"], + "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"], } ) ) - @pytest.mark.rocm_ep_test @pytest.mark.cuda_ep_test @pytest.mark.trt_ep_test @require_torch_gpu @require_diffusers def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str): + if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]: + self.skipTest("Testing a single arch for TensorrtExecutionProvider") + model_args = {"test_name": test_name, "model_arch": model_arch} self._setup(model_args) @@ -541,7 +582,17 @@ def test_safety_checker(self, model_arch: str): class ORTPipelineForInpaintingTest(ORTModelTestMixin): - SUPPORTED_ARCHITECTURES = ["stable-diffusion", "stable-diffusion-xl"] + SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + ] + if is_transformers_version(">=", "4.45"): + SUPPORTED_ARCHITECTURES += ["stable-diffusion-3"] + + CALLBACK_SUPPORTED_ARCHITECTURES = [ + "stable-diffusion", + "stable-diffusion-xl", + ] AUTOMODEL_CLASS = AutoPipelineForInpainting ORTMODEL_CLASS = ORTPipelineForInpainting @@ -558,9 +609,8 @@ def generate_inputs(self, height=128, width=128, batch_size=1, channel=3, input_ height=height, width=width, batch_size=batch_size, channel=1, input_type=input_type ) + inputs["height"], inputs["width"] = height, width inputs["strength"] = 0.75 - inputs["height"] = height - inputs["width"] = width return inputs @@ -593,14 +643,14 @@ def test_num_images_per_prompt(self, model_arch: str): pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) for batch_size in [1, 3]: - for height in [64, 128]: - for width in [64, 128]: + for height in [16, 32]: + for width in [16, 32]: for num_images_per_prompt in [1, 3]: inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images self.assertEqual(outputs.shape, (batch_size * num_images_per_prompt, height, width, 3)) - @parameterized.expand(SUPPORTED_ARCHITECTURES) + @parameterized.expand(CALLBACK_SUPPORTED_ARCHITECTURES) @require_diffusers def test_callback(self, model_arch: str): model_args = {"test_name": model_arch, "model_arch": model_arch} @@ -618,15 +668,16 @@ def __init__(self): def __call__(self, *args, **kwargs) -> None: self.has_been_called = True self.number_of_steps += 1 + return kwargs ort_pipe = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch]) auto_pipe = self.AUTOMODEL_CLASS.from_pretrained(MODEL_NAMES[model_arch]) ort_callback = Callback() auto_callback = Callback() - # callback_steps=1 to trigger callback every step - ort_pipe(**inputs, callback=ort_callback, callback_steps=1) - auto_pipe(**inputs, callback=auto_callback, callback_steps=1) + + ort_pipe(**inputs, callback_on_step_end=ort_callback) + auto_pipe(**inputs, callback_on_step_end=auto_callback) self.assertTrue(ort_callback.has_been_called) self.assertEqual(ort_callback.number_of_steps, auto_callback.number_of_steps) @@ -654,9 +705,19 @@ def test_shape(self, model_arch: str): elif output_type == "pt": self.assertEqual(outputs.shape, (batch_size, 3, height, width)) else: + out_channels = ( + pipeline.unet.config.out_channels + if pipeline.unet is not None + else pipeline.transformer.config.out_channels + ) self.assertEqual( outputs.shape, - (batch_size, 4, height // pipeline.vae_scale_factor, width // pipeline.vae_scale_factor), + ( + batch_size, + out_channels, + height // pipeline.vae_scale_factor, + width // pipeline.vae_scale_factor, + ), ) @parameterized.expand(SUPPORTED_ARCHITECTURES) @@ -674,10 +735,10 @@ def test_compare_to_diffusers_pipeline(self, model_arch: str): for output_type in ["latent", "np", "pt"]: inputs["output_type"] = output_type - ort_output = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images - diffusers_output = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images + ort_images = ort_pipeline(**inputs, generator=get_generator("pt", SEED)).images + diffusers_images = diffusers_pipeline(**inputs, generator=get_generator("pt", SEED)).images - np.testing.assert_allclose(ort_output, diffusers_output, atol=1e-4, rtol=1e-2) + np.testing.assert_allclose(ort_images, diffusers_images, atol=1e-4, rtol=1e-2) @parameterized.expand(SUPPORTED_ARCHITECTURES) @require_diffusers @@ -702,16 +763,18 @@ def test_image_reproducibility(self, model_arch: str): grid_parameters( { "model_arch": SUPPORTED_ARCHITECTURES, - "provider": ["CUDAExecutionProvider", "ROCMExecutionProvider", "TensorrtExecutionProvider"], + "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"], } ) ) - @pytest.mark.rocm_ep_test @pytest.mark.cuda_ep_test @pytest.mark.trt_ep_test @require_torch_gpu @require_diffusers def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str): + if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]: + self.skipTest("Testing a single arch for TensorrtExecutionProvider") + model_args = {"test_name": test_name, "model_arch": model_arch} self._setup(model_args) @@ -719,7 +782,7 @@ def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str): inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size) pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[test_name], provider=provider) - self.assertEqual(pipeline.device, "cuda") + self.assertEqual(pipeline.device.type, "cuda") outputs = pipeline(**inputs).images self.assertIsInstance(outputs, np.ndarray) diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py index 597eb581e2..9ea0483e35 100644 --- a/tests/onnxruntime/test_modeling.py +++ b/tests/onnxruntime/test_modeling.py @@ -16,7 +16,6 @@ import os import subprocess import tempfile -import time import unittest from pathlib import Path from typing import Dict @@ -26,12 +25,12 @@ import onnxruntime import pytest import requests -import timm import torch from huggingface_hub import HfApi from huggingface_hub.constants import default_cache_path from parameterized import parameterized from PIL import Image +from testing_utils import MODEL_NAMES, SEED, ORTModelTestMixin from transformers import ( AutoConfig, AutoFeatureExtractor, @@ -65,7 +64,6 @@ from transformers.models.swin2sr.configuration_swin2sr import Swin2SRConfig from transformers.onnx.utils import get_preprocessor from transformers.testing_utils import get_gpu_count, require_torch_gpu, slow -from utils_onnxruntime_tests import MODEL_NAMES, SEED, ORTModelTestMixin from optimum.exporters import TasksManager from optimum.exporters.onnx import MODEL_TYPES_REQUIRING_POSITION_IDS, main_export @@ -107,7 +105,7 @@ DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER, logging, ) -from optimum.utils.import_utils import check_if_transformers_greater, is_diffusers_available +from optimum.utils.import_utils import is_diffusers_available, is_transformers_version from optimum.utils.testing_utils import ( grid_parameters, remove_directory, @@ -130,21 +128,12 @@ logger = logging.get_logger() -class Timer(object): - def __enter__(self): - self.elapsed = time.perf_counter() - return self - - def __exit__(self, type, value, traceback): - self.elapsed = (time.perf_counter() - self.elapsed) * 1e3 - - class ORTModelIntegrationTest(unittest.TestCase): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.TEST_MODEL_ID = "sshleifer/tiny-distilbert-base-cased-distilled-squad" - self.LOCAL_MODEL_PATH = "assets/onnx" + self.LOCAL_MODEL_PATH = "tests/assets/onnx" self.ONNX_MODEL_ID = "philschmid/distilbert-onnx" + self.TINY_ONNX_MODEL_ID = "fxmarty/resnet-tiny-beans" self.FAIL_ONNX_MODEL_ID = "sshleifer/tiny-distilbert-base-cased-distilled-squad" self.ONNX_SEQ2SEQ_MODEL_ID = "optimum/t5-small" @@ -255,6 +244,16 @@ def test_load_model_cuda_provider(self): self.assertListEqual(model.model.get_providers(), model.providers) self.assertEqual(model.device, torch.device("cuda:0")) + @require_torch_gpu + @pytest.mark.trt_ep_test + def test_load_model_tensorrt_provider(self): + model = ORTModel.from_pretrained(self.ONNX_MODEL_ID, provider="TensorrtExecutionProvider") + self.assertListEqual( + model.providers, ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"] + ) + self.assertListEqual(model.model.get_providers(), model.providers) + self.assertEqual(model.device, torch.device("cuda:0")) + @require_torch_gpu @require_ort_rocm @pytest.mark.rocm_ep_test @@ -774,7 +773,6 @@ def test_seq2seq_model_on_gpu_id(self): model.decoder_with_past.session.get_provider_options()["CUDAExecutionProvider"]["device_id"], "1" ) - # test string device input for to() @require_torch_gpu @pytest.mark.cuda_ep_test def test_seq2seq_model_on_gpu_str(self): @@ -974,7 +972,7 @@ def test_stable_diffusion_model_on_rocm_ep_str(self): def test_load_model_from_hub_private(self): token = os.environ.get("HF_HUB_READ_TOKEN", None) - if token is None: + if not token: self.skipTest( "Test requires a read access token for optimum-internal-testing in the environment variable `HF_HUB_READ_TOKEN`." ) @@ -1054,7 +1052,7 @@ def test_save_load_ort_model_with_external_data(self): def test_save_load_decoder_model_with_external_data(self, use_cache: bool): with tempfile.TemporaryDirectory() as tmpdirname: model = ORTModelForCausalLM.from_pretrained( - "gpt2-large", use_cache=use_cache, export=True, use_merged=False, use_io_binding=False + "gpt2-large", export=True, use_cache=use_cache, use_merged=False, use_io_binding=False ) model.save_pretrained(tmpdirname) @@ -1265,9 +1263,7 @@ def test_trust_remote_code(self): ort_logits = ort_model(**inputs).logits - self.assertTrue( - torch.allclose(pt_logits, ort_logits, atol=1e-4), f" Maxdiff: {torch.abs(pt_logits - ort_logits).max()}" - ) + torch.testing.assert_close(pt_logits, ort_logits, atol=1e-4, rtol=1e-4) @parameterized.expand(("", "onnx")) def test_loading_with_config_not_from_subfolder(self, subfolder): @@ -1289,8 +1285,8 @@ class ORTModelForQuestionAnsweringIntegrationTest(ORTModelTestMixin): "albert", "bart", "bert", - # "big_bird", - # "bigbird_pegasus", + "big_bird", + "bigbird_pegasus", "camembert", "convbert", "data2vec_text", @@ -1312,6 +1308,7 @@ class ORTModelForQuestionAnsweringIntegrationTest(ORTModelTestMixin): "squeezebert", "xlm_qa", "xlm_roberta", + "rembert", ] FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES} @@ -1353,11 +1350,14 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.end_logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # Compare tensor outputs - self.assertTrue( - torch.allclose(torch.Tensor(onnx_outputs.start_logits), transformers_outputs.start_logits, atol=1e-4) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.start_logits), + transformers_outputs.start_logits, + atol=self.ATOL, + rtol=self.RTOL, ) - self.assertTrue( - torch.allclose(torch.Tensor(onnx_outputs.end_logits), transformers_outputs.end_logits, atol=1e-4) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.end_logits), transformers_outputs.end_logits, atol=self.ATOL, rtol=self.RTOL ) gc.collect() @@ -1457,14 +1457,18 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForQuestionAnswering.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) io_model = ORTModelForQuestionAnswering.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) tokenizer = get_preprocessor(model_id) - tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt") + tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**tokens) io_outputs = io_model(**tokens) @@ -1474,8 +1478,12 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.end_logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.start_logits, io_outputs.start_logits)) - self.assertTrue(torch.equal(onnx_outputs.end_logits, io_outputs.end_logits)) + torch.testing.assert_close( + torch.Tensor(io_outputs.start_logits), onnx_outputs.start_logits, atol=self.ATOL, rtol=self.RTOL + ) + torch.testing.assert_close( + torch.Tensor(io_outputs.end_logits), onnx_outputs.end_logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -1484,7 +1492,7 @@ class ORTModelForMaskedLMIntegrationTest(ORTModelTestMixin): SUPPORTED_ARCHITECTURES = [ "albert", "bert", - # "big_bird", + "big_bird", "camembert", "convbert", "data2vec_text", @@ -1502,6 +1510,7 @@ class ORTModelForMaskedLMIntegrationTest(ORTModelTestMixin): "squeezebert", "xlm", "xlm_roberta", + "rembert", ] FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES} @@ -1542,7 +1551,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -1628,16 +1639,19 @@ def test_compare_to_io_binding(self, model_arch): self._setup(model_args) model_id = MODEL_NAMES[model_arch] - onnx_model = ORTModelForMaskedLM.from_pretrained(self.onnx_model_dirs[model_arch], use_io_binding=False).to( - "cuda" + onnx_model = ORTModelForMaskedLM.from_pretrained( + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" ) - io_model = ORTModelForMaskedLM.from_pretrained(self.onnx_model_dirs[model_arch], use_io_binding=True).to( - "cuda" + io_model = ORTModelForMaskedLM.from_pretrained( + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" ) + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) + tokenizer = get_preprocessor(model_id) - MASK_TOKEN = tokenizer.mask_token - tokens = tokenizer([f"The capital of France is {MASK_TOKEN}."] * 2, return_tensors="pt") + tokens = tokenizer([f"The capital of France is {tokenizer.mask_token}."] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**tokens) io_outputs = io_model(**tokens) @@ -1645,7 +1659,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(io_outputs.logits, onnx_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -1655,8 +1669,8 @@ class ORTModelForSequenceClassificationIntegrationTest(ORTModelTestMixin): "albert", "bart", "bert", - # "big_bird", - # "bigbird_pegasus", + "big_bird", + "bigbird_pegasus", "bloom", "camembert", "convbert", @@ -1682,6 +1696,7 @@ class ORTModelForSequenceClassificationIntegrationTest(ORTModelTestMixin): "squeezebert", "xlm", "xlm_roberta", + "rembert", ] FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES} @@ -1722,7 +1737,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -1837,14 +1854,18 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForSequenceClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) io_model = ORTModelForSequenceClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) tokenizer = get_preprocessor(model_id) - tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt") + tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**tokens) io_outputs = io_model(**tokens) @@ -1852,7 +1873,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -1861,7 +1882,7 @@ class ORTModelForTokenClassificationIntegrationTest(ORTModelTestMixin): SUPPORTED_ARCHITECTURES = [ "albert", "bert", - # "big_bird", + "big_bird", "bloom", "camembert", "convbert", @@ -1882,6 +1903,7 @@ class ORTModelForTokenClassificationIntegrationTest(ORTModelTestMixin): "squeezebert", "xlm", "xlm_roberta", + "rembert", ] FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES} @@ -1922,7 +1944,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -2016,14 +2040,18 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForTokenClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) io_model = ORTModelForTokenClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) tokenizer = get_preprocessor(model_id) - tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt") + tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**tokens) io_outputs = io_model(**tokens) @@ -2031,7 +2059,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -2079,10 +2107,11 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.last_hidden_state, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue( - torch.allclose( - torch.Tensor(onnx_outputs.last_hidden_state), transformers_outputs.last_hidden_state, atol=1e-4 - ) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.last_hidden_state), + transformers_outputs.last_hidden_state, + atol=self.ATOL, + rtol=self.RTOL, ) gc.collect() @@ -2174,14 +2203,18 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForFeatureExtraction.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) io_model = ORTModelForFeatureExtraction.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) tokenizer = get_preprocessor(model_id) - tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt") + tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**tokens) io_outputs = io_model(**tokens) @@ -2189,7 +2222,9 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.last_hidden_state, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.last_hidden_state, io_outputs.last_hidden_state)) + torch.testing.assert_close( + onnx_outputs.last_hidden_state, io_outputs.last_hidden_state, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -2202,7 +2237,9 @@ def test_default_token_type_ids(self): token_type_ids = tokens.pop("token_type_ids") outs = model(token_type_ids=token_type_ids, **tokens) outs_without_token_type_ids = model(**tokens) - self.assertTrue(np.allclose(outs.last_hidden_state, outs_without_token_type_ids.last_hidden_state)) + torch.testing.assert_close( + outs.last_hidden_state, outs_without_token_type_ids.last_hidden_state, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -2211,7 +2248,7 @@ class ORTModelForMultipleChoiceIntegrationTest(ORTModelTestMixin): SUPPORTED_ARCHITECTURES = [ "albert", "bert", - # "big_bird", + "big_bird", "camembert", "convbert", "data2vec_text", @@ -2227,6 +2264,7 @@ class ORTModelForMultipleChoiceIntegrationTest(ORTModelTestMixin): "squeezebert", "xlm", "xlm_roberta", + "rembert", ] FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES} @@ -2269,7 +2307,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # Compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -2282,24 +2322,25 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForMultipleChoice.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") - io_model = ORTModelForMultipleChoice.from_pretrained(self.onnx_model_dirs[model_arch], use_io_binding=True).to( - "cuda" + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) + io_model = ORTModelForMultipleChoice.from_pretrained( + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" ) - tokenizer = get_preprocessor(model_id) + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) + num_choices = 4 - first_sentence = ["The sky is blue due to the shorter wavelength of blue light."] * num_choices start = "The color of the sky is" + tokenizer = get_preprocessor(model_id) + first_sentence = ["The sky is blue due to the shorter wavelength of blue light."] * num_choices second_sentence = [start + "blue", start + "green", start + "red", start + "yellow"] inputs = tokenizer(first_sentence, second_sentence, truncation=True, padding=True) - # Unflatten the tokenized inputs values expanding it to the shape [batch_size, num_choices, seq_length] for k, v in inputs.items(): inputs[k] = [v[i : i + num_choices] for i in range(0, len(v), num_choices)] - - inputs = dict(inputs.convert_to_tensors(tensor_type="pt")) + inputs = dict(inputs.convert_to_tensors(tensor_type="pt").to("cuda")) onnx_outputs = onnx_model(**inputs) io_outputs = io_model(**inputs) @@ -2308,7 +2349,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(io_outputs.logits, onnx_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -2318,7 +2359,6 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin): "bloom", "codegen", "falcon", - "gemma", "gpt2", "gpt_bigcode", "gpt_neo", @@ -2326,12 +2366,21 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin): "gptj", "llama", "mistral", - "mpt", "opt", ] - if check_if_transformers_greater("4.40"): - SUPPORTED_ARCHITECTURES.extend(["gemma", "phi3", "qwen2"]) + if is_transformers_version(">=", "4.37"): + SUPPORTED_ARCHITECTURES.append("qwen2") + + if is_transformers_version(">=", "4.38"): + SUPPORTED_ARCHITECTURES.append("gemma") + + # TODO: fix "mpt" for which inference fails for transformers < v4.41 + if is_transformers_version(">=", "4.41"): + SUPPORTED_ARCHITECTURES.extend(["phi3", "mpt"]) + + if is_transformers_version(">=", "4.45"): + SUPPORTED_ARCHITECTURES.append("granite") FULL_GRID = { "model_arch": SUPPORTED_ARCHITECTURES, @@ -2341,8 +2390,7 @@ class ORTModelForCausalLMIntegrationTest(ORTModelTestMixin): ORTMODEL_CLASS = ORTModelForCausalLM TASK = "text-generation" - GENERATION_LENGTH = 90 - SPEEDUP_CACHE = 1.1 + GENERATION_LENGTH = 100 @parameterized.expand([(False,), (True,)]) @pytest.mark.run_in_series @@ -2444,7 +2492,7 @@ def test_compare_to_transformers(self, test_name: str, model_arch: str, use_cach transformers_model = AutoModelForCausalLM.from_pretrained(model_id) transformers_model = transformers_model.eval() tokenizer = get_preprocessor(model_id) - tokens = tokenizer("This is a sample output", return_tensors="pt") + tokens = tokenizer("This is a sample input", return_tensors="pt") position_ids = None if model_arch.replace("_", "-") in MODEL_TYPES_REQUIRING_POSITION_IDS: input_shape = tokens["input_ids"].shape @@ -2458,15 +2506,12 @@ def test_compare_to_transformers(self, test_name: str, model_arch: str, use_cach self.assertIsInstance(onnx_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue( - torch.allclose(onnx_outputs.logits, transformers_outputs.logits, atol=1e-4), - f"Maxdiff: {(onnx_outputs.logits - transformers_outputs.logits).abs()}", - ) + torch.testing.assert_close(onnx_outputs.logits, transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL) # Compare batched generation. tokenizer.pad_token_id = tokenizer.eos_token_id tokenizer.padding_side = "left" - tokens = tokenizer(["Today is a nice day and I am longer", "This is me"], return_tensors="pt", padding=True) + tokens = tokenizer(["This is", "This is a sample input"], return_tensors="pt", padding=True) onnx_model.generation_config.eos_token_id = None transformers_model.generation_config.eos_token_id = None onnx_model.config.eos_token_id = None @@ -2503,13 +2548,11 @@ def test_compare_to_transformers(self, test_name: str, model_arch: str, use_cach set_seed(SEED) with torch.no_grad(): transformers_outputs = transformers_model.generate(**tokens, generation_config=gen_config) + set_seed(SEED) onnx_outputs = onnx_model.generate(**tokens, generation_config=gen_config) - self.assertTrue( - torch.equal(onnx_outputs, transformers_outputs), - f"Failed with generation config : {gen_config}, transformers outputs {transformers_outputs}, ONNX model outputs {onnx_outputs}", - ) + torch.testing.assert_close(onnx_outputs, transformers_outputs, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -2652,7 +2695,6 @@ def test_pipeline_on_trt_execution_provider(self, test_name: str, model_arch: st gc.collect() @parameterized.expand(SUPPORTED_ARCHITECTURES) - @pytest.mark.cuda_ep_test # mark as GPU test as well to run the without/with cache timing test on the slow tests def test_compare_with_and_without_past_key_values(self, model_arch): model_args = {"test_name": model_arch + "_False", "model_arch": model_arch, "use_cache": False} self._setup(model_args) @@ -2664,34 +2706,25 @@ def test_compare_with_and_without_past_key_values(self, model_arch): text = "My Name is Philipp and i live" tokens = tokenizer(text, return_tensors="pt", return_token_type_ids=False if model_arch == "llama" else None) + generation_length = 10 # model has a short max length + model_with_pkv = ORTModelForCausalLM.from_pretrained( self.onnx_model_dirs[model_arch + "_True"], use_cache=True, use_io_binding=False ) - _ = model_with_pkv.generate(**tokens) # warmup - with Timer() as with_pkv_timer: - outputs_model_with_pkv = model_with_pkv.generate( - **tokens, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 - ) + outputs_model_with_pkv = model_with_pkv.generate( + **tokens, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 + ) model_without_pkv = ORTModelForCausalLM.from_pretrained( self.onnx_model_dirs[model_arch + "_False"], use_cache=False, use_io_binding=False ) - _ = model_without_pkv.generate(**tokens) # warmup - with Timer() as without_pkv_timer: - outputs_model_without_pkv = model_without_pkv.generate( - **tokens, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 - ) - - self.assertTrue(torch.equal(outputs_model_with_pkv, outputs_model_without_pkv)) - self.assertEqual(outputs_model_with_pkv.shape[1], tokens["input_ids"].shape[1] + self.GENERATION_LENGTH) - self.assertEqual(outputs_model_without_pkv.shape[1], tokens["input_ids"].shape[1] + self.GENERATION_LENGTH) + outputs_model_without_pkv = model_without_pkv.generate( + **tokens, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 + ) - if os.environ.get("TEST_LEVEL", 0) == "1": - self.assertTrue( - without_pkv_timer.elapsed / with_pkv_timer.elapsed > self.SPEEDUP_CACHE, - f"With pkv latency: {with_pkv_timer.elapsed:.3f} ms, without pkv latency: {without_pkv_timer.elapsed:.3f} ms," - f" speedup: {without_pkv_timer.elapsed / with_pkv_timer.elapsed:.3f}", - ) + torch.testing.assert_close(outputs_model_with_pkv, outputs_model_without_pkv, atol=self.ATOL, rtol=self.RTOL) + self.assertEqual(outputs_model_with_pkv.shape[1], tokens["input_ids"].shape[1] + generation_length) + self.assertEqual(outputs_model_without_pkv.shape[1], tokens["input_ids"].shape[1] + generation_length) @parameterized.expand(grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True]})) def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, model_arch: str, use_cache: bool): @@ -2734,7 +2767,7 @@ def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, mode outputs_model_not_merged = model_not_merged.generate(**tokens) outputs_model_merged = model_merged.generate(**tokens) - self.assertTrue(torch.equal(outputs_model_merged, outputs_model_not_merged)) + torch.testing.assert_close(outputs_model_not_merged, outputs_model_merged, atol=self.ATOL, rtol=self.RTOL) @parameterized.expand( grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True], "use_merged": [False, True]}) @@ -2752,11 +2785,17 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForCausalLM.from_pretrained( - self.onnx_model_dirs[test_name], use_cache=use_cache, use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[test_name], + use_cache=use_cache, + use_io_binding=False, + provider="CUDAExecutionProvider", + ) io_model = ORTModelForCausalLM.from_pretrained( - self.onnx_model_dirs[test_name], use_cache=use_cache, use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[test_name], + use_cache=use_cache, + use_io_binding=True, + provider="CUDAExecutionProvider", + ) tokenizer = get_preprocessor(model_id) tokens = tokenizer(["This is a sample output"] * 2, return_tensors="pt").to("cuda") @@ -2775,7 +2814,7 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(io_outputs.logits, onnx_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -2787,10 +2826,15 @@ def test_compare_generation_to_io_binding(self, test_name: str, model_arch: str, self._setup(model_args) model_id = MODEL_NAMES[model_arch] - onnx_model = ORTModelForCausalLM.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=False).to( - "cuda" + onnx_model = ORTModelForCausalLM.from_pretrained( + self.onnx_model_dirs[test_name], use_io_binding=False, provider="CUDAExecutionProvider" + ) + io_model = ORTModelForCausalLM.from_pretrained( + self.onnx_model_dirs[test_name], use_io_binding=True, provider="CUDAExecutionProvider" ) - io_model = ORTModelForCausalLM.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=True).to("cuda") + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) tokenizer = get_preprocessor(model_id) tokens = tokenizer( @@ -2798,11 +2842,12 @@ def test_compare_generation_to_io_binding(self, test_name: str, model_arch: str, return_tensors="pt", return_token_type_ids=False if model_arch == "llama" else None, ).to("cuda") + onnx_outputs = onnx_model.generate(**tokens) io_outputs = io_model.generate(**tokens) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs, io_outputs)) + torch.testing.assert_close(io_outputs, onnx_outputs, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -2814,6 +2859,7 @@ class ORTModelForImageClassificationIntegrationTest(ORTModelTestMixin): "convnextv2", "data2vec_vision", "deit", + "dinov2", "levit", "mobilenet_v1", "mobilenet_v2", @@ -2827,8 +2873,6 @@ class ORTModelForImageClassificationIntegrationTest(ORTModelTestMixin): "vit", ] - TIMM_SUPPORTED_ARCHITECTURES = ["default-timm-config"] - FULL_GRID = {"model_arch": SUPPORTED_ARCHITECTURES} ORTMODEL_CLASS = ORTModelForImageClassification TASK = "image-classification" @@ -2854,54 +2898,6 @@ def test_load_vanilla_transformers_which_is_not_supported(self): self.assertIn("only supports the tasks", str(context.exception)) - @parameterized.expand(TIMM_SUPPORTED_ARCHITECTURES) - @pytest.mark.run_slow - @pytest.mark.timm_test - @slow - def test_compare_to_timm(self, model_arch): - model_args = {"test_name": model_arch, "model_arch": model_arch} - - self._setup(model_args) - - model_ids = self._get_model_ids(model_arch) - for model_id in model_ids: - onnx_model = ORTModelForImageClassification.from_pretrained( - self._get_onnx_model_dir(model_id, model_arch, model_arch) - ) - - self.assertIsInstance(onnx_model.model, onnxruntime.InferenceSession) - self.assertIsInstance(onnx_model.config, PretrainedConfig) - - set_seed(SEED) - timm_model = timm.create_model(model_id, pretrained=True) - timm_model = timm_model.eval() - - # get model specific transforms (normalization, resize) - data_config = timm.data.resolve_model_data_config(timm_model) - transforms = timm.data.create_transform(**data_config, is_training=False) - - url = ( - "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png" - ) - image = Image.open(requests.get(url, stream=True).raw) - inputs = transforms(image).unsqueeze(0) - - with torch.no_grad(): - timm_outputs = timm_model(inputs) - - for input_type in ["pt", "np"]: - if input_type == "np": - inputs = inputs.cpu().detach().numpy() - onnx_outputs = onnx_model(inputs) - - self.assertIn("logits", onnx_outputs) - self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) - - # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), timm_outputs, atol=1e-4)) - - gc.collect() - @parameterized.expand(SUPPORTED_ARCHITECTURES) def test_compare_to_transformers(self, model_arch): model_args = {"test_name": model_arch, "model_arch": model_arch} @@ -2932,7 +2928,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), trtfs_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), trtfs_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -3032,16 +3030,26 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForImageClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], + use_io_binding=False, + provider="CUDAExecutionProvider", + provider_options={"cudnn_conv_algo_search": "DEFAULT"}, + ) io_model = ORTModelForImageClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], + use_io_binding=True, + provider="CUDAExecutionProvider", + provider_options={"cudnn_conv_algo_search": "DEFAULT"}, + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) preprocessor = get_preprocessor(model_id) url = "http://images.cocodataset.org/val2017/000000039769.jpg" image = Image.open(requests.get(url, stream=True).raw) - inputs = preprocessor(images=[image] * 2, return_tensors="pt") + inputs = preprocessor(images=[image] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**inputs) io_outputs = io_model(**inputs) @@ -3049,10 +3057,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue( - torch.allclose(onnx_outputs.logits, io_outputs.logits, atol=1e-4), - f" Maxdiff: {torch.abs(onnx_outputs.logits - io_outputs.logits).max()}", - ) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -3099,7 +3104,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), trtfs_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), trtfs_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -3197,16 +3204,20 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForSemanticSegmentation.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) io_model = ORTModelForSemanticSegmentation.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) preprocessor = get_preprocessor(model_id) url = "http://images.cocodataset.org/val2017/000000039769.jpg" image = Image.open(requests.get(url, stream=True).raw) - inputs = preprocessor(images=[image] * 2, return_tensors="pt") + inputs = preprocessor(images=[image] * 2, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**inputs) io_outputs = io_model(**inputs) @@ -3214,10 +3225,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue( - torch.allclose(onnx_outputs.logits, io_outputs.logits, atol=1e-4), - f" Maxdiff: {torch.abs(onnx_outputs.logits - io_outputs.logits).max()}", - ) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -3282,7 +3290,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -3381,16 +3391,19 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForAudioClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) io_model = ORTModelForAudioClassification.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=True - ).to("cuda") + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) - processor = AutoFeatureExtractor.from_pretrained(model_id) data = self._generate_random_audio_data() + processor = AutoFeatureExtractor.from_pretrained(model_id) + input_values = processor(data, return_tensors="pt").to("cuda") - input_values = processor(data, return_tensors="pt") onnx_outputs = onnx_model(**input_values) io_outputs = io_model(**input_values) @@ -3398,7 +3411,7 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.allclose(onnx_outputs.logits, io_outputs.logits, atol=1e-4)) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -3461,7 +3474,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -3473,17 +3488,26 @@ def test_compare_to_io_binding(self, model_arch): self._setup(model_args) model_id = MODEL_NAMES[model_arch] - onnx_model = ORTModelForCTC.from_pretrained( self.onnx_model_dirs[model_arch], use_io_binding=False, - ).to("cuda") - onnx_model.use_io_binding = False - io_model = ORTModelForCTC.from_pretrained(self.onnx_model_dirs[model_arch], use_io_binding=True).to("cuda") + provider="CUDAExecutionProvider", + provider_options={"cudnn_conv_algo_search": "DEFAULT"}, + ) + io_model = ORTModelForCTC.from_pretrained( + self.onnx_model_dirs[model_arch], + use_io_binding=True, + provider="CUDAExecutionProvider", + provider_options={"cudnn_conv_algo_search": "DEFAULT"}, + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) - processor = AutoFeatureExtractor.from_pretrained(model_id) data = self._generate_random_audio_data() - input_values = processor(data, return_tensors="pt") + processor = AutoFeatureExtractor.from_pretrained(model_id) + input_values = processor(data, return_tensors="pt").to("cuda") + onnx_outputs = onnx_model(**input_values) io_outputs = io_model(**input_values) @@ -3491,7 +3515,9 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), io_outputs.logits, atol=1e-1)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), io_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -3549,9 +3575,11 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.embeddings, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) - self.assertTrue( - torch.allclose(torch.Tensor(onnx_outputs.embeddings), transformers_outputs.embeddings, atol=1e-4) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.embeddings), transformers_outputs.embeddings, atol=self.ATOL, rtol=self.RTOL ) gc.collect() @@ -3565,16 +3593,19 @@ def test_compare_to_io_binding(self, model_arch): model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForAudioXVector.from_pretrained( - self.onnx_model_dirs[model_arch], use_io_binding=False - ).to("cuda") - io_model = ORTModelForAudioXVector.from_pretrained(self.onnx_model_dirs[model_arch], use_io_binding=True).to( - "cuda" + self.onnx_model_dirs[model_arch], use_io_binding=False, provider="CUDAExecutionProvider" + ) + io_model = ORTModelForAudioXVector.from_pretrained( + self.onnx_model_dirs[model_arch], use_io_binding=True, provider="CUDAExecutionProvider" ) - processor = AutoFeatureExtractor.from_pretrained(model_id) + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) + data = self._generate_random_audio_data() + processor = AutoFeatureExtractor.from_pretrained(model_id) + input_values = processor(data, return_tensors="pt").to("cuda") - input_values = processor(data, return_tensors="pt") onnx_outputs = onnx_model(**input_values) io_outputs = io_model(**input_values) @@ -3583,8 +3614,8 @@ def test_compare_to_io_binding(self, model_arch): self.assertIsInstance(io_outputs.embeddings, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.allclose(onnx_outputs.logits, io_outputs.logits, atol=1e-4)) - self.assertTrue(torch.allclose(onnx_outputs.embeddings, io_outputs.embeddings, atol=1e-4)) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) + torch.testing.assert_close(onnx_outputs.embeddings, io_outputs.embeddings, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -3632,6 +3663,7 @@ def test_compare_to_transformers(self, model_arch): with torch.no_grad(): transformers_outputs = transformers_model(**input_values) + for input_type in ["pt", "np"]: input_values = processor(self._generate_random_audio_data(), return_tensors=input_type) onnx_outputs = onnx_model(**input_values) @@ -3640,7 +3672,9 @@ def test_compare_to_transformers(self, model_arch): self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -3648,7 +3682,7 @@ def test_compare_to_transformers(self, model_arch): class ORTModelForSeq2SeqLMIntegrationTest(ORTModelTestMixin): SUPPORTED_ARCHITECTURES = [ "bart", - # "bigbird_pegasus", + "bigbird_pegasus", "blenderbot", "blenderbot_small", "encoder-decoder", @@ -3671,7 +3705,6 @@ class ORTModelForSeq2SeqLMIntegrationTest(ORTModelTestMixin): TASK = "text2text-generation" GENERATION_LENGTH = 100 - SPEEDUP_CACHE = 1.1 def _get_model_ids(self, model_arch): model_ids = MODEL_NAMES[model_arch] @@ -3875,8 +3908,8 @@ def test_compare_to_transformers(self, test_name: str, model_arch: str, use_cach self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # Compare tensor outputs - self.assertTrue( - torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL ) gc.collect() @@ -4084,15 +4117,17 @@ def test_pipeline_on_trt_execution_provider(self, test_name: str, model_arch: st gc.collect() @parameterized.expand(SUPPORTED_ARCHITECTURES) - @pytest.mark.cuda_ep_test # mark as GPU test as well to run the without/with cache timing test on the slow tests def test_compare_with_and_without_past_key_values(self, model_arch: str): - if model_arch == "m2m_100": - self.skipTest("m2m_100 comparison with/without pkv fail or is not supported") model_args = {"test_name": model_arch + "_False", "model_arch": model_arch, "use_cache": False} self._setup(model_args) model_args = {"test_name": model_arch + "_True", "model_arch": model_arch, "use_cache": True} self._setup(model_args) + if model_arch == "m2m_100": + generation_length = 20 # model's predefined maximum length + else: + generation_length = self.GENERATION_LENGTH + model_ids = self._get_model_ids(model_arch) for model_id in model_ids: if ( @@ -4109,31 +4144,23 @@ def test_compare_with_and_without_past_key_values(self, model_arch: str): self._get_onnx_model_dir(model_id, model_arch, model_arch + "_True"), use_cache=True ) - _ = model_with_pkv.generate(**tokens) # warmup - with Timer() as with_pkv_timer: - outputs_model_with_pkv = model_with_pkv.generate( - **tokens, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 - ) + outputs_model_with_pkv = model_with_pkv.generate( + **tokens, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 + ) model_without_pkv = ORTModelForSeq2SeqLM.from_pretrained( self._get_onnx_model_dir(model_id, model_arch, model_arch + "_False"), use_cache=False ) - _ = model_without_pkv.generate(**tokens) # warmup - with Timer() as without_pkv_timer: - outputs_model_without_pkv = model_without_pkv.generate( - **tokens, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 - ) - self.assertTrue(torch.equal(outputs_model_with_pkv, outputs_model_without_pkv)) - self.assertEqual(outputs_model_with_pkv.shape[1], self.GENERATION_LENGTH + 1) - self.assertEqual(outputs_model_without_pkv.shape[1], self.GENERATION_LENGTH + 1) + outputs_model_without_pkv = model_without_pkv.generate( + **tokens, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 + ) - if os.environ.get("TEST_LEVEL", 0) == "1": - self.assertTrue( - without_pkv_timer.elapsed / with_pkv_timer.elapsed > self.SPEEDUP_CACHE, - f"With pkv latency: {with_pkv_timer.elapsed:.3f} ms, without pkv latency: {without_pkv_timer.elapsed:.3f} ms," - f" speedup: {without_pkv_timer.elapsed / with_pkv_timer.elapsed:.3f}", - ) + torch.testing.assert_close( + outputs_model_with_pkv, outputs_model_without_pkv, rtol=self.RTOL, atol=self.ATOL + ) + self.assertEqual(outputs_model_with_pkv.shape[1], generation_length + 1) + self.assertEqual(outputs_model_without_pkv.shape[1], generation_length + 1) @parameterized.expand(grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True]})) def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, model_arch: str, use_cache: bool): @@ -4182,7 +4209,7 @@ def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, mode outputs_model_not_merged = model_not_merged.generate(**tokens) outputs_model_merged = model_merged.generate(**tokens) - self.assertTrue(torch.equal(outputs_model_merged, outputs_model_not_merged)) + torch.testing.assert_close(outputs_model_not_merged, outputs_model_merged, rtol=self.RTOL, atol=self.ATOL) @parameterized.expand( grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True], "use_merged": [False, True]}) @@ -4212,11 +4239,17 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: continue onnx_model = ORTModelForSeq2SeqLM.from_pretrained( - self._get_onnx_model_dir(model_id, model_arch, test_name), use_io_binding=False, use_cache=use_cache - ).to("cuda") + self._get_onnx_model_dir(model_id, model_arch, test_name), + use_io_binding=False, + use_cache=use_cache, + provider="CUDAExecutionProvider", + ) io_model = ORTModelForSeq2SeqLM.from_pretrained( - self._get_onnx_model_dir(model_id, model_arch, test_name), use_io_binding=True, use_cache=use_cache - ).to("cuda") + self._get_onnx_model_dir(model_id, model_arch, test_name), + use_io_binding=True, + use_cache=use_cache, + provider="CUDAExecutionProvider", + ) self.assertFalse(onnx_model.use_io_binding) self.assertTrue(io_model.use_io_binding) @@ -4226,8 +4259,9 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: decoder_start_token_id = onnx_model.config.decoder_start_token_id if model_arch != "mbart" else 2 if model_arch == "encoder-decoder": decoder_start_token_id = tokenizer.cls_token_id - - decoder_inputs = {"decoder_input_ids": torch.ones((2, 1), dtype=torch.long) * decoder_start_token_id} + decoder_inputs = { + "decoder_input_ids": torch.ones((2, 1), dtype=torch.long).to("cuda") * decoder_start_token_id + } onnx_outputs = onnx_model(**tokens, **decoder_inputs) io_outputs = io_model(**tokens, **decoder_inputs) @@ -4236,7 +4270,7 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -4251,6 +4285,7 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: ) ) @require_torch_gpu + @pytest.mark.cuda_ep_test def test_compare_generation_to_io_binding( self, test_name: str, @@ -4281,25 +4316,34 @@ def test_compare_generation_to_io_binding( continue onnx_model = ORTModelForSeq2SeqLM.from_pretrained( - self._get_onnx_model_dir(model_id, model_arch, test_name), use_io_binding=False, use_cache=use_cache - ).to("cuda") + self._get_onnx_model_dir(model_id, model_arch, test_name), + use_io_binding=False, + use_cache=use_cache, + provider="CUDAExecutionProvider", + ) io_model = ORTModelForSeq2SeqLM.from_pretrained( - self._get_onnx_model_dir(model_id, model_arch, test_name), use_io_binding=True, use_cache=use_cache - ).to("cuda") + self._get_onnx_model_dir(model_id, model_arch, test_name), + use_io_binding=True, + use_cache=use_cache, + provider="CUDAExecutionProvider", + ) + + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) tokenizer = get_preprocessor(model_id) tokens = tokenizer("This is a sample output", return_tensors="pt").to("cuda") + onnx_outputs = onnx_model.generate(**tokens, num_beams=num_beams) io_outputs = io_model.generate(**tokens, num_beams=num_beams) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs, io_outputs)) + torch.testing.assert_close(onnx_outputs, io_outputs, atol=self.ATOL, rtol=self.RTOL) gc.collect() class ORTModelForSpeechSeq2SeqIntegrationTest(ORTModelTestMixin): - # TODO: speech_to_text should be tested SUPPORTED_ARCHITECTURES = ["whisper", "speech_to_text"] FULL_GRID = { @@ -4312,7 +4356,6 @@ class ORTModelForSpeechSeq2SeqIntegrationTest(ORTModelTestMixin): TASK = "automatic-speech-recognition" GENERATION_LENGTH = 100 - SPEEDUP_CACHE = 1.1 def _generate_random_audio_data(self): np.random.seed(10) @@ -4450,28 +4493,33 @@ def test_compare_to_transformers(self, test_name: str, model_arch: str, use_cach self.assertIsInstance(onnx_outputs.logits, self.TENSOR_ALIAS_TO_TYPE[input_type]) # Compare tensor outputs - self.assertTrue(torch.allclose(torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=1e-4)) + torch.testing.assert_close( + torch.Tensor(onnx_outputs.logits), transformers_outputs.logits, atol=self.ATOL, rtol=self.RTOL + ) - new_tokens = 20 # because tiny random speech to text model has a max_position_embeddings of 20 + if model_arch == "speech_to_text": + generation_length = 20 + else: + generation_length = self.GENERATION_LENGTH with torch.no_grad(): transformers_outputs = transformers_model.generate( **features["pt"], - max_new_tokens=new_tokens, - min_new_tokens=new_tokens, + max_new_tokens=generation_length, + min_new_tokens=generation_length, do_sample=False, num_beams=1, ) onnx_outputs = onnx_model.generate( **features["pt"], - max_new_tokens=new_tokens, - min_new_tokens=new_tokens, + max_new_tokens=generation_length, + min_new_tokens=generation_length, do_sample=False, num_beams=1, ) - self.assertTrue(torch.equal(onnx_outputs, transformers_outputs)) + torch.testing.assert_close(torch.Tensor(onnx_outputs), transformers_outputs, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -4562,7 +4610,6 @@ def test_pipeline_on_rocm_ep(self, test_name: str, model_arch: str, use_cache: b self.assertTrue(isinstance(outputs["text"], str)) @parameterized.expand(SUPPORTED_ARCHITECTURES) - @pytest.mark.cuda_ep_test # mark as GPU test as well to run the without/with cache timing test on the slow tests def test_compare_with_and_without_past_key_values(self, model_arch: str): model_args = {"test_name": model_arch + "_False", "model_arch": model_arch, "use_cache": False} self._setup(model_args) @@ -4579,40 +4626,34 @@ def test_compare_with_and_without_past_key_values(self, model_arch: str): self.onnx_model_dirs[model_arch + "_True"], use_cache=True ) - generation_length = self.GENERATION_LENGTH - self.GENERATION_LENGTH = 10 - _ = model_with_pkv.generate(**features) # warmup - with Timer() as with_pkv_timer: - outputs_model_with_pkv = model_with_pkv.generate( - **features, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 - ) + if model_arch == "speech_to_text": + generation_length = 20 + else: + generation_length = self.GENERATION_LENGTH + + outputs_model_with_pkv = model_with_pkv.generate( + **features, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 + ) model_without_pkv = ORTModelForSpeechSeq2Seq.from_pretrained( self.onnx_model_dirs[model_arch + "_False"], use_cache=False ) - _ = model_without_pkv.generate(**features) # warmup - with Timer() as without_pkv_timer: - outputs_model_without_pkv = model_without_pkv.generate( - **features, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 - ) - self.assertTrue(torch.equal(outputs_model_with_pkv, outputs_model_without_pkv)) - self.assertEqual( - outputs_model_with_pkv.shape[1], - self.GENERATION_LENGTH + 2 if model_arch == "whisper" else self.GENERATION_LENGTH + 1, - ) - self.assertEqual( - outputs_model_without_pkv.shape[1], - self.GENERATION_LENGTH + 2 if model_arch == "whisper" else self.GENERATION_LENGTH + 1, + outputs_model_without_pkv = model_without_pkv.generate( + **features, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 ) - self.GENERATION_LENGTH = generation_length - if os.environ.get("TEST_LEVEL", 0) == "1": - self.assertTrue( - without_pkv_timer.elapsed / with_pkv_timer.elapsed > self.SPEEDUP_CACHE, - f"With pkv latency: {with_pkv_timer.elapsed:.3f} ms, without pkv latency: {without_pkv_timer.elapsed:.3f} ms," - f" speedup: {without_pkv_timer.elapsed / with_pkv_timer.elapsed:.3f}", - ) + torch.testing.assert_close(outputs_model_with_pkv, outputs_model_without_pkv, rtol=self.RTOL, atol=self.ATOL) + + if model_arch == "whisper" and is_transformers_version(">=", "4.48"): + out_length = generation_length + elif model_arch == "whisper" and is_transformers_version(">=", "4.43"): + out_length = generation_length + 2 + else: + out_length = generation_length + 1 + + self.assertEqual(outputs_model_with_pkv.shape[1], out_length) + self.assertEqual(outputs_model_without_pkv.shape[1], out_length) @parameterized.expand(grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True]})) def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, model_arch: str, use_cache: bool): @@ -4651,18 +4692,16 @@ def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, mode self.assertEqual(model_merged.decoder_with_past, None) self.assertEqual(model_merged.use_merged, True) - generation_length = self.GENERATION_LENGTH - self.GENERATION_LENGTH = 10 + generation_length = 10 outputs_model_not_merged = model_not_merged.generate( - **features, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 + **features, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 ) outputs_model_merged = model_merged.generate( - **features, min_new_tokens=self.GENERATION_LENGTH, max_new_tokens=self.GENERATION_LENGTH, num_beams=1 + **features, min_new_tokens=generation_length, max_new_tokens=generation_length, num_beams=1 ) - self.GENERATION_LENGTH = generation_length - self.assertTrue(torch.equal(outputs_model_merged, outputs_model_not_merged)) + torch.testing.assert_close(outputs_model_not_merged, outputs_model_merged, rtol=self.RTOL, atol=self.ATOL) @parameterized.expand( grid_parameters({"model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True], "use_merged": [False, True]}) @@ -4670,9 +4709,6 @@ def test_compare_merged_and_not_merged_models_outputs(self, test_name: str, mode @require_torch_gpu @pytest.mark.cuda_ep_test def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: bool, use_merged: bool): - if use_cache is False and use_merged is True: - self.skipTest("use_cache=False, use_merged=True are uncompatible") - model_args = { "test_name": test_name, "model_arch": model_arch, @@ -4683,31 +4719,38 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForSpeechSeq2Seq.from_pretrained( - self.onnx_model_dirs[test_name], use_io_binding=False - ).to("cuda") - io_model = ORTModelForSpeechSeq2Seq.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=True).to( - "cuda" + self.onnx_model_dirs[test_name], + use_io_binding=False, + provider="CUDAExecutionProvider", + provider_options={ + "cudnn_conv_algo_search": "DEFAULT", + }, + ) + io_model = ORTModelForSpeechSeq2Seq.from_pretrained( + self.onnx_model_dirs[test_name], + use_io_binding=True, + provider="CUDAExecutionProvider", + provider_options={ + "cudnn_conv_algo_search": "DEFAULT", + }, ) self.assertFalse(onnx_model.use_io_binding) self.assertTrue(io_model.use_io_binding) processor = get_preprocessor(model_id) - data = self._generate_random_audio_data() - features = processor.feature_extractor([data] * 2, return_tensors="pt").to("cuda") - - decoder_start_token_id = onnx_model.config.decoder_start_token_id - decoder_inputs = {"decoder_input_ids": torch.ones((2, 1), dtype=torch.long) * decoder_start_token_id} + inputs = processor([data] * 2, return_tensors="pt").to("cuda") + inputs["decoder_input_ids"] = torch.ones((2, 1), dtype=torch.long).to("cuda") - onnx_outputs = onnx_model(**features, **decoder_inputs) - io_outputs = io_model(**features, **decoder_inputs) + onnx_outputs = onnx_model(**inputs) + io_outputs = io_model(**inputs) self.assertTrue("logits" in io_outputs) self.assertIsInstance(io_outputs.logits, torch.Tensor) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs.logits, io_outputs.logits)) + torch.testing.assert_close(onnx_outputs.logits, io_outputs.logits, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -4717,7 +4760,7 @@ def test_compare_to_io_binding(self, test_name: str, model_arch: str, use_cache: "model_arch": SUPPORTED_ARCHITECTURES, "use_cache": [True], "use_merged": [False, True], - "num_beams": [1, 5], + "num_beams": [1, 3], } ) ) @@ -4744,22 +4787,24 @@ def test_compare_generation_to_io_binding( model_id = MODEL_NAMES[model_arch] onnx_model = ORTModelForSpeechSeq2Seq.from_pretrained( - self.onnx_model_dirs[test_name], use_io_binding=False - ).to("cuda") - io_model = ORTModelForSpeechSeq2Seq.from_pretrained(self.onnx_model_dirs[test_name], use_io_binding=True).to( - "cuda" + self.onnx_model_dirs[test_name], use_io_binding=False, provider="CUDAExecutionProvider" + ) + io_model = ORTModelForSpeechSeq2Seq.from_pretrained( + self.onnx_model_dirs[test_name], use_io_binding=True, provider="CUDAExecutionProvider" ) - processor = get_preprocessor(model_id) + self.assertFalse(onnx_model.use_io_binding) + self.assertTrue(io_model.use_io_binding) + processor = get_preprocessor(model_id) data = self._generate_random_audio_data() - features = processor.feature_extractor(data, return_tensors="pt").to("cuda") + features = processor(data, return_tensors="pt").to("cuda") onnx_outputs = onnx_model.generate(**features, num_beams=num_beams) io_outputs = io_model.generate(**features, num_beams=num_beams) # compare tensor outputs - self.assertTrue(torch.equal(onnx_outputs, io_outputs)) + torch.testing.assert_close(onnx_outputs, io_outputs, atol=self.ATOL, rtol=self.RTOL) gc.collect() @@ -4809,7 +4854,9 @@ def test_compare_to_transformers(self, model_arch: str): self.assertIsInstance(onnx_outputs, ImageSuperResolutionOutput) self.assertTrue("reconstruction" in onnx_outputs) self.assertIsInstance(onnx_outputs.reconstruction, torch.Tensor) - self.assertTrue(torch.allclose(onnx_outputs.reconstruction, transformers_outputs.reconstruction, atol=1e-4)) + torch.testing.assert_close( + onnx_outputs.reconstruction, transformers_outputs.reconstruction, atol=self.ATOL, rtol=self.RTOL + ) gc.collect() @@ -4908,7 +4955,9 @@ class ORTModelForVision2SeqIntegrationTest(ORTModelTestMixin): TASK = "image-to-text" GENERATION_LENGTH = 100 - SPEEDUP_CACHE = 1.1 + + ATOL = 1e-3 + RTOL = 1e-3 def _get_sample_image(self): url = "http://images.cocodataset.org/val2017/000000039769.jpg" @@ -4980,55 +5029,46 @@ def test_compare_to_transformers(self, test_name: str, model_arch: str, use_cach self.assertIsInstance(onnx_model.config, PretrainedConfig) set_seed(SEED) + image_processor, tokenizer = self._get_preprocessors(model_id) transformers_model = AutoModelForVision2Seq.from_pretrained(model_id) - feature_extractor, tokenizer = self._get_preprocessors(model_id) data = self._get_sample_image() + inputs = image_processor(data, return_tensors="pt") + inputs["decoder_input_ids"] = tokenizer("This is a sample output", return_tensors="pt").input_ids - start_token = "