Prithvi WxC: Foundation Model for Weather and Climate
With the realization that artificial intelligence (AI) emulators can rival the performance of traditional numerical weather prediction models running on high-performance computing systems, an increasing number of large AI models — with a focus on foundation models — are addressing use cases such as weather and climate forecasting. Prithvi WxC is a 2.3 billion-parameter foundation model developed using 160 variables from the Modern-Era Retrospective Analysis for Research and Applications, Version 2 (MERRA-2). Prithvi WxC employs an encoder-decoder-based architecture, incorporating concepts from various recent transformer models to effectively capture both regional and global dependencies in the input data. By design, the model accommodates large token counts to model weather phenomena in different topologies at fine resolutions. Furthermore, it is trained with a mixed objective that combines the paradigms of masked reconstruction with forecasting. Researchers test the model on a set of challenging downstream tasks, namely: autoregressive rollout forecasting, downscaling, gravity wave flux parameterization, and extreme events estimation. The pretrained model, along with the associated fine-tuning workflows, is publicly available as an open-source contribution via Hugging Face.
Systems at the NASA Advanced Supercomputing (NAS) facility are quite good with hosting advanced package stacks and providing access to A100 GPUs with 80 gigabytes of memory on each node. The NAS support team works 24x7, which helps you to quickly check on system performance and remove any access bottlenecks like bad nodes.
Sujit Roy,
NASA Marshall Space Flight Center
Quick Facts
- Researchers used NASA Advanced Supercomputing (NAS) systems with NVIDIA A100 GPUs arranged in nodes each containing four A100 GPUs and 80 gigabytes of memory. During the initial phase, they employed 16 nodes (64 GPUs) for 24 hours to execute each experiment. For pretraining the model, they ran on 16 nodes for about 15 days.
- The zero-shot evaluation introduces reconstruction as a new benchmark, revealing that the model excels in forecasting at shorter lead times.
- Finetuning Prithvi WxC also demonstrates that large, transformer-based foundation models can effectively learn mesoscale atmospheric evolution. This advance helps to streamline, enhance, and accelerate the development of physical parameterizations in climate models, which in turn improves prediction accuracy on interannual timescales. To the team’s knowledge, this effort is the first to use a foundation for climate model parametrization.
- Next, Prithvi WxC researchers are going to further advance the model on data assimilation tasks.
Researchers
- Sujit Roy, NASA Marshall Space Flight Center Please enable Javascript to see email address.
