[Training] Can we use ORTModule for inference?

[LSC527]

Describe the issue

I see ORTModule is easy to use for pytorch model training. When it comes to inference, it seems that we have to write so much more code to do inference with torch tensors. An example code from official docs:

# X is a PyTorch tensor on device
session = onnxruntime.InferenceSession('model.onnx', providers=['CUDAExecutionProvider', 'CPUExecutionProvider']))
binding = session.io_binding()

X_tensor = X.contiguous()

binding.bind_input(
    name='X',
    device_type='cuda',
    device_id=0,
    element_type=np.float32,
    shape=tuple(x_tensor.shape),
    buffer_ptr=x_tensor.data_ptr(),
    )

## Allocate the PyTorch tensor for the model output
Y_shape = ... # You need to specify the output PyTorch tensor shape
Y_tensor = torch.empty(Y_shape, dtype=torch.float32, device='cuda:0').contiguous()
binding.bind_output(
    name='Y',
    device_type='cuda',
    device_id=0,
    element_type=np.float32,
    shape=tuple(Y_tensor.shape),
    buffer_ptr=Y_tensor.data_ptr(),
)

session.run_with_iobinding(binding)

So can ORTModule support inference with providers?

To reproduce

:)

Urgency

No response

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.17.1

PyTorch Version

2.2.0

Execution Provider

CUDA

Execution Provider Library Version

No response

[xadupre]

You should be able to ORTModule for inference but you are right, this piece of code is missing. Below the code I use. It replicates what ORTModule does. It relies on two functions only available in onnxruntime-training and not really documented. That's why it is difficult to find such code in the documentation but that's the most efficient way. We should change that.

    def _get_ortvalues_from_torch_tensors(
        self, tensors: Tuple["torch.Tensor", ...]  # noqa: F821
    ) -> Tuple[Tuple["torch.Tensor", ...], Tuple["OrtDevice", ...], Any]:  # noqa: F821
        ortvalues = self.OrtValueVector()
        ortvalues.reserve(len(tensors))
        dtypes = []
        shapes = []
        data_ptrs = []
        devices = []
        dimensions = []

        max_device = -1
        assert isinstance(max_device, int), f"unexpected type for device={max_device!r}"
        assert tensors is not None, "tensors cannot be None"
        new_tensors = []
        for tensor, (dim, rk, name) in zip(tensors, self.is_dimension_in):
            if dim:
                assert isinstance(
                    tensor, (int, torch.SymInt)
                ), f"Unexpected type {type(tensor)} for name={name!r}."
                dtypes.append(np.int64)
                ti = int(tensor)
                assert ti != 0, (
                    f"Null value for a dimension ti={ti}, "
                    f"tensor={tensor}, rk={rk}, name={name!r}, "
                    f"type(tensor)={type(tensor)}, "
                    f"dimension={[t for t in tensors if isinstance(t, (int,torch.SymInt))]}"
                )
                t = torch.tensor([ti] if rk == 1 else ti, dtype=torch.int64)
                devices.append(self.devices[-1])
                new_tensors.append(t)
                dimensions.append(t)
                shapes.append(t.size())
                data_ptrs.append(t.data_ptr())
            else:
                assert isinstance(tensor, torch.Tensor), (
                    f"Unexpected type {type(tensor)}, dim={dim}, rk={rk}, name={name!r}, "
                    f"len(tensors)={len(tensors)}, len(is_dimension_in)={len(self.is_dimension_in)}"
                )
                dtypes.append(self.TORCH_DTYPE_TO_NUMPY_DTYPE[tensor.dtype])
                shapes.append(tensor.size())
                data_ptrs.append(tensor.data_ptr())
                d = tensor.get_device()
                devices.append(self.devices[d])
                new_tensors.append(tensor)
                max_device = max(max_device, tensor.get_device())

        ortvalues.push_back_batch(new_tensors, data_ptrs, dtypes, shapes, devices)
        output_devices = []
        for dim, rk, name in self.is_dimension_out:
            if dim:
                dev = self.devices[-1]
            else:
                dev = self.devices[max_device]
            output_devices.append(dev)

        return (ortvalues, output_devices, dimensions)

    def _ortvalues_to_torch_tensor(
        self, ortvalues: "onnxruntime.OrtValueVector"  # noqa: F821
    ) -> Tuple["torch.Tensor", ...]:  # noqa: F821
        if len(ortvalues) == 0:
            return tuple()

        res = ortvalues.to_dlpacks(self.from_dlpack)
        return tuple(
            _post_process(r, d[2], d[0]) for r, d in zip(res, self.is_dimension_out)
        )

    def _run_onnx_session_with_ortvaluevector(
        self, inputs: Tuple["torch.Tensor", ...]  # noqa: F821
    ) -> Tuple["torch.Tensor"]:  # noqa: F821
        # _nvtx_range_push("contiguous")
        contiguous_inputs = tuple(
            (a.contiguous() if isinstance(a, torch.Tensor) else a) for a in inputs
        )
        # _nvtx_range_pop()

        # _nvtx_range_push("push_back_batch")
        ort_inputs, output_devices, dimensions = self._get_ortvalues_from_torch_tensors(
            contiguous_inputs
        )
        # _nvtx_range_pop()

        # _nvtx_range_push("run_with_ortvaluevector")
        ort_outputs = self.OrtValueVector()
        self.sess.run_with_ortvaluevector(
            self.run_options,
            self.input_names,
            ort_inputs,
            self.output_names,
            ort_outputs,
            output_devices,
        )

        # _nvtx_range_pop()

        # _nvtx_range_push("after run_with_ortvaluevector")
        # Map ORTValue to torch.Tensor.
        pth_outputs = self._ortvalues_to_torch_tensor(ort_outputs)
        # _nvtx_range_pop()

        return pth_outputs, dimensions

[pengwa]

Yeah, ORTModule already supports training/evaluating with a single wrap, when you want an acceleration from inside a PyTorch based model inferencing code. ORTModule internally will export the model and handle model optimization.

While if your want to inference with an exported ONNX model, the above code snippet is something you can try.

[LSC527]

@xadupre Yeah, I think onnxruntime should support pytorch model inference directly. That is a very common use case.

[LSC527]

@pengwa Thx for explanation. I wonder what optimization will ORTModule do when inferencing? Can we still choose EPs like TensorrtExecutionProvider ?

[pengwa]

@pengwa Thx for explanation. I wonder what optimization will ORTModule do when inferencing? Can we still choose EPs like TensorrtExecutionProvider ?

ORTModule internally used InferenceSession,initilize, and Run. You can understand what InferenceSession optimized will be applied also. But there indeed some differences for example all weights/parameters are treated as graph inputs instead of constant in current impl.

We did not expose too many options to users, aiming at a transparent training engine

onnxruntime/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py

Line 221 in a0d5067

if self._device.type == "cuda":

. Though ideally we can add TensorRT ep if there is a need for it. Would you mind sharing more backaround on your requirements?

[LSC527]

@pengwa I am trying to accelerate the inference of pytorch models directly.
I have tried onnxruntime's InferenceSession with several EPs like TensorrtExecutionProvider and OpenVINOExecutionProvider, which gain very promising inference speedup. But I have to export the model to onnx manully and create io_bindings when inference with torch tensor inputs.
Later I see ORTModule is convenient to train pytorch models directly. I also find a ORTInferenceModule (link) support inference directly but only support OpenVINOExecutionProvider and ORTInferenceModule seems not mentioned in onnxruntime docs. This is confusing.
So ideally it will be best that onnxruntime provide ORTInferenceModule with full EP support as onnxruntime.InferenceSession do, while more convenient to do inference with pytorch models and tensors.