FP16 and AMP Support

[toenshoff]

Hey PyG Team,

do PyG-Lib ops currently support 16-bit floating point precision in any way?

I could not find any information about this in the documentation but segment_matmul throws a runtime error when given 16-bit tensors, saying it expects 32-bit.

Are there any near future plans for adding support for 16-bit tensors or automatic mixed precision?

[rusty1s]

I think most of PyG supports 16-bit precision right now (segment_matmul might actually be the only exception).

[toenshoff]

Thank you for the quick answer!

I do get the same error when running grouped_matmul.

The use case I am currently working on would benefit greatly from an efficient version of the HeteroLinear module that also supports 16 bit or AMP. As far as i can tell this has to be implemented with either segment_matmul or grouped_matmul for proper GPU utilization, but neither seems to run with 16 bit floating point.

[rusty1s]

These two functions reference the same underlying CUDA routine which currently only supports float precision. Pinging @puririshi98 for additional thoughts.

[toenshoff]

Thank you. After reviewing the source code more closely it turns out that grouped_matmul does support FP16, but only when torch 2.0.0 (1.14.0) is installed. The method then uses torch.nested since its recent update and this makes no restrictions to dtypes.

We will implement a version of HeteroLinear around grouped_matmul and use torch 2.0.0 for now.