Add IOBinding support to ONNX Runtime module

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 IOBindingHelper, TypeHelper

optimum/onnxruntime/io_binding/io_binding_helper.py

@@ -0,0 +1,170 @@
+#  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 logging
+import traceback
+from typing import Dict
+
+import numpy as np
+import torch
+
+import onnxruntime as ort
+import pkg_resources
+from onnxruntime import InferenceSession
+from onnxruntime.capi.onnxruntime_inference_collection import OrtValue
+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):
+    """
+    TypeHelper helps get data type information of the ONNX runtime inference session and offers the mapping from
+    OrtValues data type to the data type of other data structures(numpy, torch tensor...).
+    """
+
+    # 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 not in ort_type_to_numpy_type_map:
+            raise ValueError(f"{ort_type} not found in map")
+        return ort_type_to_numpy_type_map[ort_type]
+
+    @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 not in ort_type_to_torch_type_map:
+            raise ValueError(f"{ort_type} not found in map")
+        return ort_type_to_torch_type_map[ort_type]
+
+
+# adapted from https://github.com/microsoft/onnxruntime/blob/1ab11a111ce0717bfbfaca964d04a017cb9b1752/onnxruntime/python/tools/transformers/io_binding_helper.py#L97
+class IOBindingHelper:
+    """
+    IOBindingHelper is a class that helps ORTModels to create buffers for inputs and outputs of an ONNX runtime
+    inference session when using devices like GPU for acceleration. It helps reduce memory copy between the host
+    and device.
+    """
+
+    def __init__(self, model: ort.InferenceSession, config, device, **kwargs):
+        self.model = model
+        self.config = config
+        self.device = device
+        # create {name:idx} dict for model outputs
+        self.model_inputs = {output_key.name: idx for idx, output_key in enumerate(model.get_inputs())}
+        self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(model.get_outputs())}
+        self.model_input_names = list(self.model_inputs.keys())
+        self.model_output_names = list(self.model_outputs.keys())
+
+    def prepare_io_binding(self, **kwargs):
+        """Returnas IO binding object for a session."""
+
+        name_to_np_type = TypeHelper.get_io_numpy_type_map(self.model)
+
+        # Bind inputs and outputs to onnxruntime session
+        io_binding = self.model.io_binding()
+
+        # Bind inputs
+        for input_name in self.model_input_names:
+            onnx_input = kwargs.pop(input_name)
+
+            assert onnx_input.is_contiguous()
+            io_binding.bind_input(
+                input_name,
+                onnx_input.device.type,
+                self.device.index,
+                name_to_np_type[input_name],
+                list(onnx_input.size()),
+                onnx_input.data_ptr(),
+            )
+
+        # Bind outputs
+        for name in self.model_output_names:
+            io_binding.bind_output(name, self.device.type, device_id=self.device.index)
+
+        return io_binding
+
+    @staticmethod
+    def to_pytorch(ort_value: OrtValue) -> torch.Tensor:
+        """Converts tensors held by OrtValues to torch tensor."""
+        env = {pkg.key for pkg in pkg_resources.working_set}
+        if "onnxruntime-training" in env:
+            return IOBindingHelper.to_pytorch_via_dlpack(ort_value)
+        else:
+            try:
+                return IOBindingHelper.to_pytorch_via_cupy(ort_value)
+            except Exception as e:
+                logging.error(traceback.format_exc())
+                logging.info("Unable to access output memory in CUDA, will offload to CPU")
+                return IOBindingHelper.to_pytorch_via_np(ort_value)
+
+    @staticmethod
+    def to_pytorch_via_np(ort_value: OrtValue) -> torch.Tensor:
+        ort_device = ort_value.device_name().lower()
+        return torch.tensor(ort_value.numpy()).to(ort_device)
+
+    @staticmethod
+    def to_pytorch_via_cupy(ort_value: OrtValue) -> torch.Tensor:
+        ort_device = ort_value.device_name().lower()
+        assert ort_device == "cuda", f"Convert via CuPy only when device is CUDA, got: {ort_device}"
+
+        ort_type = ort_value.data_type()
+        np_type = TypeHelper.ort_type_to_numpy_type(ort_type)
+
+        # Access CUDA memory via CuPy
+        import cupy as cp
+
+        memory = cp.cuda.UnownedMemory(ort_value.data_ptr(), 0, None)
+        memory_ptr = cp.cuda.MemoryPointer(memory, 0)
+        cp_array = cp.ndarray(shape=ort_value.shape(), memptr=memory_ptr, dtype=np_type)
+        torch_tensor = torch.from_dlpack(cp_array.toDlpack())
+
+        # If is boolean, the dtype will be uint8 and need to be convert back to bool.
+        if "bool" in ort_type:
+            torch_tensor = torch_tensor.to(torch.bool)
+
+        torch_tensor = torch_tensor.clone()
+
+        return torch_tensor
+
+    @staticmethod
+    # only `onnxruntime-training` supports dlpack for OrtValue
+    def to_pytorch_via_dlpack(ort_value: OrtValue) -> torch.Tensor:
+        from torch._C import _from_dlpack
+
+        torch_tensor = ort_value.to_dlpacks(_from_dlpack)
+        return torch_tensor
+
+    @staticmethod
+    def get_device_index(device):
+        if isinstance(device, str):
+            # could be 'cuda:0', 'cuda:1', or 'cpu'. with cpu, set index=0
+            device = torch.device(device)
+        elif isinstance(device, int):
+            return device
+        return 0 if device.index is None else device.index

optimum/onnxruntime/modeling_ort.py

@@ -49,6 +49,7 @@
 from huggingface_hub import HfApi, hf_hub_download
 
 from ..modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel
+from .io_binding import IOBindingHelper
 from .utils import ONNX_WEIGHTS_NAME, get_device_for_provider, get_provider_for_device
 
 
@@ -66,6 +67,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 use IO Binding during inference to avoid memory copy between the host and devices. Will default to True if the device is CUDA, otherwise default to False.
 """
 
 ONNX_TEXT_INPUTS_DOCSTRING = r"""
@@ -107,7 +109,7 @@ class ORTModel(OptimizedModel):
     base_model_prefix = "onnx_model"
     auto_model_class = AutoModel
 
-    def __init__(self, model: ort.InferenceSession = None, config=None, **kwargs):
+    def __init__(self, model: ort.InferenceSession = None, config=None, use_io_binding=True, **kwargs):
         self.model = model
         self.config = config
         self.model_save_dir = kwargs.get("model_save_dir", None)
@@ -121,6 +123,11 @@ def __init__(self, model: ort.InferenceSession = None, config=None, **kwargs):
                 f" Use `ort_model.to()` to send the outputs to the wanted device."
             )
 
+        if self.device.type == "cuda" and use_io_binding:
+            self.use_io_binding = True
+        else:
+            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)
@@ -859,8 +866,8 @@ 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())}
@@ -888,16 +895,34 @@ 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.use_io_binding:
+            onnx_inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
+            io_helper = IOBindingHelper(self.model, self.config, self.device)
+            io_binding = io_helper.prepare_io_binding(**onnx_inputs)
+
+            # run inference with binding
+            io_binding.synchronize_inputs()
+            self.model.run_with_iobinding(io_binding)
+            io_binding.synchronize_outputs()
+
+            # map outputs with names
+            outputs = {}
+            for name, output in zip(io_helper.model_output_names, io_binding._iobinding.get_outputs()):
+                outputs[name] = IOBindingHelper.to_pytorch(output)
+
+            # converts output to namedtuple for pipelines post-processing
+            return CausalLMOutputWithCrossAttentions(**outputs)
+        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(