Option for disable tf32

[dnnspark]

Hi, is there a way in trainer to disable tf32 for ampere architecture? It's motivated by this discussion:https://discuss.pytorch.org/t/numerical-error-on-a100-gpus/148032/2

cc @justusschock @kaushikb11 @awaelchli @Borda @rohitgr7

[akihironitta]

Hi @dnnspark! Simply setting the flags in your script doesn't work?

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

[dnnspark]

Hi @akihironitta, Yup. that works. I wanted to check if there's a way (e.g. flag of Trainer) to do this.

[krshrimali]

Hi, @dnnspark - Thanks for creating the issue and the question. From what I know, there is no such flag to do this currently with the Trainer class. Since this is specific to use-cases, I would prefer it to stay as a flag (like what @akihironitta mentioned). However, we are definitely open to feedback and suggestions (if you think that it should instead be an arg in the Trainer class).

Hope you're enjoying using PyTorch Lightning ⚡ ! :)

[justusschock]

Just leaving my two cents here: When manually setting these, it might not work with ddp_spawn, since we don't rerun the whole script question here: If `torch is reimported in the new process, are those changed attributes retained? So In general, we need a mechanism to copy all the flags set to the main process to the spawned processes as well. Not saying, this should be a trainer flag though, this should happen under the hood.

Not sure, how to properly do it either. For the beginning a hardcoded list of properties may be sufficient.

[justusschock]

thoughts @awaelchli ?

[justusschock]

@Borda just fyi: I think this should have been an issue. A lot easier to track there especially if it shows a possible limitation :)