Skip to content

Commit

Permalink
Add openvino export configs (#568)
Browse files Browse the repository at this point in the history 
* add openvino export configs

* more libs

* more libs

* mixtral and model patcher

* chatglm export

* rework chatglm config

* more testing models

* rework config registration

* add chatglm in tests

* Update tests/openvino/test_modeling.py

* fix style

* gemma

* add test models

* qwen

* fix failed tests

* add comment for gemma
  • Loading branch information
@eaidova
eaidova authored Mar 14, 2024
1 parent 2588077 commit 8880d2e
Show file tree
Hide file tree
Showing 10 changed files with 933 additions and 58 deletions.
2 changes: 2 additions & 0 deletions optimum/exporters/openvino/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import optimum.exporters.openvino.model_configs

from .__main__ import main_export
from .convert import export, export_from_model, export_models, export_pytorch_via_onnx
from .stateful import ensure_stateful_is_available, patch_stateful
Expand Down
18 changes: 9 additions & 9 deletions optimum/exporters/openvino/__main__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ def main_export(
local_files_only: bool = False,
use_auth_token: Optional[Union[bool, str]] = None,
model_kwargs: Optional[Dict[str, Any]] = None,
custom_onnx_configs: Optional[Dict[str, "OnnxConfig"]] = None,
custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None,
fn_get_submodels: Optional[Callable] = None,
compression_option: Optional[str] = None,
compression_ratio: Optional[float] = None,
Expand Down Expand Up @@ -112,11 +112,11 @@ def main_export(
when running `transformers-cli login` (stored in `~/.huggingface`).
model_kwargs (`Optional[Dict[str, Any]]`, defaults to `None`):
Experimental usage: keyword arguments to pass to the model during
the export. This argument should be used along the `custom_onnx_configs` argument
the export. This argument should be used along the `custom_export_configs` argument
in case, for example, the model inputs/outputs are changed (for example, if
`model_kwargs={"output_attentions": True}` is passed).
custom_onnx_configs (`Optional[Dict[str, OnnxConfig]]`, defaults to `None`):
Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users that desire a finer-grained control on the export. An example is available [here](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model).
custom_export_configs (`Optional[Dict[str, OnnxConfig]]`, defaults to `None`):
Experimental usage: override the default export config used for the given model. This argument may be useful for advanced users that desire a finer-grained control on the export. An example is available [here](https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model).
fn_get_submodels (`Optional[Callable]`, defaults to `None`):
Experimental usage: Override the default submodels that are used at the export. This is
especially useful when exporting a custom architecture that needs to split the ONNX (e.g. encoder-decoder). If unspecified with custom models, optimum will try to use the default submodels used for the given task, with no guarantee of success.
Expand All @@ -134,7 +134,7 @@ def main_export(
```python
>>> from optimum.exporters.openvino import main_export
>>> main_export("gpt2", output="gpt2_onnx/")
>>> main_export("gpt2", output="gpt2_ov/")
```
"""

Expand Down Expand Up @@ -206,14 +206,14 @@ def main_export(
if model_type not in TasksManager._SUPPORTED_MODEL_TYPE:
custom_architecture = True
elif task not in TasksManager.get_supported_tasks_for_model_type(
model_type, exporter="onnx", library_name=library_name
model_type, exporter="openvino", library_name=library_name
):
if original_task == "auto":
autodetected_message = " (auto-detected)"
else:
autodetected_message = ""
model_tasks = TasksManager.get_supported_tasks_for_model_type(
model_type, exporter="onnx", library_name=library_name
model_type, exporter="openvino", library_name=library_name
)
raise ValueError(
f"Asked to export a {model_type} model for the task {task}{autodetected_message}, but the Optimum OpenVINO exporter only supports the tasks {', '.join(model_tasks.keys())} for {model_type}. Please use a supported task. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the task {task} to be supported in the ONNX export for {model_type}."
Expand Down Expand Up @@ -288,7 +288,7 @@ class StoreAttr(object):
not custom_architecture
and library_name != "diffusers"
and task + "-with-past"
in TasksManager.get_supported_tasks_for_model_type(model_type, exporter="onnx", library_name=library_name)
in TasksManager.get_supported_tasks_for_model_type(model_type, exporter="openvino", library_name=library_name)
):
# Make -with-past the default if --task was not explicitely specified
if original_task == "auto":
Expand Down Expand Up @@ -319,7 +319,7 @@ class StoreAttr(object):
ov_config=ov_config,
stateful=stateful,
model_kwargs=model_kwargs,
custom_onnx_configs=custom_onnx_configs,
custom_export_configs=custom_export_configs,
fn_get_submodels=fn_get_submodels,
preprocessors=preprocessors,
device=device,
Expand Down
49 changes: 22 additions & 27 deletions optimum/exporters/openvino/convert.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,10 +32,11 @@
from optimum.exporters.onnx.convert import check_dummy_inputs_are_allowed
from optimum.exporters.onnx.convert import export_pytorch as export_pytorch_to_onnx
from optimum.exporters.onnx.convert import export_tensorflow as export_tensorflow_onnx
from optimum.exporters.utils import _get_submodels_and_export_configs
from optimum.utils import DEFAULT_DUMMY_SHAPES, is_diffusers_available
from optimum.utils.save_utils import maybe_save_preprocessors

from ...intel.utils.import_utils import is_nncf_available, is_optimum_version
from ...intel.utils.import_utils import is_nncf_available
from .model_patcher import patch_model_with_bettertransformer
from .stateful import ensure_export_task_support_stateful, ensure_stateful_is_available, patch_stateful
from .utils import (
Expand All @@ -48,13 +49,6 @@
)


if is_optimum_version(">=", "1.16.99"):
from optimum.exporters.onnx.utils import _get_submodels_and_onnx_configs

else:
from optimum.exporters.onnx.__main__ import _get_submodels_and_onnx_configs


UNSUPPORTED_TOKENIZER_CLASSES = (T5Tokenizer, T5TokenizerFast)


Expand Down Expand Up @@ -418,7 +412,7 @@ def ts_patched_forward(*args, **kwargs):


def export_models(
models_and_onnx_configs: Dict[
models_and_export_configs: Dict[
str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]
],
output_dir: Path,
Expand All @@ -434,7 +428,7 @@ def export_models(
Export the models to OpenVINO IR format
Args:
models_and_onnx_configs (Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]):
models_and_export_configs (Dict[ str, Tuple[Union["PreTrainedModel", "TFPreTrainedModel", "ModelMixin"], "OnnxConfig"]):
output_dir (Path): output directory for saving models
opset (Optional[int], optional, Default to None): ONNX export opset
output_names (Optional[List[str]], optional, Defaults to None): model output names
Expand All @@ -459,20 +453,20 @@ def export_models(

outputs = []

if output_names is not None and len(output_names) != len(models_and_onnx_configs):
if output_names is not None and len(output_names) != len(models_and_export_configs):
raise ValueError(
f"Provided custom names {output_names} for the export of {len(models_and_onnx_configs)} models. Please provide the same number of names as models to export."
f"Provided custom names {output_names} for the export of {len(models_and_export_configs)} models. Please provide the same number of names as models to export."
)

for i, model_name in enumerate(models_and_onnx_configs.keys()):
submodel, sub_onnx_config = models_and_onnx_configs[model_name]
for i, model_name in enumerate(models_and_export_configs.keys()):
submodel, sub_export_config = models_and_export_configs[model_name]
output_name = output_names[i] if output_names is not None else Path(model_name + ".xml")
output_path = output_dir / output_name
output_path.parent.mkdir(parents=True, exist_ok=True)
outputs.append(
export(
model=submodel,
config=sub_onnx_config,
config=sub_export_config,
output=output_path,
opset=opset,
device=device,
Expand All @@ -495,7 +489,7 @@ def export_from_model(
stateful: bool = True,
opset: Optional[int] = None,
model_kwargs: Optional[Dict[str, Any]] = None,
custom_onnx_configs: Optional[Dict[str, "OnnxConfig"]] = None,
custom_export_configs: Optional[Dict[str, "OnnxConfig"]] = None,
fn_get_submodels: Optional[Callable] = None,
preprocessors: List = None,
device: str = "cpu",
Expand Down Expand Up @@ -524,14 +518,14 @@ def export_from_model(
task = TasksManager._infer_task_from_model_or_model_class(model=model)
except (ValueError, KeyError) as e:
raise RuntimeError(
f"The model task could not be automatically inferred in `onnx_export_from_model`. Please provide the argument `task` with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
f"The model task could not be automatically inferred in `export_from_model`. Please provide the argument `task` with the relevant task from {', '.join(TasksManager.get_all_tasks())}. Detailed error: {e}"
)

if (
not custom_architecture
and library_name != "diffusers"
and task + "-with-past"
in TasksManager.get_supported_tasks_for_model_type(model_type, "onnx", library_name=library_name)
in TasksManager.get_supported_tasks_for_model_type(model_type, "openvino", library_name=library_name)
):
# -with-past is the default.
task = task + "-with-past"
Expand All @@ -541,9 +535,9 @@ def export_from_model(
stateful = stateful and ensure_export_task_support_stateful(task)

# TODO: support onnx_config.py in the model repo
if custom_architecture and custom_onnx_configs is None:
if custom_architecture and custom_export_configs is None:
raise ValueError(
f"Trying to export a {model_type} model, that is a custom or unsupported architecture, but no custom onnx configuration was passed as `custom_onnx_configs`. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type {model_type} to be supported natively in the ONNX export."
f"Trying to export a {model_type} model, that is a custom or unsupported architecture, but no custom export configuration was passed as `custom_export_configs`. Please refer to https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#custom-export-of-transformers-models for an example on how to export custom models. Please open an issue at https://github.com/huggingface/optimum/issues if you would like the model type {model_type} to be supported natively in the ONNX export."
)

if task.startswith("text-generation") and model.config.is_encoder_decoder:
Expand All @@ -569,18 +563,19 @@ def export_from_model(
kwargs_shapes[input_name] if input_name in kwargs_shapes else DEFAULT_DUMMY_SHAPES[input_name]
)

onnx_config, models_and_onnx_configs = _get_submodels_and_onnx_configs(
export_config, models_and_export_configs = _get_submodels_and_export_configs(
model=model,
task=task,
monolith=False,
custom_onnx_configs=custom_onnx_configs if custom_onnx_configs is not None else {},
custom_export_configs=custom_export_configs if custom_export_configs is not None else {},
custom_architecture=custom_architecture,
fn_get_submodels=fn_get_submodels,
preprocessors=preprocessors,
library_name=library_name,
model_kwargs=model_kwargs,
_variant="default",
legacy=False,
exporter="openvino",
)

if ov_config is None:
Expand Down Expand Up @@ -612,18 +607,18 @@ def export_from_model(
model_name_or_path = model.config._name_or_path
maybe_save_preprocessors(model_name_or_path, output)

files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_onnx_configs.keys()]
files_subpaths = ["openvino_" + model_name + ".xml" for model_name in models_and_export_configs.keys()]

else:
# save the subcomponent configuration
for model_name in models_and_onnx_configs:
subcomponent = models_and_onnx_configs[model_name][0]
for model_name in models_and_export_configs:
subcomponent = models_and_export_configs[model_name][0]
if hasattr(subcomponent, "save_config"):
subcomponent.save_config(output / model_name)
elif hasattr(subcomponent, "config") and hasattr(subcomponent.config, "save_pretrained"):
subcomponent.config.save_pretrained(output / model_name)

files_subpaths = [os.path.join(name_dir, OV_XML_FILE_NAME) for name_dir in models_and_onnx_configs]
files_subpaths = [os.path.join(name_dir, OV_XML_FILE_NAME) for name_dir in models_and_export_configs]

# Saving the additional components needed to perform inference.
model.scheduler.save_pretrained(output.joinpath("scheduler"))
Expand All @@ -643,7 +638,7 @@ def export_from_model(
model.save_config(output)

export_models(
models_and_onnx_configs=models_and_onnx_configs,
models_and_export_configs=models_and_export_configs,
output_dir=output,
output_names=files_subpaths,
input_shapes=input_shapes,
Expand Down
Loading

0 comments on commit 8880d2e

Please sign in to comment.