Possible to train a smaller version of NeuralGCM?

google-research / neuralgcm

Hybrid ML + physics model of the Earth's atmosphere

https://neuralgcm.readthedocs.io

Apache License 2.0

703 stars 81 forks source link

Possible to train a smaller version of NeuralGCM? #101

Open ZoudaLevy opened 4 months ago

ZoudaLevy commented 4 months ago

I am currently trying to train NeuralGCM on some custom data, but the computational resource requirements are too high for my setup. The model (NeuralGCM-2.8deg) currently requires 16 TPUs, which is beyond my available resources. I would like to know if there is a way to tweak the size of the model so that it can fit on a single A100 GPU.

shoyer commented 3 months ago

You can definitely train NeuralGCM at smaller scale, it will just take longer. Our 2.8 degree fits on a single TPU/GPU, we only used the multiple TPUs for data parallelism.