Add low-precision to TFNO

[crwhite14]

This pull request adds low precision options for TFNO. Here is a brief summary.

1. opt.amp_autocast was previously in the yaml file but was a no-op. This PR gets opt.amp_autocast working

• False (default): run the model in full precision
• True: turn on torch.amp.autocast, torch's built-in half-precision method. This turns many operations to half precision, with a few notable exceptions: reduction operations, weight updates (due to instability) and complex-valued operations (due to lack of implementation).

2. Add new parameter to the yaml file, fno_block_precision

• 'full' (default): standard full precision
• 'half': the FFT, contraction, and inverse FFT run in half precision
• 'mixed': the FFT runs in full-precision, and the contraction and inverse FFT run in half precision

3. Add new parameter to the yaml file, stabilizer

• None (default)
• 'tanh': adds a tanh just before the FFT in the fno-block. Typically this needs to be set to stabilize the FFT, when fno_block_precision='half'.

Running

python train_navier_stokes.py --opt.amp_autocast=True --tfno2d.fno_block_precision='half' --tfno2d.stabilizer='tanh'
python train_navier_darcy.py --opt.amp_autocast=True --tfno2d.fno_block_precision='half' --tfno2d.stabilizer='tanh'

Improves runtime and memory by up to 30%, depending on the GPU used, the resolution of the data (greater speedups for 64x64 resolution or higher), and other hyperparameters such as factorization and rank.

[JeanKossaifi]

This looks great, thanks!

[JeanKossaifi]

Should we remove this one or do you think we need to always explicitly cast here @crwhite14 @rtu715 ?

[crwhite14]

We can remove the "else: x = x.float()"
I just did that in this commit: bcc52e7#diff-5ae1e49af12ed16c75135c0043a08575110fd03d4c722a837b60aa0950b31e32L342-L343

[JeanKossaifi]

Thanks, great PR @crwhite14 @rtu715, merging!