Correct behavior when using precision=16 with Tensor-Cores

[SerezD]

I have notices that setting precision=16 in the Trainer raises a warning on the cluster I am using:

You are using a CUDA device ('NVIDIA A100-SXM4-40GB') that has Tensor Cores.  
To properly utilize them, you should set 
`torch.set_float32_matmul_precision('medium' | 'high')`  
which will trade-off precision for performance. 
For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision

Does this means that I should avoid precision=16 ? Where should i set the suggested command:
torch.set_float32_matmul_precision('medium' | 'high') in the PL script ?

Thanks

[awaelchli]

@SerezD Good question.
This setting is independent of the "precision" setting in the Trainer. It is regarding how matrix multiplications are performed in general. Note that precision=16 conceptually means "mixed precision", so not everything is in float-16 precision. So you can potentially benefit from this setting regardless.

Where should i set the suggested command

This is a global setting, so you can put it anywhere you want. But of course, preferably before training starts :)

Just a heads up: For future questions, consider posting them in our new Forum over at lightning.ai/forums. We are slowly beginning the migration away from GH discussions.