Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add float8_e8m0_fnu (E8M0) OCP MX scale format #166

Merged
merged 1 commit into from
Sep 12, 2024

Conversation

balancap
Copy link
Contributor

@balancap balancap commented Aug 9, 2024

Adding the OCP MX scale format E8M0, which has the following properties:

  • Unsigned format;
  • 8 exponent bits;
  • Exponent range from -127 to 127;
  • No zero and infinity;
  • Single NaN value (0xFF);

The E8M0 format is used in all MX block format definitions for scale representation (see https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf and issue #116).

ml_dtypes float8_base C++ class is extended to support floating point formats which are unsigned and with no zero (i.e. additional kIsSigned and kHasZero Traits properties).

Base on these traits, float8_e8m0_fnu has been implemented using the existing functionalities (convert, unary/binary ops, ...). Float8 Python unit tests have been extended to be able to cover unsigned floating point formats.

Points to discuss on the PR:

  • Naming: not sure about float8_e8m0_fnu. Maybe ufloat8_e8m0_fn could be better?
  • Testing: not 100% satisfied with additional branching/checks when extending unit tests for E8M0. Is there a cleaner way?
  • (Separate PR?) Additional optimization, as same convertion + ops (mul/div) could be specialized with a more efficient implementation.

@hawkinsp
Copy link
Collaborator

This looks good to me. There's a small failure in the test suite: please fix?

@hawkinsp hawkinsp self-assigned this Aug 20, 2024
@hawkinsp
Copy link
Collaborator

(Sorry about the slow review.)

@balancap
Copy link
Contributor Author

Thanks for looking at it, I'll check the failing tests in the next few days. In the meantime, do you have an opinion on the naming? Should we stick to the name in the PR?

@balancap
Copy link
Contributor Author

balancap commented Sep 3, 2024

@hawkinsp Apologies it took me a bit longer than expected. The CI issue should be fixed now (was just a trivial skipTest issue with Python 3.12). I have rebased the branch as well on the latest main

@hawkinsp
Copy link
Collaborator

This LGTM, but I'd appreciate @sergey-kozub's review, since he's in the process of adding the (complementary) MX data types in #181

@hawkinsp
Copy link
Collaborator

There's a test failure on mac os (different integer overflow behavior?)

@sergey-kozub
Copy link
Contributor

Overall looks fine to me.

This PR touches many files also affected by #181, so I'd like to upstream that one first. I'll do the rebase to fix the merge failures.

@balancap
Copy link
Contributor Author

@sergey-kozub @hawkinsp Removing 0 fixes the MacOS test run, but I am still investigating a crash on the Windows build (stack overflow).

@hawkinsp
Copy link
Collaborator

Ideally we'd like to make a new ml_dtypes release in the next few days, and it'd be great to get both of the MX dtype PRs in.

@balancap
Copy link
Contributor Author

@hawkinsp Just pushed a fix commit for MacOS and Windows. Hopefully should be the last one if @sergey-kozub is good with the PR as it is.

@hawkinsp
Copy link
Collaborator

I'm trying to merge #181 first, which will probably create a merge conflict and you'll need to rebase on top of that PR. Hopefully it will be in soon.

@hawkinsp
Copy link
Collaborator

Ok, #181 is merged.

Would you:
a) squash your commits, please, and
b) rebase on top of head?

@hawkinsp
Copy link
Collaborator

Oh, it looks like you have a test failure on Python 3.13t. I doubt the "free-threading" part is important, but we are trying to make a Python 3.13 release.

@balancap
Copy link
Contributor Author

I'll rebase & squash commits. On the Python 3.13 failure, we should probably skip this test as float8_e8m0 does not have a zero (or negative zero!). Just surprising it triggers on error only on Python 3.13

@balancap balancap force-pushed the add-e8m0-datatype branch 2 times, most recently from 2953430 to b53f481 Compare September 12, 2024 15:54
@balancap
Copy link
Contributor Author

@hawkinsp I should be hopefully the last round of rebase & fix.

I added a skip on the test testHashZero, as E8M0 does not have a zero representation. I can dig a bit more into why the test is failing on Python 3.13, but I don't think it should block the PR merge.

Note: there is a new failure on Python 3.9 introduced by the previous MX float4/float6 PR. I am looking into it.

Adding the OCP MX scale format `E8M0`, which has the following properties:
* Unsigned format;
* 8 exponent bits;
* Exponent range from -127 to 127;
* No zero and infinity;
* Single NaN value (0xFF);

`ml_dtypes` `float8_base` C++ class is extended to support floating point formats
which are unsigned and with no zero (i.e. additional `kIsSigned` and `kHasZero` Traits properties).

Base on these traits, `float8_e8m0_fnu` has been implemented using the existing functionalities (convert, unary/binary ops, ...).
Float8 Python unit tests have been extended to be able to cover unsigned floating point formats.
@balancap
Copy link
Contributor Author

The trick cls._finfo_cache[dtype] = init.__func__() seems to do it for Python 3.9. CI all green on my fork, finally!

@hawkinsp
Copy link
Collaborator

What went wrong without init.__func__()? @sergey-kozub had that in his PR, but i had to revert it because it hit a type checker error. But I'm not sure what problem it's solving.

@balancap
Copy link
Contributor Author

balancap commented Sep 12, 2024

Python 3.9 does not seem to support the direct call to staticmethod (getting this error without __func__(): https://stackoverflow.com/questions/41921255/staticmethod-object-is-not-callable)

The merge commit on the main branch has this issue: https://github.com/jax-ml/ml_dtypes/actions/runs/10833234151/job/30059511975

@copybara-service copybara-service bot merged commit d581b6f into jax-ml:main Sep 12, 2024
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants