BLS script + FORCE_CPU_ONLY_INPUT_TENSORS -> output tensor from ORT is NEVER on GPU memory

[pommedeterresautee]

Description
In a Python script (BLS), when I retrieve output of a GPU model (ONNX Runtime), it is always a CPU tensor (on 64 ONNX Runtime calls in row).
DLPack is used to move tensor from/to ONNX Runtime.
Because of that, I need to move the result back to CUDA memory to resend it to the model (generative language model use case).
This move from CUDA -> CPU HOST -> CUDA takes time during inference.
FORCE_CPU_ONLY_INPUT_TENSORS is set

The flow works (minus the memory issue) aka it doesn't crash, etc. The result is as expected.

According to this example, there is nothing special to do to keep tensor on GPU, but the example doesn't imply calling external model, it's just the script calling itself.
https://github.com/triton-inference-server/server/blob/fd7ceefeb5096add392b12f829b7d77b2b59b73c/qa/python_models/dlpack_io_identity/model.py

Triton Information
What version of Triton are you using? nvcr.io/nvidia/tritonserver:21.12-py3

Are you using the Triton container or did you build it yourself? docker image from Nvidia repo

To Reproduce

config of BLS script (i tried with kind: KIND_GPU too... in case of...)

name: "transformer_generative_model"
max_batch_size: 0
backend: "python"

input [
    {
        name: "TEXT"
        data_type: TYPE_STRING
        dims: [ -1 ]
    }
]

output [
    {
        name: "output"
        data_type: TYPE_STRING
        dims: [ -1 ]
    }
]

instance_group [
    {
      count: 1
      kind: KIND_CPU
    }
]

parameters: {
  key: "FORCE_CPU_ONLY_INPUT_TENSORS"
  value: {
    string_value:"no"
  }
}

ONNX Runtime config

name: "transformer_onnx_model"
max_batch_size: 0
platform: "onnxruntime_onnx"
default_model_filename: "model.bin"

input [
    {
        name: "input_ids"
        data_type: TYPE_INT32
        dims: [-1, -1]
    }
]

output {
    name: "output"
    data_type: TYPE_FP32
    dims: [-1, -1, 50257]
}

instance_group [
    {
      count: 1
      kind: KIND_GPU
    }
]

The BLS script

This function is called 64 times by the BLS script with input tensor which is always on CUDA memory.
There are print function calls -> they are always true (the tensor is always on CPU) which is the issue I try to fix/understand

...
        def inference_triton(input_ids: torch.Tensor) -> torch.Tensor:
            print(f"input_ids device: {input_ids.device}")  # always cuda
            input_ids = input_ids.type(dtype=torch.int32)
            inputs = [pb_utils.Tensor.from_dlpack("input_ids", torch.to_dlpack(input_ids))]
            inference_request = pb_utils.InferenceRequest(
                model_name='transformer_onnx_model',
                requested_output_names=['output'],
                inputs=inputs)
            inference_response = inference_request.exec()
            if inference_response.has_error():
                raise pb_utils.TritonModelException(inference_response.error().message())
            else:
                output = pb_utils.get_output_tensor_by_name(inference_response, 'output')
                print(f"is cpu {output.is_cpu()}")  # always true :-(
                tensor: torch.Tensor = torch.from_dlpack(output.to_dlpack())
                print(f'input device: {tensor.device}')  # always CPU :-(
                tensor = tensor.cuda()  # takes time :-(
                return tensor
...

Expected behavior

I understand that with option FORCE_CPU_ONLY_INPUT_TENSORS there is no guarantee that the output tensor from ONNX Runtime model is in CUDA memory, but I expect that it is the case most of the time, or at least some times.

[Tabrizian]

Thanks for filing a detailed issue about this. This is something that we are aware of and there is a ticket for it here to fix this:

triton-inference-server/onnxruntime_backend#73

The problem is that the ORT backend always returns the output in CPU and that's why the BLS response is always in CPU.

[pommedeterresautee]

Thank you a lot for the (very fast) response.
Is there a similar issue with TensorRT engine?
Moreover, I understand that the fix is coming in Triton 22.02, may I ask if there is a (even very approximative) date for its release as a docker image?

[Tabrizian]

I assume by TRT engine you mean the TensorRT backend models? I don't think there is a similar issue with the TRT models.

Moreover, I understand that the fix is coming in Triton 22.02, may I ask if there is a (even very approximative) date for its release as a docker image?

cc @dzier

[CoderHam]

@pommedeterresautee 22.02 will be released in a month from now.

[CoderHam]

@Tabrizian can this be closed?

[Tabrizian]

Yes, I think this can be closed since there is already a GitHub issue in onnxruntime backend repo about this.

[pommedeterresautee]

@Tabrizian just for information, tried Tensorrt and there is indeed no issue (tensors are kept on gpu memory)

[Tabrizian]

@pommedeterresautee Thanks for letting us know!