google-research / dinosaur

Dinosaur: differentiable dynamics for global atmospheric modeling

Authors: Jamie A. Smith, Dmitrii Kochkov, Peter Norgaard, Janni Yuval, Stephan Hoyer

Dinosaur is a spectral dynamical core for global atmospheric modeling written in JAX:

• Dynamics: Dinosaur solves the shallow water equations, and the primitive equations (moist and dry) on sigma coordinates.
• Auto-diff: Dinosaur supports both forward- and backward-mode automatic differentiation in JAX.
• Acceleration: Dinosaur is designed to run efficiently on modern accelerator hardware (GPU/TPU), including parallelization across multiple devices.

For more details, see our paper on Neural General Circulation Models.

Usage instructions

Dinosaur is an experimental research project that we are still working on documenting.

We currently have examples replicating two standard test-cases for dynamical cores:

• Baroclinic instability
• Held-Suarez forcing

We recommend running them using Google Colab with a GPU runtime. You can also install Dinosaur locally: pip install git+https://github.com/google-research/dinosaur

See also

If you like Dinosaur, you might also like SpeedyWeather.jl, which solves similar equations in Julia.

Contributing

See CONTRIBUTING.md for details. We are open to user contributions, but please reach out (either on GitHub or by email) to coordinate before starting significant work.

License

Apache 2.0; see LICENSE for details.

dinosaur's Issues

Running Dinosaur across multiple GPU devices

I've been testing out dinosaur for use in my research group here at the University of Exeter , and I'm attempting to see how dinosaur scales across multiple GPU devices. I've been running a version of your Held-Suarez test case on an AWS instance with 4 T4 GPUs. However, I haven't been able to get the test to make use of more than one of the GPUs - when I monitor usage across the GPUs I see GPU 0 getting 100% usage, and barely anything else on the others:

I see that there are test functions with the sharding specified across multiple devices, but I can't see that in the main body of Dinosaur. I'm new to jax, so perhaps this is my lack of experience showing here. I'm much more used to normal Python and Fortran (I did a lot of development of the Isca modelling framework (https://github.com/ExeClim/Isca). Any help would be gratefully recieved. Thanks for releasing dinosaur into the wild - it's great to be able to use it.