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

arXiv: https://arxiv.org/abs/2204.01205 · HTML: https://mlsanity.com/paper/2204.01205/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 54 references that are in the map.
- First posted on arXiv 4 Apr 2022.

## Builds on (most cited first, 16 of 16)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Neural Operator: Graph Kernel Network for Partial Differential Equations](https://mlsanity.com/paper/2003.03485/) (2020)
- [DGM: A deep learning algorithm for solving partial differential equations](https://mlsanity.com/paper/1708.07469/) (2017)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [U-FNO -- An enhanced Fourier neural operator-based deep-learning model for multiphase flow](https://mlsanity.com/paper/2109.03697/) (2021)
- [Mesh-TensorFlow: Deep Learning for Supercomputers](https://mlsanity.com/paper/1811.02084/) (2018)
- [A Robust Deep Learning Workflow to Predict Multiphase Flow Behavior during Geological CO2 Sequestration Injection and Post-Injection Periods](https://mlsanity.com/paper/2107.07274/) (2021)
- [Deep-learning-based coupled flow-geomechanics surrogate model for CO2 sequestration](https://mlsanity.com/paper/2105.01334/) (2021)
- [Learned coupled inversion for carbon sequestration monitoring and forecasting with Fourier neural operators](https://mlsanity.com/paper/2203.14396/) (2022)
- [A Linear Algebraic Approach to Model Parallelism in Deep Learning](https://mlsanity.com/paper/2006.03108/) (2020)
- [Fast parallel multidimensional FFT using advanced MPI](https://mlsanity.com/paper/1804.09536/) (2018)

## Built on by (most cited first, 13 of 13)

- [Neural Operators for Accelerating Scientific Simulations and Design](https://mlsanity.com/paper/2309.15325/) (2023)
- [Sensitivity-Constrained Fourier Neural Operators for Forward and Inverse Problems in Parametric Differential Equations](https://mlsanity.com/paper/2505.08740/) (2025)
- [Active operator learning with predictive uncertainty quantification for partial differential equations](https://mlsanity.com/paper/2503.03178/) (2025)
- [Fourier Neural Operator Surrogate Model to Predict 3D Seismic Waves Propagation](https://mlsanity.com/paper/2304.10242/) (2023)
- [Solving multiphysics-based inverse problems with learned surrogates and constraints](https://mlsanity.com/paper/2307.11099/) (2023)
- [Learned multiphysics inversion with differentiable programming and machine learning](https://mlsanity.com/paper/2304.05592/) (2023)
- [Inference of CO2 flow patterns -- a feasibility study](https://mlsanity.com/paper/2311.00290/) (2023)
- [Fourier Neural Operator based surrogates for CO2 storage in realistic geologies](https://mlsanity.com/paper/2503.11031/) (2025)
- [Localized PCA-Net Neural Operators for Scalable Solution Reconstruction of Elliptic PDEs](https://mlsanity.com/paper/2509.18110/) (2025)
- [Joint Microseismic Event Detection and Location with a Detection Transformer](https://mlsanity.com/paper/2307.09207/) (2023)
- [Radial Basis Operator Networks](https://mlsanity.com/paper/2410.04639/) (2024)
- [Neural surrogates for designing gravitational wave detectors](https://mlsanity.com/paper/2511.19364/) (2025)
- [Fast Surrogate Modeling of Excitable and Oscillatory FitzHugh-Nagumo Dynamics with Parametric Neural Operators](https://mlsanity.com/paper/2609.04549/) (2026)

## Abstract (arXiv)

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 CO2 dynamics in the Earth's subsurface.