Disable tensorfloat32 type

[scarlehoff]

Apparently the "new math" is just being sloppy as hardware instructions: https://blogs.nvidia.com/blog/tensorfloat-32-precision-format/

Note that the above is for GPU (or TPUs) but newer CPUs might be affected by the same issue in case we see it again:
https://www.intel.com/content/dam/develop/external/us/en/documents/lower-numerical-precision-deep-learning-jan2018-754765.pdf

For context here are the results, 251105-jcm-nnpdf41-mhou-001 is the new baseline and these are a few comparisons with a computer with this problem:

a) Default options, tensorfloat 32 is used for some operations with the consequence of a loss of precision

b) double_precision: true in n3fit. By setting the default type to float64, only some (not clear to me which) operations are downgraded to these tensorfloat types, but everything looks ok

c) And finally comparison to a fit in which the tensorfloat is disabled and the entire fit is done in single precision, results are ok

I'll benchmark cases b) and c) to check what is faster in cineca before merging this branch. Or if you prefer we can merge it asap and if during the benchmark we find that option b) is better (run by default in dp but keep these special features on) we can change the default.

NB: I'm only disabling it in TensorFlow, although in principle this is a hardware thing. However, apparently pytorch is more conservative in which operations have it on by default (perhaps that's why it is much slower...) and the results are actually ok. I haven't tested jax.

[Radonirinaunimi]

Tested and it works fine for me as well! Everything looks good.

[scarlehoff]

Sadly it does make the runs about 40% slower in the cluster :(

I'll see whether I can find which operations need to be run with extra precision, but I'm afraid it might be the optimisation step itself

[tgiani]

@scarlehoff do you mean that b) makes runs 40% slower?

[scarlehoff]

Running in Leonardo, 110 replicas, without MHOU:

a) 4180s
b) 5695s
c) 5364s

So the best option is strict float32, which seems to be as precise as float64 for us. Sadly the nice NVIDIA magick by which they get to announce 8x speed ups for machine learning in their presentations is, indeed, much faster, in exchange for the numbers being wrong 😅

But you get them wrong much faster!!!

[scarlehoff]

I'm running a few more tests in case I can isolate some operations for which we can lower the precision, but I think I will merge this as is.