[Feature request] WebGPU support

[loretoparisi]

WebGPU

Chrome shipped WebGPU today in Chrome 113 Beta.

Reason for request

WebGPU is currently a WIP in Firefox and Safari, in addition to the beta Chrome. Also TensorflowJS is supporting WebGPU already in several operators.

Additional context
It's worth to note the Google's project Dawn, a C++ native WebGPU implementation will support NodeJS soon. WIP of Node bindings here.

[xenova]

Thanks for the resources :) For the most part, we are waiting for onnxruntime-web to add webgpu as a supported backend.

Here is the associated PR to track its progress:

• [js/web] WebGPU backend via JSEP microsoft/onnxruntime#14579

However, we do plan to support other model formats/backends (in a similar way to how the python library supports PyTorch, tensorflow and ONNX). I don't want to spoil anything... but things are in the work 😉

[gabrielgrant]

AFAIU ORT has merged WebGPU support: microsoft/onnxruntime#11695

What's needed to take advantage of this on the transformers.js side?

[sroussey]

For reference, the webgpu operators implemented:

https://github.com/microsoft/onnxruntime/blob/main/js/web/docs/webgpu-operators.md

[gabrielgrant]

Unfortunately the WebGPU implementation is currently slower than the WASM version, though: microsoft/onnxruntime#18754 (comment)

Would be great to know what's needed to support WebGPU in transformers.js assuming that perf issue gets resolved at some point, but not super urgent/important at the moment

[DavidGOrtega]

Unfortunately the WebGPU implementation is currently slower than the WASM version,

I have some models running in jsep webGPU and are 10 times faster than wasm. I.E. clip

To me, the main problem is the current backend design: It's global (as far as I know). We should be able to setup the preferred backend to our model.

[gabrielgrant]

@DavidGOrtega that's great news! to be clear, are you running your models directly on ORT? or using JSEP through transformers.js somehow? would love to hear more details about exactly what your setup looks like, and which other models you've found this perf improvement on!

[DavidGOrtega]

Im running them with vanilla onnx.

I can do a PR to support WebGPU here (I did node), its trivial. However I think that we should rethink the backend a bit to be more flexible and be able to choose the backend and options per model. Also the onnx fallback is not perfect i.e. I have models that despite the session can be loaded the infer do not work, thats a step after the onnx fallback...

@xenova can also do webgpu and its testing it among other backends like candle. Probably not done yet just because not all the models supports wgpu?

[luweigen]

Unfortunately the WebGPU implementation is currently slower than the WASM version,

I have some models running in jsep webGPU and are 10 times faster than wasm. I.E. clip

To me, the main problem is the current backend design: It's global (as far as I know). We should be able to setup the preferred backend to our model.

@DavidGOrtega What model can you run?

I tried some BERT and got "cannot resolve operator 'Erf' with opsets: ai.onnx v11" with direct call of https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/esm/ort.min.js and the cache of model weights by transformers.js 2.14.2.

[luweigen]

I also tried the v3 branch of transformers.js and got a syntax error. It seems that the commit 66da130 was overwritten by 8c465a9. A simple fix as follows leads to other errors. It seems still a long way to go?

[xenova]

@luweigen the v3 branch is still a work-in-progress, and will be marked as non-draft when reading for testing 👍