Explictly specify fp32 precision to avoid defaulting to 'highest'

[O-J1]

We just need to decide on 'high' or 'medium'. Habishi reccomended medium. Highest is exceedingly slow.

https://docs.pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html

[dxqb]

summary as discussed on Discord:

• 'medium' is bfloat16 precision, which the user can do by setting this train dtype. If he did set float32 as train dtype, float32 should be used, not a lower precision

• 'high' seems interesting and might be a precision we don't support yet:

“high”, float32 matrix multiplications either use the TensorFloat32 datatype (10 mantissa bits explicitly stored) or treat each float32 number as the sum of two bfloat16 numbers (approximately 16 mantissa bits with 14 bits explicitly stored), if the appropriate fast matrix multiplication algorithms are available. Otherwise float32 matrix multiplications are computed as if the precision is “highest”. See below for more information on the bfloat16 approach.

I'm not sure if this is better than tfloat32, which we do support. If we want to support this 2xbfloat16 algorithm, it should be a separate setting the user chooses.