Fix bf16 Support

[manbearian]

Fix bf16 Support
Fix bfloat16 support in helper and numpy_helper.py.

Motivation and Context
Without native bfloat16 support in numpy implementing bfloat16 support for ONNX is a bit complicated. Recently support was added by encoding bfloat16 as float16. this is not correct as this is a lossy encoding; it's also a bit confusing as a design and indeed broke some bespoke tools on my end that didn't expect the bfloat16 to be encoded as float16 in the ONNX file.

Open Questions

I've implemented some changes here that give rudimentary support (not completely broken) but i do have a two questions that i believe need to be answered before committing these changes.

1. what is the proper binary encoding of bfloat16 in ONNX protobuf format (is this documented? should it be?)
2. it appears that "raw" encoding and normal encoding for 16-bit values are of different sizes, is this intended?

Not implemented Changes

1. support for big-endian encodings is not implemented in my code
2. support for "correct" truncation of NaN payloads is not implemented when converting from f32->bf16 in my code

Addresses #4189

[manbearian]

How can we move forward on this? I have some time this week to continue to hack on this if there's more folks would like to see.

[gramalingam]

It would be great if we could make this work for both big-endian and small-endian. Is that hard? Is there some reference for the transformation you are using? I am trying to understand why trunction of 32-bit to 16-bit is incorrect.

[manbearian]

It would be great if we could make this work for both big-endian and small-endian. Is that hard? Is there some reference for the transformation you are using? I am trying to understand why trunction of 32-bit to 16-bit is incorrect.

Endianness makes by brain hurt, but i'll see if i can work through the pain and implement it. :) (i was kind of hoping someone would say it isn't needed ;)

See my other reply as to why i chose rounding over truncation for converting from fp32 to bf16 and a pointer to the pytorch bfloat16 code which is a correct implementation that i drew inspiration from.

[manbearian]

NaN is now implemented correctly (my previous implementation could potentially convert NaN to infinity)
Big-Endian is now 'supported'... at least i believe i have the right implementation, but i have no way to test.

[manbearian]

@gramalingam and @souptc thoughts?

[souptc]

i see. i think it is correct. I actually wondering how does our cpu/gpu EP handle it.

[lgtm-com]

This pull request introduces 2 alerts when merging 9e0d884 into 28ed7f5 - view on LGTM.com

new alerts:

• 1 for Use of the return value of a procedure
• 1 for Unused local variable

[gramalingam]

Nit: if we decide that rounding is the right behavior instead of truncating, it would be good to fix the "Cast" implementation in the test-case generator also.

[manbearian]

Nit: if we decide that rounding is the right behavior instead of truncating, it would be good to fix the "Cast" implementation in the test-case generator also.

@gramalingam
i'm looking at this code now, and suspect i see how it can be changed, but i'm not sure how it gets used. Are there any pointers to how it works?

[manbearian]

i see. i think it is correct. I actually wondering how does our cpu/gpu EP handle it.

do you have a pointer to where the code might be, i'm happy to take a look.

[jcwchen]

Nit: if we decide that rounding is the right behavior instead of truncating, it would be good to fix the "Cast" implementation in the test-case generator also.

@gramalingam i'm looking at this code now, and suspect i see how it can be changed, but i'm not sure how it gets used. Are there any pointers to how it works?

https://github.com/onnx/onnx/blob/main/onnx/backend/test/case/node/cast.py This file is used to create the node test model (For instance, https://github.com/onnx/onnx/tree/main/onnx/backend/test/data/node/test_cast_FLOAT_to_BFLOAT16). You can update the code with correct behavior in cast.py first and then use tools/update_doc.sh to regenerate related node test models.

[jcwchen]

i see. i think it is correct. I actually wondering how does our cpu/gpu EP handle it.

do you have a pointer to where the code might be, i'm happy to take a look.

Probably here. However, different EPs seem to have different behaviors. We are trying to reach out the code owner to understand the inconsistency and which behavior is correct.

[manbearian]

@jwchen and @gramalingam i updated cast.py/castlike.py to use rounding for f32->bf16. Let me know if this is what you were looking for.

I was expecting to see some tests fail, but i believe the tests aren't checking for exact values, so the single bit difference isn't detectable.

[gramalingam]

I suspect that the checkin for test-data for ops other than cast / cast-like are spurious, I don't think they should change?

[jcwchen]

Thank you @manbearian for the quick update! As @gramalingam mentioned, please remove other irrelevant operators' update (like updating output.pb due to different numpy.random behaviors by different machines). tools/update_doc.sh will update every node test data and we should only take Cast/CastLike related tests.

I was expecting to see some tests fail, but i believe the tests aren't checking for exact values, so the single bit difference isn't detectable.

You are right -- these models/input.pb/output.pb just exactly follow what onnx/backend/test/case/node/[operator_name].py defines. The CI only checks that whether the uploaded model can be reproduced by CI environments.

[manbearian]

after talking with @gramalingam offline we decided to drop the CAST/CASTLIKE changes for now. I have them stored in a branch if we want to bring them back. There was also a suggestion that the f32->bf16 helper support both rounding and truncation modes. I can add support for that with a future PR if there is interest (please open an issue and assign it to me).

[gramalingam]

Thanks Ian! The rationale for our decision was the question about what Cast should do (rounding or truncation) is separate (it involves a tradeoff between efficiency and precision) and will require more time to get consensus. However, the helper function make_tensor is not as critical (users can round or truncate as they wish), but it is important to fix the error in make_tensor's current handling for this release.