ONNX Runtime中io_binding.bind_input报错：无设备间数据传输注册

问题场景 #

在Nvidia Triton推理服务器环境中，使用ONNX Runtime的IO Binding绑定PyTorch CUDA张量进行推理时，触发以下错误：

RuntimeError: Error when binding input: There's no data transfer registered for copying tensors from Device:[DeviceType:1 MemoryType:0 DeviceId:0] to Device:[DeviceType:0 MemoryType:0 DeviceId:0]

核心原因 #

ONNX Runtime的CUDA Execution Provider（EP）无法直接识别PyTorch的CUDA内存指针，本质是两者的CUDA上下文未共享，或是ORT未启用PyTorch CUDA张量的直接绑定支持。

解决方法 #

方法1：配置共享CUDA上下文并启用默认流拷贝 #

通过ORT Session配置，确保与PyTorch共享CUDA上下文，同时启用默认流中的数据拷贝，让ORT能识别PyTorch的CUDA内存：

import onnxruntime as ort
import torch
import numpy as np

# 配置Session选项，启用图优化并指定CUDA EP参数
so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(
    "path/to/onnx.model",
    sess_options=so,
    providers=[
        ("CUDAExecutionProvider", {
            "device_id": 0,
            "do_copy_in_default_stream": True,  # 关键：启用默认流拷贝，对齐PyTorch流
            "gpu_mem_limit": 2 * 1024 * 1024 * 1024
        }),
        "CPUExecutionProvider"
    ]
)

io_binding = session.io_binding()
pt_script_embeddings = torch.rand(
    size=(100, 2010), dtype=torch.float32, device="cuda:0"
).contiguous()

# 绑定输入，保持原参数配置
io_binding.bind_input(
    name="np_script_embeddings",
    device_type="cuda",
    device_id=0,
    element_type=np.float32,
    shape=tuple(pt_script_embeddings.shape),
    buffer_ptr=pt_script_embeddings.data_ptr(),
)

# 绑定输出逻辑不变
logit_output_shape = (100, 2)
logit_output = torch.empty(logit_output_shape, dtype=torch.float32, device='cuda:0').contiguous()
io_binding.bind_output(
    name="np_logits",
    device_type="cuda",
    device_id=0,
    element_type=np.float32,
    shape=tuple(logit_output.shape),
    buffer_ptr=logit_output.data_ptr()
)

# 同步CUDA流确保操作顺序
torch.cuda.synchronize()
session.run_with_iobinding(io_binding)
torch.cuda.synchronize()

outputs = logit_output.cpu().numpy()

方法2：临时中转CPU张量（低性能要求场景） #

如果方法1不生效，可先将PyTorch CUDA张量拷贝到CPU，再让ORT自动将数据传输到GPU：

io_binding = session.io_binding()
pt_script_embeddings = torch.rand(
    size=(100, 2010), dtype=torch.float32, device="cuda:0"
).contiguous()

# 中转到CPU，ORT自动拷贝至GPU
cpu_embeddings = pt_script_embeddings.cpu().numpy()
io_binding.bind_input(name="np_script_embeddings", value=cpu_embeddings)

# 输出绑定逻辑不变
logit_output_shape = (100, 2)
logit_output = torch.empty(logit_output_shape, dtype=torch.float32, device='cuda:0').contiguous()
io_binding.bind_output(
    name="np_logits",
    device_type="cuda",
    device_id=0,
    element_type=np.float32,
    shape=tuple(logit_output.shape),
    buffer_ptr=logit_output.data_ptr()
)

session.run_with_iobinding(io_binding)
outputs = logit_output.cpu().numpy()

方法3：对齐ONNX Runtime与CUDA版本兼容性 #

Triton 23.12镜像内置CUDA 12.1，需安装适配该版本的onnxruntime-gpu，修改Dockerfile：

FROM --platform=linux/amd64 nvcr.io/nvidia/tritonserver:23.12-py3

# 指定适配CUDA 12.1的onnxruntime-gpu版本
RUN pip install torch transformers onnx onnxruntime-gpu==1.16.3 onnxruntime

关键注意点 #

• 确保device_type参数严格为小写'cuda'，与ORT的设备命名规则一致。
• 绑定的输入输出名称必须与ONNX模型中的节点名称完全匹配。
• 推理前后执行torch.cuda.synchronize()，避免CUDA流异步操作导致的上下文冲突。

内容的提问来源于stack exchange，提问作者Dan M