Tensor Core enabled for mixed precision on GPUs that allow it

ORNL / HydraGNN

Distributed PyTorch implementation of multi-headed graph convolutional neural networks

BSD 3-Clause "New" or "Revised" License

68 stars 29 forks source link

Open allaffa opened 11 months ago

jychoi-hpc commented 11 months ago

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 commented 11 months ago

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

allaffa commented 11 months ago

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