FP8 support on H100

mixed-precision
1

with H100 supporting FP8, is there any plan to support FP8 training in pytorch?
Looks like the only alternative is to use FP8 through nvidia transformer engine. Even transformer engine doesn’t support all nn modules like conv layers for FP8 support.

2

@drisspg has been working on this

3

@marksaroufim is there a timeline for the feature to be available?

4

Hey @navmarri we actually have landed some fp8 primitives already in PyTorch. Specifically we have landed fp8e4m3 and fp8e5m2 datatypes and a private function (with no BC/FC guarantees) for doing scaled matmul on h100 machines.

For a quick example script of how you can use it, check out:

5

@drisspg awesome!
does this work with FP8 precision training using FSDP on H100s?

6

As I know (and also be mentioned in Nvidia/TransformerEngine)
RTX 4000 series (compute capibility 8.9) should also support FP8 computing.
Does pytorch support that too?
or for now it is only for H100 and need further works for sm8.9 cards?

7

Is the mixed precision support planned for FP8?

There was an announcement in JAN in the dev forum:

8

TransformerEngine supports FP8 for a long time now. Alternatively you can also play around with pytorch-labs/float8_experimental.

9

Thanks, my point was more related to pytorch AMP support.