1 of 15

Karthik Kashinath, Principal Engineer and Scientist, AI-HPC

Earth-2: Towards kilometer-scale Digital Twins for Weather and Climate

2 of 15

NVIDIA Earth-2

AI emulates global weather at 25km-scale

Generative AI emulates the kilometer-scale

Agenda

3 of 15

Earth-2: An NVIDIA Initiative to build a Climate Information System

To boost climate science & climate tech

SIMULATION

AI

VISUALIZATION

4 of 15

The Triple Bind of a Realistic Climate Information System

5 of 15

Earth-2's mandate: “Achieve 3 miracles” - Jensen Huang 

Remarks at The Berlin Summit for Earth Virtualization Engines, Berlin, July 2023.

Km-scale Simulations

30,000 SYPY

30 MW

Generate hi-fidelity data

energy efficiently

AI Emulation of the Full State Vector

Any Region

Any Time Period

Full State Vector Visualization

From Cloud

To enable low-latency interactivity with exabytes of data

Put climate information on your fingertips

MIRACLE #1

TO

MIRACLE #2

TO

MIRACLE #3

TO

6 of 15

FourCastNet

  • Scope Global
  • Model Type Full-Atmosphere Emulator
  • Architecture Spherical Fourier Neural Operator
  • Resolution 25km, 6-hourly (up to 10km, 1-hourly)
  • Training Data ERA5 Reanalysis
  • Training Time O(1000) GPU-hours
  • Inference Time 3 sec (10-day forecast)
  • Calibration Bred Vector + Multi-model checkpoint *NEW*
  • Speedup vs NWP 1000x
  • Power Savings 3000x
  • Max Stable Rollout Decades
  • Scalability Over 90% up to 4000 GPUs, 140 petaFLOPS

ICML 2023, Bonev et al., Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere

7 of 15

Spherical Fourier Neural Operator (SFNO)

Similar to a classical spectral solver

spectral transforms pointwise operations

https://github.com/NVIDIA/torch-harmonics

8 of 15

O(1000)-member Ensembles for Low-Likelihood High-Impact Extreme Events

  • Motivation: Understand climate drivers of LLHI’s
  • Problem: LLHI events poorly sampled in climate record.
  • Only alternatives are simulations of plausible observations.
  • Huge ensembles can quantify the probability of tail events and help assess the associated risk​

Like the 2021 Pacific Northwest heat wave

Bill Collins

Ankur Mahesh

9 of 15

Jensen's 2nd miracle calls for AI that can synthesize km-scale detail

Km-scale Simulations

30,000 SYPY

30 MW

Generate hi-fidelity data

energy efficiently

AI Emulation of the Full State Vector

Any Region

Any Time Period

Full State Vector Visualization

From Cloud

To enable low-latency interactivity with exabytes of data

Put climate information on your fingertips

MIRACLE #1

TO

MIRACLE #2

TO

MIRACLE #3

TO

10 of 15

Key challenges towards global km-scale emulation

  • Model size and data should be scaled equally – transformers and some others [Chinchilla, 2022].
  • Increased spatial resolution requires finer timesteps, but error accumulates with autoregressive rollout.
  • Training SFNO on km-scale global data from scratch requires at least 12,000 H100s!

11 of 15

The fidelity and flexibility of modern generative AI

Diffusion models are revolutionizing Generative AI in visual domains

SORA, OpenAI

Prompt: Photorealistic closeup video of two pirate ships battling each other as they sail inside a cup of coffee.

Can diffusion models learn high-res details of Earth's weather and climate dynamics too?

12 of 15

CorrDiff: Generative diffusion modeling for regional km-scale downscaling

Tapping into extensive gen AI research and optimizations being developed

Mardani et al. (2023), Generative Residual Diffusion Modeling for Km-scale Atmospheric Downscaling, arXiv: 2309.15214

Residual learning inspired by the Reynolds decomposition in fluid dynamics

13 of 15

14 of 15

Coupled with global ML emulators, CorrDiff offers a scalable pathway to km-scale prediction…

25km global ML emulator

Km-scale regional predictions

GenAI

Lead-time aware, fine-tuned to distribution shifts

  • 1000-member ensemble in 8 minutes. Massive ensembles for free (sampling from distribution)
  • 200x data compression, 1000x faster, and 3000x more energy efficient than a WRF simulation at 2km
  • Could be adapted for climate prediction by conditioning on hi-fidelity regional climate data, e.g., CORDEX

15 of 15

Km-scale training datasets are key!