diff --git a/optimum/onnxruntime/io_binding/__init__.py b/optimum/onnxruntime/io_binding/__init__.py new file mode 100644 index 0000000000..e0810d5e80 --- /dev/null +++ b/optimum/onnxruntime/io_binding/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from .io_binding_helper import TypeHelper diff --git a/optimum/onnxruntime/io_binding/io_binding_helper.py b/optimum/onnxruntime/io_binding/io_binding_helper.py new file mode 100644 index 0000000000..e8005188be --- /dev/null +++ b/optimum/onnxruntime/io_binding/io_binding_helper.py @@ -0,0 +1,60 @@ +# Copyright 2022 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import numpy as np +import torch + +from onnxruntime.transformers.io_binding_helper import TypeHelper as ORTTypeHelper + + +# Adapted from https://github.com/microsoft/onnxruntime/blob/93e0a151177ad8222c2c95f814342bfa27f0a64d/onnxruntime/python/tools/transformers/io_binding_helper.py#L12 +class TypeHelper(ORTTypeHelper): + """ + Gets data type information of the ONNX Runtime inference session and provides the mapping from + `OrtValue` data types to the data types of other frameworks (NumPy, PyTorch, etc). + """ + + # TODO: Current DLPack doesn't support boolean tensor, use uint8 as workaround, remove after it is supported. + @staticmethod + def ort_type_to_numpy_type(ort_type: str): + ort_type_to_numpy_type_map = { + "tensor(int64)": np.int64, + "tensor(int32)": np.int32, + "tensor(int8)": np.int8, + "tensor(float)": np.float32, + "tensor(float16)": np.float16, + "tensor(bool)": np.uint8, + } + if ort_type in ort_type_to_numpy_type_map: + return ort_type_to_numpy_type_map[ort_type] + else: + raise ValueError( + f"{ort_type} is not supported. Here is a list of supported data type: {ort_type_to_numpy_type_map.keys()}" + ) + + @staticmethod + def ort_type_to_torch_type(ort_type: str): + ort_type_to_torch_type_map = { + "tensor(int64)": torch.int64, + "tensor(int32)": torch.int32, + "tensor(int8)": torch.int8, + "tensor(float)": torch.float32, + "tensor(float16)": torch.float16, + "tensor(bool)": torch.bool, + } + if ort_type in ort_type_to_torch_type_map: + return ort_type_to_torch_type_map[ort_type] + else: + raise ValueError( + f"{ort_type} is not supported. Here is a list of supported data type: {ort_type_to_torch_type_map.keys()}" + ) diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py index ed96f53f65..367c58acae 100644 --- a/optimum/onnxruntime/modeling_ort.py +++ b/optimum/onnxruntime/modeling_ort.py @@ -18,6 +18,7 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Union +import numpy as np import torch from transformers import ( AutoConfig, @@ -49,6 +50,7 @@ from huggingface_hub import HfApi, hf_hub_download from ..modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel +from .io_binding import TypeHelper from .utils import ONNX_WEIGHTS_NAME, get_device_for_provider, get_provider_for_device, parse_device @@ -66,6 +68,7 @@ Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~onnxruntime.modeling_ort.ORTModel.from_pretrained`] method to load the model weights. model (`onnxruntime.InferenceSession`): [onnxruntime.InferenceSession](https://onnxruntime.ai/docs/api/python/api_summary.html#inferencesession) is the main class used to run a model. Check out the [`~onnxruntime.modeling_ort.ORTModel.load_model`] method for more information. + use_io_binding (`bool`, *optional*): Whether to use IOBinding during inference to avoid memory copy between the host and devices. Defaults to `True` if the device is CUDA, otherwise defaults to `False`. """ ONNX_TEXT_INPUTS_DOCSTRING = r""" @@ -107,9 +110,16 @@ class ORTModel(OptimizedModel): base_model_prefix = "onnx_model" auto_model_class = AutoModel - def __init__(self, model: ort.InferenceSession = None, config: PretrainedConfig = None, **kwargs): + def __init__( + self, + model: ort.InferenceSession = None, + config: PretrainedConfig = None, + use_io_binding: bool = True, + **kwargs + ): self.model = model self.config = config + self.use_io_binding = use_io_binding self.model_save_dir = kwargs.get("model_save_dir", None) self.latest_model_name = kwargs.get("latest_model_name", "model.onnx") self.providers = model.get_providers() @@ -121,6 +131,12 @@ def __init__(self, model: ort.InferenceSession = None, config: PretrainedConfig f" Use `ort_model.to()` to send the outputs to the wanted device." ) + if "TensorrtExecutionProvider" in self.providers and self.use_io_binding: + logger.warning( + "There is no need to do IO binding for TensorrtExecutionProvider, `use_io_binding` will be set to False." + ) + self.use_io_binding = False + # registers the ORTModelForXXX classes into the transformers AutoModel classes # to avoid warnings when create a pipeline https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/pipelines/base.py#L863 AutoConfig.register(self.base_model_prefix, AutoConfig) @@ -438,10 +454,78 @@ class ORTModelForFeatureExtraction(ORTModel): export_feature = "default" auto_model_class = AutoModel - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) # create {name:idx} dict for model outputs self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None + + def prepare_output_buffer(self, batch_size, sequence_length, hidden_size): + """Prepare the buffer of output(`last_hidden_state`) with a 1D tensor on shape: (batch_size, sequence_length, hidden_size).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + output_shape = (batch_size, sequence_length, hidden_size) + output_buffer = torch.empty(np.prod(output_shape), dtype=torch_type, device=self.device) + + return output_shape, output_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + ): + io_binding = self.model.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self.device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self.device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + if token_type_ids is not None: + # bind token type ids + io_binding.bind_input( + "token_type_ids", + token_type_ids.device.type, + self.device.index, + self.name_to_np_type["token_type_ids"], + tuple(token_type_ids.shape), + token_type_ids.data_ptr(), + ) + + # bind logits + output_shape, output_buffer = self.prepare_output_buffer( + batch_size=input_ids.size(0), + sequence_length=input_ids.size(1), + hidden_size=self.config.hidden_size, + ) + io_binding.bind_output( + "last_hidden_state", + output_buffer.device.type, + self.device.index, + self.name_to_np_type["last_hidden_state"], + output_shape, + output_buffer.data_ptr(), + ) + output_shapes = {"last_hidden_state": output_shape} + output_buffers = {"last_hidden_state": output_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") @@ -458,18 +542,35 @@ def forward( token_type_ids: Optional[torch.Tensor] = None, **kwargs, ): - # converts pytorch inputs into numpy inputs for onnx - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "attention_mask": attention_mask.cpu().detach().numpy(), - } - if token_type_ids is not None: - onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() - # run inference - outputs = self.model.run(None, onnx_inputs) - last_hidden_state = torch.from_numpy(outputs[self.model_outputs["last_hidden_state"]]).to(self.device) - # converts output to namedtuple for pipelines post-processing - return BaseModelOutput(last_hidden_state=last_hidden_state) + 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 + ) + + # 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() + + # converts output to namedtuple for pipelines post-processing + return BaseModelOutput( + last_hidden_state=output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) + ) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "attention_mask": attention_mask.cpu().detach().numpy(), + } + if token_type_ids is not None: + onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() + + # run inference + outputs = self.model.run(None, onnx_inputs) + last_hidden_state = torch.from_numpy(outputs[self.model_outputs["last_hidden_state"]]).to(self.device) + + # converts output to namedtuple for pipelines post-processing + return BaseModelOutput(last_hidden_state=last_hidden_state) QUESTION_ANSWERING_EXAMPLE = r""" @@ -523,10 +624,87 @@ class ORTModelForQuestionAnswering(ORTModel): export_feature = "question-answering" auto_model_class = AutoModelForQuestionAnswering - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) # create {name:idx} dict for model outputs self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None + + def prepare_logits_buffer(self, batch_size, sequence_length): + """Prepare the buffer of logits with a 1D tensor on shape: (batch_size, sequence_length).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + logits_shape = (batch_size, sequence_length) + logits_buffer = torch.empty(np.prod(logits_shape), dtype=torch_type, device=self.device) + + return logits_shape, logits_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + ): + io_binding = self.model.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self.device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self.device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + if token_type_ids is not None: + # bind token type ids + io_binding.bind_input( + "token_type_ids", + token_type_ids.device.type, + self.device.index, + self.name_to_np_type["token_type_ids"], + tuple(token_type_ids.shape), + token_type_ids.data_ptr(), + ) + + # bind logits + start_logits_shape, start_logits_buffer = self.prepare_logits_buffer( + batch_size=input_ids.size(0), sequence_length=input_ids.size(1) + ) + end_logits_shape, end_logits_buffer = self.prepare_logits_buffer( + batch_size=input_ids.size(0), sequence_length=input_ids.size(1) + ) + io_binding.bind_output( + "start_logits", + start_logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + start_logits_shape, + start_logits_buffer.data_ptr(), + ) + io_binding.bind_output( + "end_logits", + end_logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + end_logits_shape, + end_logits_buffer.data_ptr(), + ) + output_shapes = {"start_logits": start_logits_shape, "end_logits": end_logits_shape} + output_buffers = {"start_logits": start_logits_buffer, "end_logits": end_logits_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") @@ -543,19 +721,41 @@ def forward( token_type_ids: Optional[torch.Tensor] = None, **kwargs, ): - # converts pytorch inputs into numpy inputs for onnx - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "attention_mask": attention_mask.cpu().detach().numpy(), - } - if token_type_ids is not None: - onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() - # run inference - outputs = self.model.run(None, onnx_inputs) - start_logits = torch.from_numpy(outputs[self.model_outputs["start_logits"]]).to(self.device) - end_logits = torch.from_numpy(outputs[self.model_outputs["end_logits"]]).to(self.device) - # converts output to namedtuple for pipelines post-processing - return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits) + 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 + ) + + # 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() + + # map outputs with names + start_logits = io_binding._iobinding.get_outputs()[0] + end_logits = io_binding._iobinding.get_outputs()[1] + + # converts output to namedtuple for pipelines post-processing + return QuestionAnsweringModelOutput( + start_logits=output_buffers["start_logits"].view(output_shapes["start_logits"]), + end_logits=output_buffers["end_logits"].view(output_shapes["end_logits"]), + ) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "attention_mask": attention_mask.cpu().detach().numpy(), + } + if token_type_ids is not None: + onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() + + # run inference + outputs = self.model.run(None, onnx_inputs) + start_logits = torch.from_numpy(outputs[self.model_outputs["start_logits"]]).to(self.device) + end_logits = torch.from_numpy(outputs[self.model_outputs["end_logits"]]).to(self.device) + + # converts output to namedtuple for pipelines post-processing + return QuestionAnsweringModelOutput(start_logits=start_logits, end_logits=end_logits) SEQUENCE_CLASSIFICATION_EXAMPLE = r""" @@ -623,11 +823,78 @@ class ORTModelForSequenceClassification(ORTModel): export_feature = "sequence-classification" auto_model_class = AutoModelForSequenceClassification - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) # create {name:idx} dict for model outputs self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} self.model_inputs = {input_key.name: idx for idx, input_key in enumerate(self.model.get_inputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None + + def prepare_logits_buffer(self, batch_size, num_labels): + """Prepare the buffer of logits with a 1D tensor on shape: (batch_size, config.num_labels).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + logits_shape = (batch_size, num_labels) + logits_buffer = torch.empty(np.prod(logits_shape), dtype=torch_type, device=self.device) + + return logits_shape, logits_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + ): + io_binding = self.model.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self.device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self.device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + if token_type_ids is not None: + # bind token type ids + io_binding.bind_input( + "token_type_ids", + token_type_ids.device.type, + self.device.index, + self.name_to_np_type["token_type_ids"], + tuple(token_type_ids.shape), + token_type_ids.data_ptr(), + ) + + # bind logits + logits_shape, logits_buffer = self.prepare_logits_buffer( + batch_size=input_ids.size(0), + num_labels=self.config.num_labels, + ) + io_binding.bind_output( + "logits", + logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + logits_shape, + logits_buffer.data_ptr(), + ) + output_shapes = {"logits": logits_shape} + output_buffers = {"logits": logits_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") @@ -644,19 +911,36 @@ def forward( token_type_ids: Optional[torch.Tensor] = None, **kwargs, ): - # converts pytorch inputs into numpy inputs for onnx - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "attention_mask": attention_mask.cpu().detach().numpy(), - } + 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 + ) - if token_type_ids is not None: - onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() - # run inference - outputs = self.model.run(None, onnx_inputs) - logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) - # converts output to namedtuple for pipelines post-processing - return SequenceClassifierOutput(logits=logits) + # 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() + + # map outputs with names + logits = io_binding._iobinding.get_outputs()[0] + + # converts output to namedtuple for pipelines post-processing + return SequenceClassifierOutput(logits=output_buffers["logits"].view(output_shapes["logits"])) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "attention_mask": attention_mask.cpu().detach().numpy(), + } + if token_type_ids is not None: + onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() + + # run inference + outputs = self.model.run(None, onnx_inputs) + logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) + + # converts output to namedtuple for pipelines post-processing + return SequenceClassifierOutput(logits=logits) TOKEN_CLASSIFICATION_EXAMPLE = r""" @@ -709,10 +993,78 @@ class ORTModelForTokenClassification(ORTModel): export_feature = "token-classification" auto_model_class = AutoModelForTokenClassification - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) # create {name:idx} dict for model outputs self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None + + def prepare_logits_buffer(self, batch_size, sequence_length, num_labels): + """Prepare the buffer of logits with a 1D tensor on shape: (batch_size, sequence_length, config.num_labels).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + logits_shape = (batch_size, sequence_length, num_labels) + logits_buffer = torch.empty(np.prod(logits_shape), dtype=torch_type, device=self.device) + + return logits_shape, logits_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + ): + io_binding = self.model.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self.device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self.device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + if token_type_ids is not None: + # bind token type ids + io_binding.bind_input( + "token_type_ids", + token_type_ids.device.type, + self.device.index, + self.name_to_np_type["token_type_ids"], + tuple(token_type_ids.shape), + token_type_ids.data_ptr(), + ) + + # bind logits + logits_shape, logits_buffer = self.prepare_logits_buffer( + batch_size=input_ids.size(0), + sequence_length=input_ids.size(1), + num_labels=self.config.num_labels, + ) + io_binding.bind_output( + "logits", + logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + logits_shape, + logits_buffer.data_ptr(), + ) + output_shapes = {"logits": logits_shape} + output_buffers = {"logits": logits_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") @@ -729,18 +1081,36 @@ def forward( token_type_ids: Optional[torch.Tensor] = None, **kwargs, ): - # converts pytorch inputs into numpy inputs for onnx - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "attention_mask": attention_mask.cpu().detach().numpy(), - } - if token_type_ids is not None: - onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() - # run inference - outputs = self.model.run(None, onnx_inputs) - logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) - # converts output to namedtuple for pipelines post-processing - return TokenClassifierOutput(logits=logits) + 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 + ) + + # 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() + + # map outputs with names + logits = io_binding._iobinding.get_outputs()[0] + + # converts output to namedtuple for pipelines post-processing + return TokenClassifierOutput(logits=output_buffers["logits"].view(output_shapes["logits"])) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "attention_mask": attention_mask.cpu().detach().numpy(), + } + if token_type_ids is not None: + onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() + + # run inference + outputs = self.model.run(None, onnx_inputs) + logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) + + # converts output to namedtuple for pipelines post-processing + return TokenClassifierOutput(logits=logits) MULTIPLE_CHOICE_EXAMPLE = r""" @@ -788,9 +1158,75 @@ class ORTModelForMultipleChoice(ORTModel): export_feature = "multiple-choice" auto_model_class = AutoModelForMultipleChoice - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None + + def prepare_logits_buffer(self, batch_size, num_choices): + """Prepare the buffer of logits with a 1D tensor on shape: (batch_size, num_choices).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + logits_shape = (batch_size, num_choices) + logits_buffer = torch.empty(np.prod(logits_shape), dtype=torch_type, device=self.device) + + return logits_shape, logits_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + ): + io_binding = self.model.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self.device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self.device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + if token_type_ids is not None: + # bind token type ids + io_binding.bind_input( + "token_type_ids", + token_type_ids.device.type, + self.device.index, + self.name_to_np_type["token_type_ids"], + tuple(token_type_ids.shape), + token_type_ids.data_ptr(), + ) + + # bind logits + logits_shape, logits_buffer = self.prepare_logits_buffer( + batch_size=input_ids.size(0), num_choices=input_ids.size(1) + ) + io_binding.bind_output( + "logits", + logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + logits_shape, + logits_buffer.data_ptr(), + ) + output_shapes = {"logits": logits_shape} + output_buffers = {"logits": logits_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") @@ -807,20 +1243,33 @@ def forward( token_type_ids: Optional[torch.Tensor] = None, **kwargs, ): - # Converts pytorch inputs into numpy inputs - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "attention_mask": attention_mask.cpu().detach().numpy(), - } + 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 + ) - if token_type_ids is not None: - onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() + # 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() + + # converts output to namedtuple for pipelines post-processing + return MultipleChoiceModelOutput(logits=output_buffers["logits"].view(output_shapes["logits"])) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "attention_mask": attention_mask.cpu().detach().numpy(), + } + if token_type_ids is not None: + onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy() - # Run inference - outputs = self.model.run(None, onnx_inputs) - logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) + # run inference + outputs = self.model.run(None, onnx_inputs) + logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) - return MultipleChoiceModelOutput(logits=logits) + # converts output to namedtuple for pipelines post-processing + return MultipleChoiceModelOutput(logits=logits) TEXT_GENERATION_EXAMPLE = r""" @@ -872,11 +1321,12 @@ class ORTModelForCausalLM(ORTModel, GenerationMixin): export_feature = "causal-lm" auto_model_class = AutoModelForCausalLM - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) # create {name:idx} dict for model outputs self.main_input_name = "input_ids" self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None def prepare_inputs_for_generation(self, input_ids: torch.LongTensor, **kwargs) -> Dict[str, Any]: """ @@ -887,6 +1337,59 @@ def prepare_inputs_for_generation(self, input_ids: torch.LongTensor, **kwargs) - inputs["attention_mask"] = kwargs["attention_mask"] return inputs + def prepare_logits_buffer(self, batch_size, sequence_length): + """Prepare the buffer of logits with a 1D tensor on shape: (batch_size, sequence_length, config.vocab_size).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + logits_shape = (batch_size, sequence_length, self.config.vocab_size) + logits_buffer = torch.empty(np.prod(logits_shape), dtype=torch_type, device=self.device) + + return logits_shape, logits_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ): + io_binding = self.model.io_binding() + + # bind input_ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self.device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self.device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + # bind logits + logits_shape, logits_buffer = self.prepare_logits_buffer( + batch_size=input_ids.size(0), sequence_length=input_ids.size(1) + ) + io_binding.bind_output( + "logits", + logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + logits_shape, + logits_buffer.data_ptr(), + ) + output_shapes = {"logits": logits_shape} + output_buffers = {"logits": logits_buffer} + + return io_binding, output_shapes, output_buffers + @add_start_docstrings_to_model_forward( ONNX_TEXT_INPUTS_DOCSTRING.format("batch_size, sequence_length") + TEXT_GENERATION_EXAMPLE.format( @@ -901,16 +1404,29 @@ def forward( attention_mask: Optional[torch.Tensor] = None, **kwargs, ): - # converts pytorch inputs into numpy inputs for onnx - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "attention_mask": attention_mask.cpu().detach().numpy(), - } - # run inference - outputs = self.model.run(None, onnx_inputs) - logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) - # converts output to namedtuple for pipelines post-processing - return CausalLMOutputWithCrossAttentions(logits=logits) + if self.device.type == "cuda" and self.use_io_binding: + io_binding, output_shapes, output_buffers = self.prepare_io_binding(input_ids, attention_mask) + + # 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() + + # converts output to namedtuple for pipelines post-processing + return CausalLMOutputWithCrossAttentions(logits=output_buffers["logits"].view(output_shapes["logits"])) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "attention_mask": attention_mask.cpu().detach().numpy(), + } + + # run inference + outputs = self.model.run(None, onnx_inputs) + logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device) + + # converts output to namedtuple for pipelines post-processing + return CausalLMOutputWithCrossAttentions(logits=logits) # Adapted from https://github.com/huggingface/transformers/blob/99289c08a1b16a805dd4ee46de029e9fd23cba3d/src/transformers/generation_utils.py#L490 def _prepare_attention_mask_for_generation( @@ -990,10 +1506,52 @@ class ORTModelForImageClassification(ORTModel): export_feature = "image-classification" auto_model_class = AutoModelForImageClassification - def __init__(self, model=None, config=None, **kwargs): - super().__init__(model, config, **kwargs) + def __init__(self, model=None, config=None, use_io_binding=True, **kwargs): + super().__init__(model, config, use_io_binding, **kwargs) # create {name:idx} dict for model outputs self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())} + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model) if self.use_io_binding else None + + def prepare_logits_buffer(self, batch_size): + """Prepare the buffer of logits with a 1D tensor on shape: (batch_size, config.num_labels).""" + ort_type = TypeHelper.get_output_type(self.model, "logits") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + logits_shape = (batch_size, self.config.num_labels) + logits_buffer = torch.empty(np.prod(logits_shape), dtype=torch_type, device=self.device) + + return logits_shape, logits_buffer + + def prepare_io_binding( + self, + pixel_values: torch.Tensor, + ): + io_binding = self.model.io_binding() + + # bind pixel values + io_binding.bind_input( + "pixel_values", + pixel_values.device.type, + self.device.index, + self.name_to_np_type["pixel_values"], + tuple(pixel_values.shape), + pixel_values.data_ptr(), + ) + + # bind logits + logits_shape, logits_buffer = self.prepare_logits_buffer(batch_size=pixel_values.size(0)) + io_binding.bind_output( + "logits", + logits_buffer.device.type, + self.device.index, + self.name_to_np_type["logits"], + logits_shape, + logits_buffer.data_ptr(), + ) + output_shapes = {"logits": logits_shape} + output_buffers = {"logits": logits_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward( ONNX_IMAGE_INPUTS_DOCSTRING.format("batch_size, num_channels, height, width") @@ -1008,16 +1566,28 @@ def forward( pixel_values: torch.Tensor, **kwargs, ): - # converts pytorch inputs into numpy inputs for onnx - onnx_inputs = { - "pixel_values": pixel_values.cpu().detach().numpy(), - } - # run inference - outputs = self.model.run(None, onnx_inputs) - # converts output to namedtuple for pipelines post-processing - return ImageClassifierOutput( - logits=torch.from_numpy(outputs[self.model_outputs["logits"]]), - ) + if self.device.type == "cuda" and self.use_io_binding: + io_binding, output_shapes, output_buffers = self.prepare_io_binding(pixel_values) + + # 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() + + # converts output to namedtuple for pipelines post-processing + return ImageClassifierOutput(logits=output_buffers["logits"].view(output_shapes["logits"])) + else: + # converts pytorch inputs into numpy inputs for onnx + onnx_inputs = { + "pixel_values": pixel_values.cpu().detach().numpy(), + } + + # run inference + outputs = self.model.run(None, onnx_inputs) + logits = torch.from_numpy(outputs[self.model_outputs["logits"]]) + + # converts output to namedtuple for pipelines post-processing + return ImageClassifierOutput(logits=logits) CUSTOM_TASKS_EXAMPLE = r""" @@ -1069,6 +1639,12 @@ class ORTModelForCustomTasks(ORTModel): def __init__(self, model=None, config=None, **kwargs): super().__init__(model, config, **kwargs) + if kwargs.pop("use_io_binding", False): + logger.warning( + "ORTModelForCustomTasks doesn't support IO Binding yet, and the inference will be done without IO binding which could cause" + " significant overhead on data copying. If you want us to enable IO binding for custom use case, please open an issue in " + "Optimum: https://github.com/huggingface/optimum." + ) @add_start_docstrings_to_model_forward( CUSTOM_TASKS_EXAMPLE.format( diff --git a/optimum/onnxruntime/modeling_seq2seq.py b/optimum/onnxruntime/modeling_seq2seq.py index e79c7d847e..7353eac288 100644 --- a/optimum/onnxruntime/modeling_seq2seq.py +++ b/optimum/onnxruntime/modeling_seq2seq.py @@ -18,6 +18,7 @@ from pathlib import Path from typing import Any, DefaultDict, Dict, List, Mapping, Optional, Set, Tuple, Union +import numpy as np import torch import transformers from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer @@ -32,11 +33,13 @@ from ..onnx.configuration import DecoderOnnxConfig, EncoderOnnxConfig from ..onnx.modeling_seq2seq import _DecoderWithLMhead +from .io_binding import TypeHelper from .modeling_ort import ORTModel from .utils import ( ONNX_DECODER_NAME, ONNX_DECODER_WITH_PAST_NAME, ONNX_ENCODER_NAME, + ORTConfigManager, _is_gpu_available, get_device_for_provider, get_provider_for_device, @@ -67,6 +70,9 @@ decoder_with_past_file_name(`str`, *optional*): The decoder with past key values model file name overwriting the default file name, allowing to save the decoder model with a different name. + use_io_binding (`bool`, *optional*): + Whether use IOBinding during inference to avoid memory copy between the host and devices. Defaults to `True` + if the device is CUDA, otherwise defaults to `False`. """ ENCODER_INPUTS_DOCSTRING = r""" @@ -153,7 +159,7 @@ ONNX_INPUTS_DOCSTRING, ) class ORTModelForConditionalGeneration(ORTModel): - # Used in from_transformers to export model to onnx + # Used in from_transformers to export model to onnxORTEncoder base_model_prefix = "onnx_model" export_feature = "seq2seq-lm" auto_model_class = AutoModelForSeq2SeqLM @@ -164,21 +170,40 @@ def __init__( decoder_session: onnxruntime.InferenceSession = None, decoder_with_past_session: onnxruntime.InferenceSession = None, config: transformers.PretrainedConfig = None, + use_io_binding: bool = True, **kwargs ): self.config = config + self.use_io_binding = use_io_binding self.model_save_dir = kwargs.get("model_save_dir", None) self.providers = encoder_session.get_providers() self._device = get_device_for_provider(encoder_session.get_providers()[0]) - self.encoder = ORTEncoder(session=encoder_session, device=self._device) - self.decoder = ORTDecoder(session=decoder_session, device=self._device) + if "TensorrtExecutionProvider" in self.providers and self.use_io_binding: + logger.warning( + "There is no need to do IO binding for TensorrtExecutionProvider, `use_io_binding` will be set to False." + ) + self.use_io_binding = False + + self.encoder = ORTEncoder( + session=encoder_session, config=self.config, device=self._device, use_io_binding=self.use_io_binding + ) + self.decoder = ORTDecoder( + session=decoder_session, config=self.config, device=self._device, use_io_binding=self.use_io_binding + ) self.use_cache = decoder_with_past_session is not None # If a decoder_with_past_path is provided, an inference session for the decoder with past key/values as inputs # will be enabled self.decoder_with_past = ( - ORTDecoder(session=decoder_with_past_session, device=self._device) if self.use_cache else None + ORTDecoder( + session=decoder_with_past_session, + config=self.config, + device=self._device, + use_io_binding=self.use_io_binding, + ) + if self.use_cache + else None ) self.encoder_file_name = kwargs.get("last_encoder_model_name", ONNX_ENCODER_NAME) self.decoder_file_name = kwargs.get("last_decoder_model_name", ONNX_DECODER_NAME) @@ -513,12 +538,78 @@ class ORTEncoder: The ONNX Runtime inference session associated to the encoder. """ - def __init__(self, session: onnxruntime.InferenceSession, device: torch.device): + def __init__( + self, + session: onnxruntime.InferenceSession, + config: transformers.PretrainedConfig, + device: torch.device, + use_io_binding: bool = True, + **kwargs + ): self.session = session + self.config = config self._device = device + self.use_io_binding = use_io_binding self.main_input_name = "input_ids" 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.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.session) if self.use_io_binding else None + + def prepare_output_buffer(self, batch_size, sequence_length): + """Prepare the buffer of output(`last_hidden_state`) with a 1D tensor on shape: (batch_size, sequence_length, hidden_size).""" + ort_type = TypeHelper.get_output_type(self.session, "last_hidden_state") + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + + hidden_size = getattr(self.config, ORTConfigManager.get_hidden_size_name(self.config.model_type)) + output_shape = (batch_size, sequence_length, hidden_size) + output_buffer = torch.empty(np.prod(output_shape), dtype=torch_type, device=self._device) + + return output_shape, output_buffer + + def prepare_io_binding( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + ): + io_binding = self.session.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self._device.index, + self.name_to_np_type["input_ids"], + tuple(input_ids.shape), + input_ids.data_ptr(), + ) + if "attention_mask" in self.input_names: + # bind attention mask + io_binding.bind_input( + "attention_mask", + attention_mask.device.type, + self._device.index, + self.name_to_np_type["attention_mask"], + tuple(attention_mask.shape), + attention_mask.data_ptr(), + ) + + # bind logits + output_shape, output_buffer = self.prepare_output_buffer( + batch_size=input_ids.size(0), + sequence_length=input_ids.size(1), + ) + io_binding.bind_output( + "last_hidden_state", + output_buffer.device.type, + self._device.index, + self.name_to_np_type["last_hidden_state"], + output_shape, + output_buffer.data_ptr(), + ) + output_shapes = {"last_hidden_state": output_shape} + output_buffers = {"last_hidden_state": output_buffer} + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward(ENCODER_INPUTS_DOCSTRING) def forward( @@ -528,17 +619,30 @@ def forward( **kwargs, ) -> BaseModelOutput: - onnx_inputs = {"input_ids": input_ids.cpu().detach().numpy()} + if self._device.type == "cuda" and self.use_io_binding: + io_binding, output_shapes, output_buffers = self.prepare_io_binding(input_ids, attention_mask) - # Add the attention_mask inputs when needed - if "attention_mask" in self.input_names: - onnx_inputs["attention_mask"] = attention_mask.cpu().detach().numpy() + # run inference with binding & synchronize in case of multiple CUDA streams + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() + + # converts output to namedtuple for pipelines post-processing + return BaseModelOutput( + last_hidden_state=output_buffers["last_hidden_state"].view(output_shapes["last_hidden_state"]) + ) + else: + onnx_inputs = {"input_ids": input_ids.cpu().detach().numpy()} + + # Add the attention_mask inputs when needed + if "attention_mask" in self.input_names: + onnx_inputs["attention_mask"] = attention_mask.cpu().detach().numpy() - # Run inference - outputs = self.session.run(None, onnx_inputs) - last_hidden_state = torch.from_numpy(outputs[self.output_names["last_hidden_state"]]).to(self._device) + # Run inference + outputs = self.session.run(None, onnx_inputs) + last_hidden_state = torch.from_numpy(outputs[self.output_names["last_hidden_state"]]).to(self._device) - return BaseModelOutput(last_hidden_state=last_hidden_state) + return BaseModelOutput(last_hidden_state=last_hidden_state) def __call__(self, *args, **kwargs): return self.forward(*args, **kwargs) @@ -553,12 +657,208 @@ class ORTDecoder: The ONNX Runtime inference session associated to the decoder. """ - def __init__(self, session: onnxruntime.InferenceSession, device: torch.device): + def __init__( + self, + session: onnxruntime.InferenceSession, + config: transformers.PretrainedConfig, + device: torch.device, + use_io_binding: bool = True, + **kwargs + ): self.session = session + self.config = config self._device = device - 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.key_value_input_names = [key for key in self.input_names if "key_values" in key] + self.use_io_binding = use_io_binding + self.session_inputs = {output_key.name: idx for idx, output_key in enumerate(self.session.get_inputs())} + self.session_outputs = {output_key.name: idx for idx, output_key in enumerate(self.session.get_outputs())} + self.session_input_names = list(self.session_inputs.keys()) + self.session_output_names = list(self.session_outputs.keys()) + self.key_value_input_names = [key for key in self.session_input_names if "key_values" in key] + self.key_value_output_names = [key for key in self.session_output_names if "key_values" in key] + self.name_to_np_type = TypeHelper.get_io_numpy_type_map(self.session) if self.use_io_binding else None + + def prepare_output_buffer( + self, + output_name, + batch_size=None, + sequence_length=None, + encoder_sequence_length=None, + past_sequence_length=None, + is_self_attn=False, + ): + """ + Prepare the buffer of outputs(`logits`/`key_values`/`loss`) with 1D tensors. + """ + ort_type = TypeHelper.get_output_type(self.session, output_name) + torch_type = TypeHelper.ort_type_to_torch_type(ort_type) + if output_name == "loss": + output_shape = (1,) + output_buffer = torch.empty(1, dtype=torch_type, device=self._device) + elif output_name == "logits": + output_shape = (batch_size, sequence_length, self.config.vocab_size) + output_buffer = torch.empty(np.prod(output_shape), dtype=torch_type, device=self._device) + elif "key_values" in output_name: + num_heads = getattr(self.config, ORTConfigManager.get_num_heads_name(self.config.model_type)) + hidden_size = getattr(self.config, ORTConfigManager.get_hidden_size_name(self.config.model_type)) + embed_size_per_head = hidden_size // num_heads + if is_self_attn: + if past_sequence_length is not None: + sequence_length += past_sequence_length + output_shape = (batch_size, num_heads, sequence_length, embed_size_per_head) + else: + output_shape = (batch_size, num_heads, encoder_sequence_length, embed_size_per_head) + + output_buffer = torch.empty(np.prod(output_shape), dtype=torch_type, device=self._device) + + return output_shape, output_buffer + + def prepare_io_binding( + self, + input_ids: torch.LongTensor, + encoder_hidden_states: torch.FloatTensor, + encoder_attention_mask: Optional[torch.LongTensor] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + labels: Optional[torch.LongTensor] = None, + ): + io_binding = self.session.io_binding() + + # bind input ids + io_binding.bind_input( + "input_ids", + input_ids.device.type, + self._device.index, + self.name_to_np_type["input_ids"], + list(input_ids.size()), + input_ids.data_ptr(), + ) + + # bind encoder attention mask + io_binding.bind_input( + "encoder_attention_mask", + encoder_attention_mask.device.type, + self._device.index, + self.name_to_np_type["encoder_attention_mask"], + list(encoder_attention_mask.size()), + encoder_attention_mask.data_ptr(), + ) + + # bind encoder hidden states + if "encoder_hidden_states" in self.session_input_names: + io_binding.bind_input( + "encoder_hidden_states", + encoder_hidden_states.device.type, + self._device.index, + self.name_to_np_type["encoder_hidden_states"], + list(encoder_hidden_states.size()), + encoder_hidden_states.data_ptr(), + ) + + # bind past key values + if past_key_values is not None: + for input_name, past_key_value in zip(self.key_value_input_names, past_key_values): + io_binding.bind_input( + input_name, + past_key_value.device.type, + self._device.index, + self.name_to_np_type[input_name], + list(past_key_value.size()), + past_key_value.data_ptr(), + ) + + # bind labels + if "labels" in self.session_input_names: + io_binding.bind_input( + "labels", + labels.device.type, + self._device.index, + self.name_to_np_type["labels"], + list(labels.size()), + labels.data_ptr(), + ) + + # bind outputs + # bind logits + logits_shape, logits_buffer = self.prepare_output_buffer( + output_name="logits", + batch_size=input_ids.size(0), + sequence_length=input_ids.size(1), + ) + io_binding.bind_output( + "logits", + logits_buffer.device.type, + self._device.index, + self.name_to_np_type["logits"], + logits_shape, + logits_buffer.data_ptr(), + ) + output_shapes = {"logits": logits_shape} + output_buffers = {"logits": logits_buffer} + # bind loss + if "loss" in self.session_output_names: + loss_shape, loss_buffer = self.prepare_output_buffer(output_name="loss") + io_binding.bind_output( + "loss", + loss_buffer.device.type, + self._device.index, + self.name_to_np_type["loss"], + loss_shape, + loss_buffer.data_ptr(), + ) + output_shapes["loss"] = loss_shape + output_buffers["loss"] = loss_buffer + + # bind past key values + num_pkv = 4 # number of self-attention and cross-attention per decoder layer + for pkv_names_per_layer in [ + self.key_value_output_names[i : i + num_pkv] for i in range(0, len(self.key_value_output_names), num_pkv) + ]: + # bind a self attention and a cross-attention each time(2) + for i in range(2): + # bind self-attention past key values(2) + self_name = pkv_names_per_layer[i] + self_pkv_shape, self_pkv_buffer = self.prepare_output_buffer( + output_name=self_name, + batch_size=input_ids.size(0), + sequence_length=input_ids.size(1), + past_sequence_length=past_key_values[0].size(2) + if past_key_values + else None, # sequence length of self-attention key for layer.0 + is_self_attn=True, + ) + io_binding.bind_output( + self_name, + self_pkv_buffer.device.type, + self._device.index, + self.name_to_np_type[self_name], + self_pkv_shape, + self_pkv_buffer.data_ptr(), + ) + # set -1 for sequence_length as it could be larger than the real sequence_length for creating buffer + self_pkv_shape = self_pkv_shape[:2] + (-1,) + self_pkv_shape[3:] + output_shapes[self_name] = self_pkv_shape + output_buffers[self_name] = self_pkv_buffer + + # bind cross-attention past key values(2) + cross_name = pkv_names_per_layer[i + 2] + cross_pkv_shape, cross_pkv_buffer = self.prepare_output_buffer( + output_name=cross_name, + batch_size=input_ids.size(0), + encoder_sequence_length=encoder_hidden_states.size(1), + ) + io_binding.bind_output( + cross_name, + cross_pkv_buffer.device.type, + self._device.index, + self.name_to_np_type[cross_name], + cross_pkv_shape, + cross_pkv_buffer.data_ptr(), + ) + # set -1 for sequence_length as it could be larger than the real sequence_length for creating buffer + cross_pkv_shape = cross_pkv_shape[:2] + (-1,) + cross_pkv_shape[3:] + output_shapes[cross_name] = cross_pkv_shape + output_buffers[cross_name] = cross_pkv_buffer + + return io_binding, output_shapes, output_buffers @add_start_docstrings_to_model_forward(DECODER_INPUTS_DOCSTRING) def forward( @@ -569,47 +869,77 @@ def forward( past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, labels: Optional[torch.LongTensor] = None, ) -> Seq2SeqLMOutput: + # Flatten the past_key_values + if past_key_values is not None: + past_key_values = [past_key_value for pkv_per_layer in past_key_values for past_key_value in pkv_per_layer] - onnx_inputs = { - "input_ids": input_ids.cpu().detach().numpy(), - "encoder_attention_mask": encoder_attention_mask.cpu().detach().numpy(), - } + if self._device.type == "cuda" and self.use_io_binding: + io_binding, output_shapes, output_buffers = self.prepare_io_binding( + input_ids, encoder_hidden_states, encoder_attention_mask, past_key_values, labels + ) - # Add the encoder_hidden_states inputs when needed - if "encoder_hidden_states" in self.input_names: - onnx_inputs["encoder_hidden_states"] = encoder_hidden_states.cpu().detach().numpy() + # run inference with binding & synchronize in case of multiple CUDA streams + io_binding.synchronize_inputs() + self.session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() - if past_key_values is not None: - # Flatten the past_key_values - past_key_values = [past_key_value for pkv_per_layer in past_key_values for past_key_value in pkv_per_layer] - # Add the past_key_values to the decoder inputs - for input_name, past_key_value in zip(self.key_value_input_names, past_key_values): - onnx_inputs[input_name] = past_key_value.cpu().detach().numpy() + # 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) + past_key_values = tuple() + for name in self.session_output_names: + if "key_values" in name: + past_key_values += (output_buffers[name].view(output_shapes[name]),) - if "labels" in self.input_names: - # TODO: Any preprocessing like `self._shift_right(labels)`? - onnx_inputs["labels"] = labels.cpu().detach().numpy() + # Tuple of tuple of length `n_layers`, with each tuple of length equal to the number of self-attention and + # cross-attention per decoder layer + num_pkv = 4 + past_key_values = tuple(past_key_values[i : i + num_pkv] for i in range(0, len(past_key_values), num_pkv)) - # Run inference - outputs = self.session.run(None, onnx_inputs) + logits = output_buffers["logits"].view(output_shapes["logits"]) - # 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) - past_key_values = tuple( - torch.from_numpy(outputs[self.output_names[key]]).to(self._device) - for key in self.output_names - if "key_values" in key - ) + loss = None + if "loss" in self.session_output_names: + loss = output_buffers["loss"].view(output_shapes["loss"]) + else: + onnx_inputs = { + "input_ids": input_ids.cpu().detach().numpy(), + "encoder_attention_mask": encoder_attention_mask.cpu().detach().numpy(), + } + + # Add the encoder_hidden_states inputs when needed + if "encoder_hidden_states" in self.session_input_names: + onnx_inputs["encoder_hidden_states"] = encoder_hidden_states.cpu().detach().numpy() + + if past_key_values is not None: + # Add the past_key_values to the decoder inputs + for input_name, past_key_value in zip(self.key_value_input_names, past_key_values): + onnx_inputs[input_name] = past_key_value.cpu().detach().numpy() + + if "labels" in self.session_input_names: + # TODO: Any preprocessing like `self._shift_right(labels)`? + onnx_inputs["labels"] = labels.cpu().detach().numpy() + + # Run inference + outputs = self.session.run(None, onnx_inputs) + # 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) + past_key_values = tuple( + torch.from_numpy(outputs[self.session_outputs[key]]).to(self._device) + for key in self.session_output_names + if "key_values" in key + ) + + # Tuple of tuple of length `n_layers`, with each tuple of length equal to the number of self-attention and + # cross-attention per decoder layer + num_pkv = 4 + past_key_values = tuple(past_key_values[i : i + num_pkv] for i in range(0, len(past_key_values), num_pkv)) + logits = torch.from_numpy(outputs[self.session_outputs["logits"]]).to(self._device) - # Tuple of tuple of length `n_layers`, with each tuple of length equal to the number of self-attention and - # cross-attention per decoder layer - num_pkv = 4 - past_key_values = tuple(past_key_values[i : i + num_pkv] for i in range(0, len(past_key_values), num_pkv)) - logits = torch.from_numpy(outputs[self.output_names["logits"]]).to(self._device) + loss = None + if "loss" in self.session_output_names: + loss = torch.from_numpy(outputs[self.session_outputs["loss"]]).to(self._device) - loss = None - if "loss" in self.output_names: - loss = torch.from_numpy(outputs[self.output_names["loss"]]).to(self._device) + # converts output to namedtuple for pipelines post-processing return Seq2SeqLMOutput(loss=loss, logits=logits, past_key_values=past_key_values) def __call__(self, *args, **kwargs): diff --git a/optimum/onnxruntime/optimization.py b/optimum/onnxruntime/optimization.py index f625569a6b..c0f1a70428 100644 --- a/optimum/onnxruntime/optimization.py +++ b/optimum/onnxruntime/optimization.py @@ -112,7 +112,7 @@ def optimize( save_dir = Path(save_dir) save_dir.mkdir(parents=True, exist_ok=True) model_type = self.config.model_type - ORTConfigManager.check_supported_model_or_raise(model_type) + ORTConfigManager.check_optimization_supported_model_or_raise(model_type) # Save the model configuration self.config.save_pretrained(save_dir) diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py index da6a2e0974..47bdf25a8f 100644 --- a/optimum/onnxruntime/utils.py +++ b/optimum/onnxruntime/utils.py @@ -62,6 +62,7 @@ class ORTConfigManager: "bart": ("encoder_attention_heads", "d_model", "bart"), "bert": ("num_attention_heads", "hidden_size", "bert"), "big_bird": ("num_attention_heads", "hidden_size", "bert"), + "bigbird_pegasus": ("encoder_attention_heads", "d_model", None), # bug in `fusion_skiplayernorm.py` "camembert": ("num_attention_heads", "hidden_size", "bert"), "codegen": ("n_head", "n_embd", "gpt2"), "deberta": ("num_attention_heads", "hidden_size", "bert"), @@ -70,9 +71,12 @@ class ORTConfigManager: "electra": ("num_attention_heads", "hidden_size", "bert"), "gpt2": ("n_head", "n_embd", "gpt2"), "gpt_neo": ("num_heads", "hidden_size", "gpt2"), - "mt5": ("num_heads", "d_model", "bart"), "marian": ("encoder_attention_heads", "d_model", "bart"), + "mbart": ("encoder_attention_heads", "d_model", "bart"), + "mt5": ("num_heads", "d_model", "bart"), + "m2m_100": ("encoder_attention_heads", "d_model", "bart"), "roberta": ("num_attention_heads", "hidden_size", "bert"), + "t5": ("num_heads", "d_model", "t5"), "xlm-roberta": ("num_attention_heads", "hidden_size", "bert"), } @@ -116,6 +120,15 @@ def check_supported_model_or_raise(cls, model_type: str) -> bool: f"If you want to support {model_type} please propose a PR or open up an issue." ) + @classmethod + def check_optimization_supported_model_or_raise(cls, model_type: str) -> bool: + supported_model_types_for_optimization = ["bert", "gpt2", "bart"] + if (model_type not in cls._conf) or (cls._conf[model_type][2] not in supported_model_types_for_optimization): + raise KeyError( + f"ONNX Runtime doesn't support the graph optimization of {model_type} yet. Only {supported_model_types_for_optimization} are supported. " + f"If you want to support {model_type} please propose a PR or open up an issue in ONNX Runtime:https://github.com/microsoft/onnxruntime." + ) + def generate_identified_filename(filename, identifier): return filename.parent.joinpath(filename.stem + identifier).with_suffix(filename.suffix) diff --git a/tests/onnxruntime/test_modeling.py b/tests/onnxruntime/test_modeling.py index 58d9265cb9..b2ac9490d9 100644 --- a/tests/onnxruntime/test_modeling.py +++ b/tests/onnxruntime/test_modeling.py @@ -527,6 +527,33 @@ def test_pipeline_on_gpu(self, model_arch): gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForQuestionAnswering.from_pretrained( + model_id, from_transformers=True, use_io_binding=False + ) + set_seed(SEED) + io_model = ORTModelForQuestionAnswering.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model(**tokens) + io_outputs = io_model(**tokens) + + self.assertTrue("start_logits" in io_outputs) + self.assertTrue("end_logits" in io_outputs) + self.assertIsInstance(io_outputs.start_logits, torch.Tensor) + 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)) + + gc.collect() + class ORTModelForSequenceClassificationIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES = ( @@ -632,6 +659,32 @@ def test_pipeline_zero_shot_classification(self): self.assertTrue(all(score > 0.0 for score in outputs["scores"])) self.assertTrue(all(isinstance(label, str) for label in outputs["labels"])) + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForSequenceClassification.from_pretrained( + model_id, from_transformers=True, use_io_binding=False + ) + set_seed(SEED) + io_model = ORTModelForSequenceClassification.from_pretrained( + model_id, from_transformers=True, use_io_binding=True + ) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model(**tokens) + io_outputs = io_model(**tokens) + + 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)) + + gc.collect() + class ORTModelForTokenClassificationIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES = ( @@ -715,6 +768,32 @@ def test_pipeline_on_gpu(self, model_arch): gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForTokenClassification.from_pretrained( + model_id, from_transformers=True, use_io_binding=False + ) + set_seed(SEED) + io_model = ORTModelForTokenClassification.from_pretrained( + model_id, from_transformers=True, use_io_binding=True + ) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model(**tokens) + io_outputs = io_model(**tokens) + + 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)) + + gc.collect() + class ORTModelForFeatureExtractionIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES = ( @@ -795,6 +874,30 @@ def test_pipeline_on_gpu(self, model_arch): gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForFeatureExtraction.from_pretrained( + model_id, from_transformers=True, use_io_binding=False + ) + set_seed(SEED) + io_model = ORTModelForFeatureExtraction.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model(**tokens) + io_outputs = io_model(**tokens) + + self.assertTrue("last_hidden_state" in io_outputs) + 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)) + + gc.collect() + class ORTModelForMultipleChoiceIntegrationTest(unittest.TestCase): # Multiple Choice tests are conducted on different models due to mismatch size in model's classifier @@ -843,6 +946,38 @@ def test_compare_to_transformers(self, model_id): gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_id): + set_seed(SEED) + onnx_model = ORTModelForMultipleChoice.from_pretrained(model_id, from_transformers=True, use_io_binding=False) + set_seed(SEED) + io_model = ORTModelForMultipleChoice.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + 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" + 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")) + + 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)) + + gc.collect() + class ORTModelForCausalLMIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES = ("gpt2",) @@ -944,6 +1079,47 @@ def test_pipeline_on_gpu(self, model_arch): gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True, use_io_binding=False) + set_seed(SEED) + io_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model(**tokens) + io_outputs = io_model(**tokens) + + 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)) + + gc.collect() + + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_generation_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True, use_io_binding=False) + set_seed(SEED) + io_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model.generate(**tokens) + io_outputs = io_model.generate(**tokens) + + # compare tensor outputs + self.assertTrue(torch.equal(onnx_outputs, io_outputs)) + + gc.collect() + class ORTModelForImageClassificationIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = { @@ -1027,6 +1203,34 @@ def test_pipeline_on_gpu(self, *args, **kwargs): gc.collect() + @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items()) + @require_torch_gpu + def test_compare_to_io_binding(self, *args, **kwargs): + model_arch, model_id = args + set_seed(SEED) + onnx_model = ORTModelForImageClassification.from_pretrained( + model_id, from_transformers=True, use_io_binding=False + ) + set_seed(SEED) + io_model = ORTModelForImageClassification.from_pretrained( + model_id, from_transformers=True, use_io_binding=True + ) + + 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, return_tensors="pt") + 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)) + + gc.collect() + class ORTModelForSeq2SeqLMIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES = ( @@ -1169,6 +1373,50 @@ def test_compare_with_and_without_past_key_values_model_outputs(self): outputs_model_without_pkv = model_without_pkv.generate(**tokens) self.assertTrue(torch.equal(outputs_model_with_pkv, outputs_model_without_pkv)) + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True, use_io_binding=False) + set_seed(SEED) + io_model = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + decoder_start_token_id = onnx_model.config.decoder_start_token_id if model_arch != "mbart" else 2 + decoder_inputs = {"decoder_input_ids": torch.ones((1, 1), dtype=torch.long) * decoder_start_token_id} + + onnx_outputs = onnx_model(**tokens, **decoder_inputs) + io_outputs = io_model(**tokens, **decoder_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)) + + gc.collect() + + @parameterized.expand(SUPPORTED_ARCHITECTURES) + @require_torch_gpu + def test_compare_generation_to_io_binding(self, model_arch): + model_id = MODEL_NAMES[model_arch] + set_seed(SEED) + onnx_model = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True, use_io_binding=False) + set_seed(SEED) + io_model = ORTModelForSeq2SeqLM.from_pretrained(model_id, from_transformers=True, use_io_binding=True) + + tokenizer = get_preprocessor(model_id) + tokens = tokenizer("This is a sample output", return_tensors="pt") + onnx_outputs = onnx_model.generate(**tokens) + io_outputs = io_model.generate(**tokens) + + # compare tensor outputs + self.assertTrue(torch.equal(onnx_outputs, io_outputs)) + + gc.collect() + class ORTModelForCustomTasksIntegrationTest(unittest.TestCase): SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = { diff --git a/tests/onnxruntime/test_optimization.py b/tests/onnxruntime/test_optimization.py index 315e727bd7..673506b1b8 100644 --- a/tests/onnxruntime/test_optimization.py +++ b/tests/onnxruntime/test_optimization.py @@ -77,7 +77,12 @@ def test_compare_original_model_with_optimized_model(self, model_cls, model_name (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-bart", True), (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-marian", False), (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-marian", True), + (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-mbart", False), + (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-mbart", True), (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-onnx-mt5", False), + (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-onnx-mt5", True), + (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-m2m_100", False), + (ORTModelForSeq2SeqLM, "hf-internal-testing/tiny-random-m2m_100", True), ) @parameterized.expand(SUPPORTED_SEQ2SEQ_ARCHITECTURES_WITH_MODEL_ID)