is this intended to be used with scale? The current autocast design is stateless, so scaled float8 doesn't cleanly fit in. If this is intended to be used without scale, I'd love to learn on how one might achieve competitive accuracy.

You're right, this shouldn't be there. We were exploring possibilities to use fp8 with autocast, but finally decided to use custom solutions.
In fact, I see that AT_FORAUTOCAST_SCALAR_TYPES is not use anywhere, the only usage were removed quite long time ago. I will clean this during further work.

just curious, are you folks using these unscaled arithmetic operations on float8 for real use cases, or is this just for completeness?

This is just for completeness and simple tests. We use scaled float8 operations.

nice! I didn't read line by line yet, but high level question - what's your level of confidence on the correctness of these functions? Have these been matched up with NVIDIA's casts or https://github.com/pytorch/FBGEMM/blob/277677039bae25b2570a73013b03bfaa9d2a523e/fbgemm_gpu/include/fbgemm_gpu/quantize_ops_utils.h ?

It was quite extensively tested against Gaudi2, which supports these variants natively. However, in next commits, I'll add comparative tests for all bits configurations against some reference implementation.
Can you guarantee that fbgemm implementation is fully aligned with the specification? Also, do you know if there is some open fp8_e5m2 reference implementation I could use for testing?

the fbgemm implementation has been used extensively at Meta but it hasn't been matched up against real hardware afaik.

Your test cases look good, I'd say if future changes on numerics are needed that can be done in future PRs and not block this one. LGTM.

We probably will need to replace this with a primitive at some point as this is highly unlikely to be optimal.

Changed to helper method

inline C10_HOST_DEVICE bool Float8_e4m3fn::isnan() const {
  return (x & 0b01111111) == 0b01111111;
}

b8/h8 are funny names to call these lol.

Are there funny enough to be changed, or can I leave there?

This says that when you combine Float8_e5m2 and Float8_e4m3, it promotes to Float8_e4m3. This seems suspicious to me, as it's inconsistent with float16 and bfloat16 promotion, which promotes to float32 per the RFC at #43049 Additionally, JAX currently doesn't support promotion between these two types. It would be safer to disallow promotion to start and let folks figure out what they want to do later.

Done.

@vkuzo, folks pointed out to me on Twitter https://twitter.com/ezyang/status/1665399428809629696 that there are multiple e4m3 variants; GraphCore has its own separate NaN variant. In JAX, these variants are distinguished, see https://github.com/jax-ml/ml_dtypes#specifications-of-implemented-floating-point-formats for their naming convention. Since I don't think we're expecting people to actually use these dtype names directly, it seems... OK to keep inline with the LLVM/JAX naming conventions?

I agree, I think this should be named Float8_e4m3fn to match other frameworks and since there is no single float8 standard yet, I commented on that yesterday (on the PR itself, not on the code).

Done.

I didn't carefully audit these

Looks good.
Checked with

import jax.numpy as np
vals = np.arange(0, 2**8, dtype=np.uint8).reshape(-1, 4)

print("e4m3 limits")
a = np.array([0x08, 0xFE, 0x7E, 0x20, 0x30, 0x7F, 0x01], dtype=np.uint8)
print(a.view(np.float8_e4m3fn))
print(vals.view(np.float8_e4m3fn))


print("e5m2 limits")
a = np.array([0x4, 0x7B, 0xFB, 0x34, 0x38, 0x7C, 0x01], dtype=np.uint8)
print(a.view(np.float8_e5m2))

print(vals.view(np.float8_e5m2))

I just want to point out that this is completely unnecessary; as uint8 is byte sized you can just read it directly out of the pointer, no need to memcpy lol

Removed.

don't forget to fix the missing eol

I did only a cursory inspection of this file

Looks good

Code motion only right?

Yes.

return false; haha

Nice catch, done :)

Looks good

Looks good.
Checked with

import jax.numpy as np
vals = np.arange(0, 2**8, dtype=np.uint8).reshape(-1, 4)

print("e4m3 limits")
a = np.array([0x08, 0xFE, 0x7E, 0x20, 0x30, 0x7F, 0x01], dtype=np.uint8)
print(a.view(np.float8_e4m3fn))
print(vals.view(np.float8_e4m3fn))


print("e5m2 limits")
a = np.array([0x4, 0x7B, 0xFB, 0x34, 0x38, 0x7C, 0x01], dtype=np.uint8)
print(a.view(np.float8_e5m2))

print(vals.view(np.float8_e5m2))

You can use parametrize from torch.testing._internal.common_utils

Done.

We don't directly use pytest in tests but unittest.
You can take https://github.com/pytorch/pytorch/blob/a2049dbf8f921c925b70bd3714c1f5dbcd8200ed/test/quantization/core/experimental/test_linear.py as an example for the class creation and the __main__ clause.
Also all plain asserts can be replaced with sefl.assert*

Done.

Not sure what this is testing? Aren't 0s supposed to be the same?
Also once you use our test framework, you can do self.assertEqual(x, x8) !

Just first sanity test. I can remove it if it's too basic.

Anyone remembers, what was the original idea behind ALL_TYPES_AND_COMPLEX_AND?
But perhaps there is a time to add ABSOLUTELY_ALL_TYPES and ALL_TYPES_AND_ALL_FLOATS? (as I can't imagine want wants to add a support for F8_E5M2 but not for F8_E4M3)

to clarify, this is not blocking the PR - we can figure out how to clean up these abstractions independently

Can I ask if PyTorch is only implementing the non-saturation mode for casts, as specified here ?

The current default behavior for ONNX is to saturate: https://github.com/onnx/onnx/blob/main/docs/Operators.md#attributes-7.

You're right, currently it implements only non-saturation mode. And, in consequence, is not aligned with float8 implementation from FBGEMM https://github.com/pytorch/FBGEMM/blob/main/fbgemm_gpu/include/fbgemm_gpu/quantize_ops_utils.h

@vkuzo, @albanD, was it intended to have saturation enabled by default in casting to float8 dtypes? If so, it needs to be added, because currently values out of float8 range are being casted to inf/nan. I can do it next week after vacation.

Also, should it be configurable via parameter of .to() operator or some other way?

Thanks for the comment! During PR review we punted these questions to future PRs just to get the initial dtype in, I guess now it's the time to discuss them :)

I would make an educated guess that saturating by default would make sense if people are doing delayed scaling, but let me think on that / talk to some of the research teams using fp8.

For if/how to make this configurable...let me also ask around on what would make sense.

In terms of consistency, when given integer inputs, we raise an error for too large numbers.
For floating point inputs, all our dtypes will return inf.

So I think inf is the right call here.

I think that a_half = a.clamp(max=torch.finfo(torch.half).max).to(dtype=torch.half) is a good way to get the behavior you want if needed.

@albanD , there is some nuance here worth thinking about. In https://arxiv.org/pdf/2209.05433.pdf which these dtypes are expected to match, in Section 2, the authors say "Details of the heuristics to select
the scaling factors are beyond the scope of this paper, but the general idea is to choose a scaling factor such that the
maximum magnitude in the tensor becomes close to the maximum representable magnitude in the corresponding format. Values that overflow are then saturated to the maximum representable value." Specifically, if someone is using delayed scaling and there is an outlier, saturation is probably practically what they want.

NVIDIA also provides an enum for saturation behavior for fp8 (https://docs.nvidia.com/cuda/cuda-math-api/group__CUDA__MATH__FP8__MISC.html), which isn't there for wider floating point dtypes.

IMO a reasonable default behavior here should be what is practical for fp8 training with delayed scaling, and we'll have to work out the inconsistencies of how this is different with what we have for wider fp dtypes.

What do you mean by private API? Will it be exposed to the user on the python level? Or only on the level of hardware provider's PyTorch plugin?

It will be exposed to Python with a name preceded by an underscore, but there will initially be no BC guarantees.

Thank you for the information.

Hi @vkuzo @albanD @australopitek ,

May I know the current status of enabling fp8 delayed scaling and saturation in PyTorch?

For fp8 delayed scaling, it is a per-tensor one. So, is the implementation similar with Quantized operators? It cannot just use CPU dispatch, right? Because scaling factor need to be stored in Tensor. Please correct me if anything wrong.

And you mentioned that fp8 training with delayed scaling, what about fp8 inference? In https://github.com/NVIDIA/TransformerEngine, fp8 inference is also with delayed scaling as training. Will you take this into consideration of fp8 enabling?

Thanks so much! Hope to hear from you guys!

cc @mingfeima

Hi @yanbing-j , please feel free to message @gchanan on PyTorch slack for context on your question. I am on parental leave until early December but Greg can help on this while I'm out.