[Performance] GPU op placement control when some ops must be on the CPU

[Craigacp]

Describe the issue

We export ONNX transformer encoder models in OML4Py with the tokenizer attached to the bottom of the model, so the ONNX model accepts a string tensor input and returns the embedding vector. We use the tokenizer operations from onnxruntime-extensions, which are CPU only, and have wrapped around an ONNX graph which batches, pads and truncates the tokenized representation using the SequenceMap operation.

When increasing the batch size we've noticed that much of the runtime is spent in the Loop op which SequenceMap uses, which is very odd considering it doesn't actually do very much. After some investigation we determined that this was due to most of the ops in the tokenizer graph being placed on the GPU rather than the CPU, even though the subgraph we're looping over must be placed on the CPU due to the presence of the BertTokenizer op. We would like the whole tokenization graph including the Loop op to be placed on the CPU EP, but there doesn't appear to be a way to control op placement in the C API. Alternatively there might be a bug in the way that ops are placed on the CPU, as this code looks like it should fallback the whole subgraph & loop to CPU, but I don't understand it well enough to see if that's part of the issue.

To reproduce

Run the supplied bert-tok.onnx graph with a batch size of 100 with the CUDA EP enabled, most of the runtime is spent in the tokenization loop operation transferring data between CPU and GPU.

Urgency

This performance issue prevents the use of GPUs to accelerate our models.

Platform

Linux

OS Version

OL8

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.20.0

ONNX Runtime API

C

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

No response

Model File

bert-tok.onnx.zip

Is this a quantized model?

No

[github-actions]

This issue has been automatically marked as stale due to inactivity and will be closed in 30 days if no further activity occurs. If further support is needed, please provide an update and/or more details.

[Craigacp]

@snnn we would appreciate someone looking at this, it blocks us using GPUs for inference.

[skottmckay]

The second SequenceMap subgraph looks off to me as I don't think we should be going to GPU for a single Sub node that is downstream of a Shape.

This is what it looks like with CUDA EP enabled

@yuslepukhin did this potentially change as part of #22766?

[yuslepukhin]

@yuslepukhin did this potentially change as part of #22766?

Unlikely, since Resource Aware Partitioning must be explicitly enabled.

[Craigacp]

I have a prototype which allows CPU pinning of specific named ops, we're testing it internally at the moment and I'd like to keep this issue open to refer to in the PR when we're done.