Model-Parallel Fourier Neural Operators as Learned Surrogates for Large-Scale Parametric PDEs

Abstract

Fourier neural operators (FNOs) are a recently introduced neural network architecture for learning solution operators of partial differential equations (PDEs), which have been shown to perform significantly better than comparable deep learning approaches. Once trained, FNOs can achieve speed-ups of multiple orders of magnitude over conventional numerical PDE solvers. However, due to the high dimensionality of their input data and network weights, FNOs have so far only been applied to two-dimensional or small three-dimensional problems. To remove this limited problem-size barrier, we propose a model-parallel version of FNOs based on domain-decomposition of both the input data and network weights. We demonstrate that our model-parallel FNO is able to predict time-varying PDE solutions of over 2.6 billion variables on Perlmutter using up to 512 A100 GPUs and show an example of training a distributed FNO on the Azure cloud for simulating multiphase CO${}_2$ dynamics in the Earth's subsurface.

References in corpus (4)

• Array Programming with NumPy
• FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators
• Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model
• A Robust Deep Learning Workflow to Predict Multiphase Flow Behavior during Geological CO2 Sequestration Injection and Post-Injection Periods

Cited by in corpus (7)

• Fourier Neural Operator Surrogate Model to Predict 3D Seismic Waves Propagation
• Comparison of Neural FEM and Neural Operator Methods for applications in Solid Mechanics
• Solving multiphysics-based inverse problems with learned surrogates and constraints
• Learned multiphysics inversion with differentiable programming and machine learning
• An uncertainty-aware Digital Shadow for underground multimodal CO2 storage monitoring
• Space-time parallel scaling of Parareal with a physics-informed Fourier Neural Operator coarse propagator applied to the Black-Scholes equation
• Superstep wavefield propagation