How can I enable TensorFloat-32 when using A100?

[austinmw]

Thanks

[HAOCHENYE]

Hi~ sorry for my late reply. We were going to solve this issue by PR (Setting any environment variable by configuration file). But this solution seems a bit dangered, do you have any opinion about this?

[austinmw]

Hi, I'm not sure about the best way to do it. A (maybe naive) way may just be like this?

if cfg.cuda_matmul_tf32:
    torch.backends.cuda.matmul.allow_tf32 = True
else:
    torch.backends.cuda.matmul.allow_tf32 = False

if cfg.cudnn_tf32:
    torch.backends.cudnn.allow_tf32 = True
else:
    torch.backends.cudnn.allow_tf32 = False

The feature would be greatly appreciated in any form though, since the speedups are very significant

[austinmw]

Also I think these default to True for PyTorch 1.7-1.11 and default to False in PyTorch >= 1.12

[HAOCHENYE]

Initially, this PR wanted to support configuring PyTorch environment variable, the example configuration is like this:

cfg.env_cfg.torch_env = {
'backends.cudnn.allow_tf32'=True
}

However, others think that the environment variable of other 3rd party libraries, such as OpenCV, should also be configured. Then the PR turns out like that 🤣. We'll reconsider your proposal and support this feature ASAP!

[austinmw]

Hehe, thanks so much!