`torch.set_float32_matmul_precision` deprecation

[cw-tan]

Outline & Motivation

On PyTorch 2.9 RC, torch.get_float32_matmul_precision() will cause the warning

UserWarning: Please use the new API settings to control TF32 behavior, such as torch.backends.cudnn.conv.fp32_precision = 'tf32' or torch.backends.cuda.matmul.fp32_precision = 'ieee'. Old settings, e.g, torch.backends.cuda.matmul.allow_tf32 = True, torch.backends.cudnn.allow_tf32 = True, allowTF32CuDNN() and allowTF32CuBLAS() will be deprecated after Pytorch 2.9. Please see https://pytorch.org/docs/main/notes/cuda.html#tensorfloat-32-tf32-on-ampere-and-later-devices (Triggered internally at /pytorch/aten/src/ATen/Context.cpp:80.)

For my use case, it seems to be coming from

pytorch-lightning/src/lightning/fabric/accelerators/cuda.py

Line 165 in dfbee80

if torch.get_float32_matmul_precision() == "highest": # default

and

pytorch-lightning/src/lightning/fabric/utilities/throughput.py

Line 610 in dfbee80

if _is_ampere_or_later() and torch.get_float32_matmul_precision() != "highest":

naively testing torch.set_float32_matmul_precision("high") in my REPL didn't seem to throw any errors.

I am unsure about the true extent of the changes necessary to account for this fp32 precision API change imposed since PyTorch 2.9 (some notes: https://docs.pytorch.org/docs/2.9/notes/cuda.html) beyond the examples presented above.

Pitch

No response

Additional context

No response

cc @lantiga @justusschock

[SkafteNicki]

@cw-tan thanks for raising this issue. The warning seems weird to me because if it indeed is torch.get_float32_matmul_precision() that triggers the warning, then that has nothing to do with what the function does. The function does not change the precision only fetches the current set one. I wonder if this is a bug on Pytorch side that should be fixed.

[cw-tan]

potentially, but it looks like it'll probably make its way into the final PyTorch 2.9 release this month. I'm guessing the query calls
https://github.com/pytorch/pytorch/blob/b5c4f46bb9ede8dc6adf11975c93b9f285d9ed67/aten/src/ATen/Context.cpp#L375
which throws the warning

[SkafteNicki]

Alright, I guess it makes sense since the get_float32_matmul_precision function is part of the old API and not the new one. Either there is a new way to get the precision in the new API or we would need for the time being filter the warning.