[Feature request] Support bfloat16/float8 inputs in `session.run()`

[justinchuby]

session.run() currently does not support bfloat16 inputs because numpy does not support bfloat16. Would it be possible to support the input via np.uint16? session.run() can accept an uint16 value which is the bit representation of the bfloat16 value. ORT should be able to interpret the value correctly because it knows the expected input type of the graph. Same can be done for the float8* types.

[yuslepukhin]

ORT supports bfloat16 internally. I would think a separate per type approach is not the best way to go. Rather, I would implement a new universal method that does not depend on numpy so much, and also include other types, such as quantized 8 and 4.

[yuslepukhin]

I would use a built-in array module and provide shape and type as separate argument for maximum flexibility.
However, along with that, we should have a capability to return those types as results. And right now, we only return numpy arrays.

I think OrtValue based interfaces can help us. Simpe run() returns numpy.
However, we can add methods to OrtValue to accept other types in python array and convert the result to python arrays as well.

[xadupre]

I suggest using run_with_iobinding and make a helper to use run(sess, None, feeds) where is a dictionary dict[str, torch.Tensor].

[yuslepukhin]

IOBinding has its own specifics as to when the input is actually copied to the device.
Also it is cumbersome for every day use.

[justinchuby]

Thoughts on supporting dlpack protocols? This way we don't need to know if the inputs are torch or numpy or anything else. The data can be shared using the same protocol.

[yuslepukhin]

Thoughts on supporting dlpack protocols? This way we don't need to know if the inputs are torch or numpy or anything else. The data can be shared using the same protocol.

Does it not come with its own set of data type enums that seem to be even more restrictive than numpy?

[justinchuby]

Hmm. I tested it and there are constraints. dlpack supports all numpy types except for strings. It supports bfloat16 additionally, but not float8 types. Although as a ML data exchange protocol I can foresee float8 support being added.

A major advantage is that it is supported by all major frameworks (numpy, torch, jax, tensorflow, mlx etc.) so ORT can be compatible with many different tensors easily.

[justinchuby]

import torch
import jax

t = torch.tensor(42, dtype=torch.bfloat16)
jax.dlpack.from_dlpack(t)

Array(42, dtype=bfloat16)

[yuslepukhin]

We can dlpack support, in fact, I have seen some code in pybind referring to dlpack, but I am not seeing it support what we need. Should we not design something that we can own, maintain and deliver what we need?

[yuslepukhin]

It seems, we can extend numpy types support for our needs. See for example a project that implements bfloat16 type and makes it 1st class citizen for numpy. https://github.com/GreenWaves-Technologies/bfloat16/tree/main

[justinchuby]

There is also https://github.com/jax-ml/ml_dtypes (from Google) that implements all relevant types, which we use in the ONNX IR. Upstreaming seems challenging based on the context this project gave.