use torch.float8_e8m0fnu in mx_formats

[vkuzo]

Summary:

Switches our MX code to use the new torch.float8_e8m0fnu dtype
directly where appropriate. This will allow for easier numerical
debugging of scales, as we can easily see the numerical values when they
are printed.

Test Plan:

pytest test/prototype/mx_formats/ -s -x

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• Make printing of MXLinear and MXInferenceLinear more descriptive #1883
• -> use torch.float8_e8m0fnu in mx_formats #1882

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Unrelated Failures

As of commit ba470ae with merge base 9ef2f06 ():

NEW FAILURE - The following job has failed:

• Run Float8nocompile Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t 293a1177caac9d74b6b63b7e40c5a76df8cc7333d4e7b160d451ea7b0ddd6291 /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh) (trunk failure)
  torchao/experimental/tests/test_int8_dynamic_activation_intx_weight.py::TestInt8DynamicActivationIntxWeight::test_accuracy_{'layout':
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh) (trunk failure)
  ModuleNotFoundError: No module named 'importlib_metadata'

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

[vkuzo]

need to fix CI

[vkuzo]

need to reverse #1608 and get a fresh nightly for this to work, will take a look

[vkuzo]

Looks like this needs something like the inductor changes from pytorch/pytorch#148461 for tests to pass

[vkuzo]

closing in favor of #1966