PaperlayerView exact revision on arXiv

Source-linked AI summary

FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators

Thorsten Kurth, Shashank Subramanian, Peter Harrington, Jaideep Pathak, Morteza Mardani, David Hall, Andrea Miele, Karthik Kashinath, Animashree Anandkumar

arXiv:2208.05419v1physics.ao-phcs.AIcs.CVcs.LGcs.PF

TL;DR

High-resolution weather forecasting is limited by the computational cost of numerical prediction, especially for ensembles and finer resolutions. FourCastNet uses a data-driven AFNO-based Earth-system emulator to generate global forecasts, reporting 80,000-times-faster ensemble inference and higher resolution than prior deep-learning surrogates.

  • Problem

    High-resolution numerical weather prediction is computationally expensive, while ensembles and finer resolutions require substantially greater resources.

  • Method

    FourCastNet uses an AFNO-based transformer that tokenizes global weather fields, mixes spatial and channel information, and reconstructs the next forecast frame.

  • Results

    80,000 times faster than IFS, FourCastNet generates a 100-member ensemble forecast in 12.41 node-seconds instead of 984,000 node-seconds.

  • Takeaways & Limitations

    FourCastNet supports higher-resolution forecasting and large ensembles that can resolve extreme events such as tropical cyclones and atmospheric rivers.

Abstract

from arXiv · show

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.

I. HIGHLIGHTS

FourCastNet targets the computational and scalability barriers that limit high-resolution weather and climate modeling, especially for ensembles, scenarios, and interactive analysis. It reports substantially faster global weather emulation while scaling to thousands of GPUs and higher resolution.

  • 67.4 minutes trains FourCastNet on 3,072 GPUs, while scalability extends to 3,808 GPUs at 8X higher resolution than state-of-the-art deep learning models.
  • 80,000X reduces the 24-hour 100-member ensemble forecast time from 984,000 to 12.4 node-seconds versus state-of-the-art numerical simulations.
  • Climate change amplifies extreme weather frequency and intensity, increasing stress on natural and human ecosystems.
  • High-resolution Earth-system simulation is constrained by coupled nonlinear physics, enormous dimensionality, and computational costs that grow sharply with resolution.
  • Large ensembles are needed to characterize uncertainty and confidently predict rare events at the distribution tails.
  • Weather and climate workflows require substantial throughput for repeated forecasts, scenarios, experiments, and multidecadal projections on thousands of nodes.

III. STATE OF THE ART

The current weather and climate modeling landscape is shaped by two approaches: numerical prediction and deep learning.

  • Weather and climate modeling currently centers on numerical prediction and deep learning.

A. Numerical Prediction

Numerical prediction has advanced through high-performance computing, but kilometer-scale global simulations remain costly and difficult to scale. Deep learning and neural operators provide alternative computational approaches for modeling complex physical systems.

  • A. Numerical Prediction: HPC advances have enabled global weather and climate simulations at kilometer scale, with simulation performance doubling every two years.
  • A. Numerical Prediction: 0.3 SYPD was achieved for hydrostatic IFS simulations on close to 4,000 Summit nodes, while non-hydrostatic simulations reached 0.085 SYPD on close to 5,000 Piz Daint nodes.
  • A. Numerical Prediction: Kilometer-scale ICON global simulations require approximately 1,000 nodes at minimum because of memory requirements, with scalability also limited by load imbalances.
  • A. Numerical Prediction: Single numerical weather and climate runs approach 1 SYPD on thousands of leadership-class nodes, while large ensembles and scenarios multiply energy requirements.
  • A. Numerical Prediction: Neural operators learn mappings between infinite-dimensional spaces, and Fourier Neural Operators efficiently implement global convolution using FFTs for nonlinear PDEs.

1) Operator learning and transformers:

High-resolution weather data makes standard vision-transformer attention computationally prohibitive, motivating AFNO-based architectures. FourCastNet combines this efficient approach with strong accuracy and substantially faster ensemble inference.

  • 1) Operator learning and transformers:: At kilometer-scale weather resolution, vision-transformer self-attention becomes prohibitive because computation and memory scale quadratically with image resolution.
  • 1) Operator learning and transformers:: Sparse, low-rank, and kernel attention approximations address transformer inefficiency but fall short at the desired kilometer-scale weather resolution.
  • 1) Operator learning and transformers:: AFNO handles continuous inputs and high resolutions efficiently, and FourCastNet uses it to achieve unprecedented accuracy and resolution.
  • 1) Operator learning and transformers:: FourCastNet provides 8X greater resolution than the leading deep-learning surrogate, 25 km versus 200 km, while resolving tropical cyclones and atmospheric rivers.
  • 1) Operator learning and transformers:: 80,000 times faster than IFS, FourCastNet generates a 100-member ensemble forecast in 12.41 node-seconds instead of 984,000 node-seconds.
  • 1) Operator learning and transformers:: Figure 1 compares ACC and RMSE for Z500 across FourCastNet, IFS, and a baseline DLWP model, reporting FourCastNet's superior DLWP performance at 8X higher resolution.

IV. INNOVATIONS

The AFNO transformer converts high-resolution weather fields into token embeddings, applies Fourier-domain spatial mixing and channel mixing, then reconstructs the forecast field. FourCastNet is trained for autoregressive multi-step prediction using its own intermediate output.

  • AFNO transformer: FourCastNet divides each 20×720×1440 input field into 180×360 patches and embeds the resulting tokens into high-dimensional vectors.The patch size is p = 4, and an example embedding has 1,024 numerical entries.
  • AFNO transformer: The AFNO spatial mixer applies an FFT, block-diagonal Fourier-space channel mixing, and an inverse FFT.The block-diagonal MLP processes each token’s 1,024 channels, while the inverse transform returns the representation to the spatial domain.
  • AFNO transformer: The final linear decoder reconstructs pixels for the forecast frame X(t+∆t) from the last transformer-layer embedding.Each transformer layer performs both token mixing and channel mixing; the channel mixer is a single-layer MLP.
  • Parallel forecast architecture: The architecture produces a 6-hour forecast with the same 20×720×1440 dimensions as its input.Model parallelism splits channels across GPUs; FFT spatial mixing operates on disjoint blocks, while channel mixing requires communication.
  • Autoregressive training: For two-step training, the model predicts X(t+1), feeds that prediction back as input, and minimizes the summed losses against both ground-truth outputs.This procedure starts from a previously pre-trained model.

C. Implementation Details

FourCastNet’s implementation combines PyTorch-based spectral operations, feature and data parallelism, optimized data movement, and GPU execution techniques designed to reduce communication and launch overhead.

  • Framework: FourCastNet is implemented in PyTorch to combine rapid prototyping with low-level performance optimization.PyTorch is described as a domain-specific language for deep learning in the supplied passage.
  • Parallelism: Feature parallelism splits channels across GPUs, while data parallelism splits the global batch across independent model instances.Feature-mixing operations require communication, whereas FFTs that do not mix feature maps are embarrassingly parallel.
  • Spectral convolution: The spectral-convolution layer uses real-to-complex and complex-to-real FFTs, with batched transforms over batch and feature dimensions.Its spectral-space MLP is implemented with torch.einsum.
  • Spectral convolution: Complex tensor layouts are chosen so PyTorch can generate a single batched complex GEMM through cuBLAS.The torch.einsum kernel cost can vary by up to 400% depending on weight and input data layout.
  • MLP: MLP matrix multiplications use pointwise convolutions, with row-wise parallelism in the first multiplication and column-wise parallelism in the second.The hidden dimension has 4,096 units compared with a 1,024-unit embedding size.
  • I/O pipeline: The I/O pipeline loads yearly HDF5 arrays, normalizes variables by mean and variance, adds Gaussian noise, and uses NVIDIA DALI with double buffering.The supplied passage ends while describing pinned host memory.
  • Execution optimization: CUDA graphs capture the full forward and backward training step to reduce CPU interference with kernel launches.The optimizer accounts for 5–8% of overall training time, while I/O runs concurrently with training.
  • Execution optimization: torch.jit.script fuses lightweight operations in composite layers and metrics to reduce launch latency and improve memory locality.This is especially used for validation code not captured in a CUDA graph.

V. PERFORMANCE MEASUREMENT

Performance is evaluated primarily through time-to-solution and measured FLOPS, using full training runs where feasible and short scaled runs to estimate epoch-time behavior across systems.

  • Time-to-solution: 80 epochs define FourCastNet pre-training for single-step prediction, with the target chosen empirically to optimize 192-hour forecasting skill after fine-tuning.Pre-training is later followed by a more expensive multi-timestep training procedure.
  • Time-to-solution: Because 80-epoch runs are computationally prohibitive at every scale, scaled configurations are trained for 5 epochs to obtain per-epoch timing statistics.The method assumes the number of epochs to convergence is the same across scales, making per-epoch time directly correlated with total training time.
  • Timing procedure: Epoch timers include validation and exposed I/O overhead, while scaffolding time is excluded from measurements.Scaffolding includes NCCL communicator startup, 20 warm-up steps, and CUDA graph capture.
  • FLOPS measurement: Floating-point operations are measured with NVIDIA Nsight Compute using instruction-specific weighting factors.Weights include 1 for multiplication or addition, 2 for FP32 and FP64 FMA, and 4 for FP16 double FMA.
  • FLOPS measurement: Table I summarizes the FLOPS weighting factors for GPU instructions relevant to the implementation.
  • FLOPS measurement: The total FLOPS estimate multiplies per-iteration FLOPS by the number of iterations, assuming perfect load balance across model- and data-parallel ranks.Validation FLOPS are also included.
  • Measurement caveat: The reported peak FLOPS rate is underestimated because it is averaged over a full epoch and excludes CPU work and data-parallel weight-gradient reductions.The omitted contributions are described as a small fraction of the overall FLOPS count.

B. Systems and Environment

FourCastNet is evaluated on Selene, Perlmutter, and JUWELS Booster, three A100-based supercomputing environments with different node counts, interconnects, memory capacities, and storage systems.

  • Perlmutter: Perlmutter’s GPU partition contains 1,536 nodes with four NVIDIA A100-40 GB GPUs per node connected by NVLink3.It uses the HPE Cray Slingshot 10 interconnect and an all-flash Lustre global file system.
  • JUWELS Booster: JUWELS Booster contains 936 nodes with four NVIDIA A100-40 GB GPUs per node and uses NVLink3 plus DragonFly+ HDR200 InfiniBand.Its global file system is GPFS-based, with a peak bandwidth of 1 TB/s, and experiments use an NVMe-based cache layer.
  • Selene: Selene comprises 560 DGX-A100 nodes, each with eight NVIDIA A100-80 GB GPUs connected through NVLink3 groups.The GPUs are partitioned into two groups of four, with additional NVLink3 connections between the groups.

3) Selene:

The Selene experiments quantify FourCastNet’s single-instance performance and show that increasing model size lowers per-GPU throughput because communication becomes more significant.

  • Measurement: FLOPS per iteration are reported per GPU for different model instance sizes and precisions, with tensor-instruction FLOPS combined into one category.The measurement uses batch size 1 and includes data loading, preprocessing, forward and backward passes, weight updates, and scheduling.
  • Performance: 160 TFLOP/s on one GPU corresponds to 51% of dense FP16 tensor-core peak performance.The average step time is about 517 ms for model instance size 1.
  • Performance: Per-GPU performance decreases from 160 TFLOP/s to 115 TFLOP/s and 76 TFLOP/s for model instance sizes 1, 2, 4, and 8.The reported decline is associated with additional communication in forward and backward passes.
  • Operation breakdown: Communication cannot be overlapped with computation because the network is strictly feed-forward and lacks branches with substantial compute.The time breakdown is reported in Table III for different model instance sizes.
  • Operation breakdown: Communication becomes the dominant part of each iteration for larger model instance sizes.FFT time remains subdominant because AFNO patch embedding produces grids with very small FFT dimensions.

B. Scalability

FourCastNet is evaluated across JUWELS Booster, Perlmutter, and Selene to characterize scaling across model and data parallelism. The results show system-dependent scaling behavior and high aggregate performance at thousands of GPUs.

  • Experimental setup: FourCastNet scaling experiments bind GPUs to nearby cores, memory, and interconnects across all three systems.Selene additionally partitions model instances so GPUs on the same socket ideally share an instance; this is unnecessary on the all-to-all-connected Perlmutter and JUWELS Booster systems.
  • Scaling behavior: Model instances of size 4 outperform size 2 on JUWELS Booster and Perlmutter, whereas Selene gains more from increased data parallelism than model parallelism.Selene’s balanced intranode and internode bandwidth supports hybrid NVLink/InfiniBand algorithms and reduces the impact of large collective operations.
  • Figures: Validation loss is plotted against wall-clock time for multiple FourCastNet configurations on JUWELS Booster, Perlmutter, and Selene.The figure indicates reduced solution times as parallelism increases.
  • Aggregate performance: 140.8 PFLOP/s is reached on 3,808 GPUs, compared with 99.1 PFLOP/s on 3,072 GPUs and 91.8 PFLOP/s on 1,024 GPUs.These are the peak performances measured on Perlmutter, JUWELS Booster, and Selene, respectively.

C. Convergence and Time-to-Solution

FourCastNet’s scalable training and inference substantially reduce weather-emulation time while preserving a pathway toward higher-resolution, larger-scale Earth-system modeling. The approach also exposes a key boundary: changing climate conditions make extrapolation beyond observed data a central challenge.

  • Training time-to-solution falls to just over an hour through increased model parallelism across major supercomputing systems.A baseline configuration without model parallelism is estimated at roughly 40 hours.
  • Efficient scaling enables larger models, faster experimentation, and more rapid hyper-parameter exploration.The authors report that this is achieved by consuming more resources for shorter periods.
  • Five orders-of-magnitude speedup over traditional numerical simulations is achieved while maintaining high accuracy.
  • FourCastNet’s resolution, dimensionality, and scalability position it for kilometer-scale emulation, although current weather emulation operates at 25 km.The proposed path includes smaller AFNO patch sizes, larger batch sizes, and more GPUs.
  • Inference efficiency addresses throughput and data-movement bottlenecks, reaching close to 100,000 simulated years per day at the cost of 1 simulated year per day conventionally.A century-long kilometer-scale climate simulation could generate an exabyte of data when outputs are written every 30 simulated minutes.
  • Because Earth’s climate is changing and observations degrade backward in time, future FourCastNet generations will need physical laws to extrapolate into unprecedented futures.

B. Implications for Future Supercomputing Systems

FourCastNet is presented as a foundation for AI-enabled Earth-system computing beyond conventional forecasting workflows. Its throughput and scalability could support digital twins, virtualization, and interactive access to weather and climate information.

  • FourCastNet makes in-situ deep-learning training and analytics a near-term possibility for exascale weather and climate computing.The passage connects this possibility to emerging AI supercomputers and digital twins of Earth.
  • FourCastNet could help build digital twins that let users access information and intervene despite exabyte-scale weather and climate data.
Loading 2208.05419v1…