Tensor Core enabled for mixed precision on GPUs that allow it

[allaffa]

No description provided.

[jychoi-hpc]

I ran a few test runs using qm9 on Frontier. I found the following related to this AMP feature (Automatic Mixed Precision):

• (Almost) no speed up in time: it may differ by cases but at least in the qm9 case the wallclock time reduction was very very small.
• But, memory saving was obvious: 114 MB with AMP vs 180 MB without AMP
• Numeric artifacts can occur with AMP. I.e., we don't get same numbers with AMP and without AMP. Here are the loss plots:
  

[jychoi-hpc]

We may need to make this feature optional as it changes the accuracy or losses.

[allaffa]

We may need to make this feature optional as it changes the accuracy or losses.

@jychoi-hpc I agree. The differences between full precision and mixed-precision seem to be drastic, even for a relatively simple dataset like QM9