[Bug] `NaN`s produced when `env.backends.onnx.wasm.proxy = true`?

[josephrocca]

I think I'm doing something silly here, but setting env.backends.onnx.wasm.proxy = true makes this produce NaNs:

let { pipeline, env } = await import('https://cdn.jsdelivr.net/npm/@xenova/transformers@2.8.0');
env.allowLocalModels = false;
// env.backends.onnx.wasm.proxy = true; // uncomment this line to produce NaNs
let extractor = await pipeline('feature-extraction', 'Xenova/bge-base-en-v1.5');
let result = await extractor("test 123", { pooling: 'mean', normalize: true });
console.log(result.data);

I thought maybe it was a 2.8.0 regression, but I tried all versions back to 2.5.4 and it's didn't work in any of them, so that makes me lean toward assuming that I'm doing something wrong. If that's the case, then an informative error message would be helpful :)

[kungfooman]

Hey @josephrocca, thank you for the issue!

I just happen to work on something that helps here: https://runtimetypeinspector.org/

However, I still need to work on the webpack integration (webpack plugins are so complicated...), so right now I can't make the runtime type inspection process reproducible (I'm using my own rollup build script), so I will just note what I discovered so far.

The first "miscalculation" RTI found was here:

Then it's easy to set a breakpoint and inspect the data, first big clue:

Hey, the BigInt64Array is... empty!

a bit of tinkering later...

Go to:

https://github.com/xenova/transformers.js/blob/4e4148cb5ce7f4a9265f58b4eeb660c64bed0386/src/pipelines.js#L128-L129

Add this code:

        // Run model
        console.log("model_inputs.input_ids.data", model_inputs.input_ids);
        //debugger;
        let outputs = await this.model(model_inputs)
        console.log("model_inputs.input_ids.data", model_inputs.input_ids);

Output:

(this means whatever happens in await this.model(model_inputs) invalidates our typed array)

You can uncomment the //debugger; and eventually I figured the BigInt64Array is invalidated in ONNXRuntime:

I did a bunch of WebAssembly programming some years ago and one reason I can think of is memory growth: whenever the WASM memory is increased, it invalidates the typed arrays pointing into that.

However, I don't have the time right now to dig deeper into this, I have to prepare a bit more for RTI to integrate it into the webpack build script of Transformers.js... I hope this sheds a little bit more light into the issue so far 💡

[josephrocca]

Ah RTI looks like a very cool project! I had a project where I desperately needed something like that in the past. BTW, very tangential, but this issue is exciting to me along these lines.

RE this issue - thanks for looking into it! From your investigation, it sounds like this comes from the ORT Web side of things, so a minimal pure ORT Web example is probably required?

Currently my workaround is just to manually put it in a web worker myself.

[kungfooman]

Ah RTI looks like a very cool project! I had a project where I desperately needed something like that in the past. BTW, very tangential, but this issue is exciting to me along these lines.

Thank you! Don't forget to like and subscribe 🤪 I'm kind of against TS-style type annotations in the first place, but I don't want to go into that here.

RE this issue - thanks for looking into it! From your investigation, it sounds like this comes from the ORT Web side of things, so a minimal pure ORT Web example is probably required?

We blamed ONNXRuntime before, but then it was our own fault - so before we point that out, I would like to read @xenova's take on this issue aswell.

I debugged a bit further into ONNXRuntime and I'm pretty sure it's an ownership issue - when the BigInt64Array crosses the Worker border, it's loosing ownership, and when transformers.js wants to do further calculations with it - the tensor data array is empty.

So if we rewrite this code:

https://github.com/microsoft/onnxruntime/blob/a5537f2f563d4975c7e6121a7eb260bbbfd9455a/js/common/lib/inference-session-impl.ts#L108

const results = await this.handler.run(feeds, fetches, options);

Into:

    console.log("feeds before", feeds);
    const feedsCopy = {};
    for (const key in feeds) {
      const feed = feeds[key];
      feedsCopy[key] = feed.clone();
    }
    const results = await this.handler.run(feedsCopy, fetches, options);
    console.log("feeds after", feeds);
    console.log("feedsCopy after", feedsCopy);

The output is:

So we can just clone the feeds into feedsCopy and "sacrificing" them for the session... probably we can do this clone-sacrifize-trick in transformers.js aswell? Something to figure out for a PR either here or at ONNXRuntime.

Documentation for ONNXRuntime run:

https://github.com/microsoft/onnxruntime/blob/a5537f2f563d4975c7e6121a7eb260bbbfd9455a/js/common/lib/inference-session.ts#L337-L356

I doesn't really warn anyone that it could kill our typed arrays 🤔 😅 So maybe it should be reported...

Currently my workaround is just to manually put it in a web worker myself.

Interesting, so just load all of transformers.js into a worker and then it's not crossing the worker boundary, because it's all in a worker? What is the disadvantage of this, maybe poorer debugging via F12/DevTools?

[xenova]

Thanks so much for your investigation @kungfooman! I'm not sure if it is related, but after the inference session returns, I replace all ONNX tensors with a custom tensor class (which subclasses the ONNX tensor class), to add certain functions to the object. Perhaps this has something to do with the invalidation of the existing tensor?

In the meantime, we could add a simple fix which (1) tests if proxy is set to true, and if it is, we make copies before sending to the inference session. @josephrocca or @kungfooman would you be interested in making a PR?

I doesn't really warn anyone that it could kill our typed arrays 🤔 😅 So maybe it should be reported...

cc @guschmue @fs-eire for this. I would expect that the proxy should handle the cloning, as to avoid issues like this in the users' code (in case this is not specific to transformers.js).