Accelerating Earth System Emulation with Spherical Fourier Neural Operators
DescriptionNVIDIA's Earth-2 initiative aims to build digital twins of Earth and has two central goals: (i) Computational: enable high-resolution weather and climate predictions with principled physics-inspired ML; and (ii) Societal: nimbly serve high-fidelity, high-resolution predictions via an intuitive interactive system. The above-mentioned goals depend on achieving orders-of-magnitude speedup and data compression via a combination of advances in AI and computing technologies, in particular, AI at supercomputing scale. In this presentation we discuss the latest developments in our Fourier Neural Operator-based Earth System emulator. We show how a novel geometric deep learning algorithm called the Spherical Fourier Neural Operator (SFNO) that respects the spherical geometry of Earth achieves stable year-long autoregressive rollouts whilst maintaining remarkable skill for high-resolution medium-range global weather predictions. This advancement has important implications for data-driven subseasonal-to-seasonal prediction, and eventually for climate prediction, although many challenges need to be addressed. We discuss the optimizations, parallelism, I/O and memory challenges that were addressed to enable efficient prediction on multi-GPU, multi-node systems. Finally, we discuss opportunities and challenges towards scaling these developments to km-scale weather and climate simulations. We conclude with machine learning algorithmic and engineering innovations required for breakthroughs that will truly help realize digital twins of Earth.
TimeTuesday, June 2716:00 - 16:30 CEST
Climate, Weather and Earth Sciences