[Paper Review] FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators

Thorsten Kurth, Shashank Subramanian+7 more|arXiv (Cornell University)|Aug 8, 2022

Meteorological Phenomena and Simulations24 citations

TL;DR

FourCastNet uses Adaptive Fourier Neural Operators to emulate Earth system dynamics and forecast global weather at high resolution, achieving about 80,000x faster inference and strong accuracy compared with traditional NWP. Training scales to 3,808 GPUs with 140.8 petaFLOPS in mixed precision.

ABSTRACT

Extreme weather amplified by climate change is causing increasingly devastating impacts across the globe. The current use of physics-based numerical weather prediction (NWP) limits accuracy due to high computational cost and strict time-to-solution limits. We report that a data-driven deep learning Earth system emulator, FourCastNet, can predict global weather and generate medium-range forecasts five orders-of-magnitude faster than NWP while approaching state-of-the-art accuracy. FourCast-Net is optimized and scales efficiently on three supercomputing systems: Selene, Perlmutter, and JUWELS Booster up to 3,808 NVIDIA A100 GPUs, attaining 140.8 petaFLOPS in mixed precision (11.9%of peak at that scale). The time-to-solution for training FourCastNet measured on JUWELS Booster on 3,072GPUs is 67.4minutes, resulting in an 80,000times faster time-to-solution relative to state-of-the-art NWP, in inference. FourCastNet produces accurate instantaneous weather predictions for a week in advance, enables enormous ensembles that better capture weather extremes, and supports higher global forecast resolutions.

Motivation & Objective

• Motivate the need for faster, high-resolution weather forecasting amid increasing extreme weather due to climate change.
• Develop a data-driven Earth system emulator to surpass the throughput of traditional numerical weather prediction (NWP) models.
• Demonstrate scalable training and inference performance on multiple supercomputers using GPU acceleration.
• Show that the model can generate accurate, large ensembles and higher-resolution forecasts efficiently.

Proposed method

• Adopt Adaptive Fourier Neural Operator (AFNO) architecture to model mappings between Earth-system state tensors across time.
• Use spectral (FFT-based) spatial mixing within AFNO blocks to enable high-resolution, global forecasting.
• Train on ERA5 reanalysis data at native 0.25-degree resolution with 20 prognostic variables, predicting 6-hourly steps (t -> t+6h) with pre-training and fine-tuning phases.
• Implement model and data parallelism (feature and data parallelism) with NCCL communications and CUDA graphs to maximize throughput and minimize CPU overhead.
• Preprocess data with NVIDIA DALI, and optimize via JIT fusion to reduce kernel launch overheads; employ FFT-based spectral convolutions and MLPs in the AFNO transformer blocks.

Experimental results

Research questions

• RQ1 Can AFNO-based architectures match or approach state-of-the-art weather forecast accuracy at high global resolutions?
• RQ2 How fast and scalable can FourCastNet train and infer on leadership-class HPC systems?
• RQ3 What are the energy and throughput implications of large-scale DL surrogates for global weather forecasting and ensembles?
• RQ4 Does the model support large ensembles and multi-step forecasts with acceptable skill at lead times (e.g., up to a week)?

Key findings

• Training on multiple HPC systems scales to 3,808 NVIDIA A100 GPUs with 140.8 petaFLOPS (11.9% of peak).
• Pre-training time on JUWELS Booster with 3,072 GPUs is 67.4 minutes, enabling 80,000x faster time-to-solution for inference relative to state-of-the-art NWP (IFS).
• Inference produces accurate instantaneous forecasts for up to a week and supports large ensembles; a 100-member 24-hour ensemble forecast runs in 12.41 node-seconds on a Selene node with 8 A100 GPUs, vs. 984,000 node-seconds for IFS (L91 18 km).
• Energy efficiency is dramatically higher, with an estimated 10,000x lower energy for FourCastNet ensemble runs compared to IFS ensembles.
• FourCastNet achieves high-resolution (8x the resolution of prior DL models) and improved accuracy over prior DL weather surrogates, including resolving fine-scale phenomena like tropical cyclones and atmospheric rivers.