diff --git a/optimum/exporters/openvino/__main__.py b/optimum/exporters/openvino/__main__.py index 128643eb1d..65648b2753 100644 --- a/optimum/exporters/openvino/__main__.py +++ b/optimum/exporters/openvino/__main__.py @@ -46,6 +46,7 @@ MULTI_MODAL_TEXT_GENERATION_MODELS, clear_class_registry, deduce_diffusers_dtype, + patch_not_check_trace, ) @@ -234,6 +235,7 @@ def main_export( do_gptq_patching = False custom_architecture = False patch_16bit = False + patch_trace = False loading_kwargs = model_loading_kwargs or {} if library_name == "transformers": config = AutoConfig.from_pretrained( @@ -350,6 +352,7 @@ class StoreAttr(object): if dtype in [torch.float16, torch.bfloat16]: loading_kwargs["torch_dtype"] = dtype patch_16bit = True + patch_trace = True if library_name == "open_clip": model = _OpenClipForZeroShotImageClassification.from_pretrained(model_name_or_path, cache_dir=cache_dir) @@ -417,21 +420,22 @@ class StoreAttr(object): model_name_or_path, subfolder=subfolder, trust_remote_code=trust_remote_code ) - submodel_paths = export_from_model( - model=model, - output=output, - task=task, - ov_config=ov_config, - stateful=stateful, - model_kwargs=model_kwargs, - custom_export_configs=custom_export_configs, - fn_get_submodels=fn_get_submodels, - preprocessors=preprocessors, - device=device, - trust_remote_code=trust_remote_code, - patch_16bit_model=patch_16bit, - **kwargs_shapes, - ) + with patch_not_check_trace(patch_trace): + submodel_paths = export_from_model( + model=model, + output=output, + task=task, + ov_config=ov_config, + stateful=stateful, + model_kwargs=model_kwargs, + custom_export_configs=custom_export_configs, + fn_get_submodels=fn_get_submodels, + preprocessors=preprocessors, + device=device, + trust_remote_code=trust_remote_code, + patch_16bit_model=patch_16bit, + **kwargs_shapes, + ) if convert_tokenizer: maybe_convert_tokenizers(library_name, output, model, preprocessors, task=task) diff --git a/optimum/exporters/openvino/utils.py b/optimum/exporters/openvino/utils.py index db4df6d0d6..5de2aec5da 100644 --- a/optimum/exporters/openvino/utils.py +++ b/optimum/exporters/openvino/utils.py @@ -15,6 +15,8 @@ import inspect import logging from collections import namedtuple +from contextlib import contextmanager +from functools import partial from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -288,3 +290,15 @@ def save_preprocessors( logger.error(f"Saving {type(processor)} failed with {ex}") else: maybe_save_preprocessors(model_name_or_path, output, trust_remote_code=trust_remote_code) + + +@contextmanager +def patch_not_check_trace(to_patch): + original_trace = torch.jit.trace + if to_patch: + torch.jit.trace = partial(torch.jit.trace, check_trace=False) + try: + yield + finally: + if to_patch: + torch.jit.trace = original_trace