add the torch.float8_e8m0fnu dtype to PyTorch

[vkuzo]

Summary:

Continuing the work from #146427

Adds the torch.float8_e8m0fnu dtype to PyTorch, as detailed in
#146414 . Please see the issue for a detailed definition of the format. Example of basic functionality:

import torch

# round trip
x0 = torch.randn(4, 4, dtype=torch.float32)
x1 = x0.to(torch.float8_e8m0fnu)  # RNE rounding
x2 = x1.to(torch.float32)  # 2 ** exponent

# creation with empty
x0 = torch.empty(4, 4, dtype=torch.float8_e8m0fnu)

# printing
print(x0)

Done in this PR:

• numerical correctness
• op coverage (except for torch._scaled_mm): create tensor, cast to/from float32
• printing a tensor works

For future PRs:

• performance optimizations for casting
• torch._scaled_mm
• PT2
• various cleanups (detailed in comments with issue numbers)

Test Plan:

pytest test/quantization/core/experimental/test_float8.py -s

Reviewers:

Subscribers:

Tasks:

Tags:

cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/147466

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d635775 with merge base 303ad19 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vkuzo]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 jobs have failed, first few of them are: linux-binary-manywheel / manywheel-py3_9-cuda11_8-build / build

Details for Dev Infra team

Raised by workflow job

[vkuzo]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[henrylhtsang]

false alarm my bad

[yiakwy-xpu-ml-framework-team]

Hi @vkuzo are you still working on the problem ? I think the way you used MXFP8_E8M0_FNU could be discussed.

The number is used inside group_quantize function each warp extract 32 fp3_exponent from 32 fp32. So the ocp_fp8e8m0fnu_from_fp32 is simply equal to extract fp32 exponent.

No mantssa and speical number should be taken careful because exponent itself is unsigned. We don't need to care about it.

Fp32 -> fp8 + mxfp8_e8m0_fnu

The second problem is that this mxfp8_e8m0_fnu will be shared by a group consecutive 32 (group) elements (no exponent, only mantissa w/wo implict-1). (because mantissa multiply the exponent is the fp32).

Note fp8 scalar multiple a fp8 data type is as easy as fp32 = fp8_scale << fp32_t::M | (fp8 & fp32_t::INT32_M_MASK));

Here is my implementation :

template<>
HOST_DEVICE_INLINE OutType ocp_fp8e8m0fnu_from_fp32(float fval) {
    using fp32_t = Float;
    using fp8_t = Float8_E8M0_FNU;
    using fp8_storage_t = Float8_E8M0_FNU::Datum;

    union {
        float fval;
        int32_t i32val;
        uint32_t ui32val;
    } val;

    val.fval = fval;

    fp8_storage_t ui8val = (val.i32val & fp32_t::INT32_E_MASK) >> fp32_t::M;
    return fp8_t::from_bits(ui8val.ui8val);
}

[vkuzo]

hi @yiakwy-xpu-ml-framework-team , sorry for late reply, I was on a long leave and am now catching up on what I missed.

The number is used inside group_quantize function each warp extract 32 fp3_exponent from 32 fp32. So the ocp_fp8e8m0fnu_from_fp32 is simply equal to extract fp32 exponent.

Correct! Please check out #146414, "E8M0 detailed proposal" for details. The default cast to e8m0 in PyTorch is using RNE to match the IEEE-754 spec and other floating point dtypes, which does not match what is described in the OCP spec. It's up to the user to specify a different casting/rounding behavior (such as floor) if they would like to do so - this is 100% valid.