TF32 flags when using AMP

mixed-precision
1

Hi! I’m trying out TF32 and mixed precision training. Can they be used at the same time? Namely, is setting the flags like this ok when I enable mixed precision training, or should I actually set them back to False?

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

Thanks!

2

You can keep TF32 enabled and would use it for convs and matmuls outside of the amp.autocast region.

3

Just to check my understanding, the flags would only affect things outside the amp.autocast context?

4

Yes, since inside the autocast context (assuming it’s enabled) float16 or bfloat16 ops will be used for TensorCore-eligible operations.

1 Like