torch.backends.cuda.matmul.allow_tf32 = True
When this is done CUDA will automatically switch to using tf32 instead of fp32 where it's possible. This, of course, assumes that the used GPU is from the Ampere series.
Like all cases with reduced precision this may or may not be satisfactory for your needs, so you have to experiment and see. According to NVIDIA research the majority of machine learning training shouldn't be impacted and showed the same perplexity and convergence as the fp32 training.
If you're already using fp16 or bf16 mixed precision it may help with the throughput as well.
You can enable this mode in the 🤗 Trainer with --tf32, or disable it with --tf32 0 or --no_tf32.
By default the PyTorch default is used.
Note: tf32 mode is internal to CUDA and can't be accessed directly via tensor.to(dtype=torch.tf32) as torch.tf32 doesn't exit.
