Model-parallel Fourier neural operators as learned surrogates for large-scale parametric PDEs

Grady, Thomas J.; Khan, Rishi; Louboutin, Mathias; Yin, Ziyi; Witte, Philipp A.; Chandra, Ranveer; Hewett, Russell J.; Herrmann, Felix J.

doi:10.1016/j.cageo.2023.105402

Model-parallel Fourier neural operators as learned surrogates for large-scale parametric PDEs

Journal Article · Sat Jun 10 00:00:00 EDT 2023 · Computers and Geosciences

DOI:https://doi.org/10.1016/j.cageo.2023.105402· OSTI ID:2421940

^[1]; Khan, Rishi ^[2]; Louboutin, Mathias ^[3]; Yin, Ziyi ^[3]; Witte, Philipp A. ^[4]; Chandra, Ranveer ^[4]; Hewett, Russell J. ^[5]; Herrmann, Felix J. ^[3]

Georgia Institute of Technology, Atlanta, GA (United States); OSTI
Extreme Scale Solutions, New Castle, DE (United States)
Georgia Institute of Technology, Atlanta, GA (United States)
Microsoft Corporation, Redmond, WA (United States)
Microsoft Corporation, Redmond, WA (United States); Virginia Polytechnic Inst. and State Univ. (Virginia Tech), Blacksburg, VA (United States)

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. Here, 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₂ dynamics in the Earth’s subsurface.

View Accepted Manuscript (DOE)

Research Organization:: Extreme Scale Solutions, LLC, Newark, DE (United States); Lawrence Berkeley National Laboratory (LBNL), Berkeley, CA (United States). National Energy Research Scientific Computing Center (NERSC); Oak Ridge National Laboratory (ORNL), Oak Ridge, TN (United States). Oak Ridge Leadership Computing Facility (OLCF); Univ. of California, Oakland, CA (United States); Virginia Polytechnic Inst. and State Univ. (Virginia Tech), Blacksburg, VA (United States)

Sponsoring Organization:: USDOE Office of Science (SC)

Grant/Contract Number:: AC02-05CH11231; AC05-00OR22725; SC0021515; SC0022041

OSTI ID:: 2421940

Journal Information:: Computers and Geosciences, Journal Name: Computers and Geosciences Journal Issue: C Vol. 178; ISSN 0098-3004

Publisher:: ElsevierCopyright Statement

Country of Publication:: United States

Language:: English

References (18)

Application of a particle swarm optimization algorithm for determining optimum well location and type Onwunalu, Jérôme E.; Durlofsky, Louis J. Computational Geosciences, Vol. 14, Issue 1 https://doi.org/10.1007/s10596-009-9142-1	journal	May 2009
MIN3P-HPC: A High-Performance Unstructured Grid Code for Subsurface Flow and Reactive Transport Simulation Su, Danyang; Mayer, K. Ulrich; MacQuarrie, Kerry T. B. Mathematical Geosciences https://doi.org/10.1007/s11004-020-09898-7	journal	October 2020
PARAMESH: A parallel adaptive mesh refinement community toolkit MacNeice, Peter; Olson, Kevin M.; Mobarry, Clark Computer Physics Communications, Vol. 126, Issue 3 https://doi.org/10.1016/S0010-4655(99)00501-9	journal	April 2000
Full-waveform inversion on heterogeneous HPC systems Gokhberg, Alexey; Fichtner, Andreas Computers & Geosciences, Vol. 89 https://doi.org/10.1016/j.cageo.2015.12.013	journal	April 2016
The Open Porous Media Flow reservoir simulator Rasmussen, Atgeirr Flø; Sandve, Tor Harald; Bao, Kai Computers & Mathematics with Applications, Vol. 81 https://doi.org/10.1016/j.camwa.2020.05.014	journal	January 2021
20 Years of Monitoring CO2-injection at Sleipner Furre, Anne-Kari; Eiken, Ola; Alnes, Håvard Energy Procedia, Vol. 114 https://doi.org/10.1016/j.egypro.2017.03.1523	journal	July 2017
Carbon capture and storage Gibbins, Jon; Chalmers, Hannah Energy Policy, Vol. 36, Issue 12 https://doi.org/10.1016/j.enpol.2008.09.058	journal	December 2008
DGM: A deep learning algorithm for solving partial differential equations Sirignano, Justin; Spiliopoulos, Konstantinos Journal of Computational Physics, Vol. 375 https://doi.org/10.1016/j.jcp.2018.08.029	journal	December 2018
Well placement optimization: A survey with special focus on application for gas/gas-condensate reservoirs Nasrabadi, Hadi; Morales, Adrian; Zhu, Ding Journal of Natural Gas Science and Engineering, Vol. 5 https://doi.org/10.1016/j.jngse.2011.10.002	journal	March 2012
Fast parallel multidimensional FFT using advanced MPI Dalcin, Lisandro; Mortensen, Mikael; Keyes, David E. Journal of Parallel and Distributed Computing, Vol. 128 https://doi.org/10.1016/j.jpdc.2019.02.006	journal	June 2019
Efficient well placement optimization based on theory-guided convolutional neural network Wang, Nanzhe; Chang, Haibin; Zhang, Dongxiao Journal of Petroleum Science and Engineering, Vol. 208 https://doi.org/10.1016/j.petrol.2021.109545	journal	January 2022
Multi-solution well placement optimization using ensemble learning of surrogate models Salehian, Mohammad; Haghighat Sefat, Morteza; Muradov, Khafiz Journal of Petroleum Science and Engineering, Vol. 210 https://doi.org/10.1016/j.petrol.2021.110076	journal	March 2022
Array programming with NumPy Harris, Charles R.; Millman, K. Jarrod; van der Walt, Stéfan J. Nature, Vol. 585, Issue 7825 https://doi.org/10.1038/s41586-020-2649-2	journal	September 2020
Physics-informed machine learning Karniadakis, George Em; Kevrekidis, Ioannis G.; Lu, Lu Nature Reviews Physics, Vol. 3, Issue 6 https://doi.org/10.1038/s42254-021-00314-5	journal	May 2021
Reflecting on the Goal and Baseline for Exascale Computing: A Roadmap Based on Weather and Climate Simulations Schulthess, Thomas C.; Bauer, Peter; Wedi, Nils Computing in Science & Engineering, Vol. 21, Issue 1 https://doi.org/10.1109/MCSE.2018.2888788	journal	January 2019
PFFT: An Extension of FFTW to Massively Parallel Architectures Pippig, Michael SIAM Journal on Scientific Computing, Vol. 35, Issue 3 https://doi.org/10.1137/120885887	journal	January 2013
DeepSpeed Rasley, Jeff; Rajbhandari, Samyam; Ruwase, Olatunji Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining https://doi.org/10.1145/3394486.3406703	conference	August 2020
Devito (v3.1.0): an embedded domain-specific language for finite differences and geophysical exploration Louboutin, Mathias; Lange, Michael; Luporini, Fabio Geoscientific Model Development, Vol. 12, Issue 3 https://doi.org/10.5194/gmd-12-1165-2019	journal	January 2019

Similar Records

Fourier-MIONet: Fourier-enhanced multiple-input neural operators for multiphase modeling of geological carbon sequestration

Journal Article · Fri Jul 26 20:00:00 EDT 2024 · Reliability Engineering and System Safety · OSTI ID:2527395

Solving Seismic Wave Equations on Variable Velocity Models With Fourier Neural Operator

Journal Article · Wed Nov 15 19:00:00 EST 2023 · IEEE Transactions on Geoscience and Remote Sensing · OSTI ID:2507005

Learning the boundary-to-domain mapping using Lifting Product Fourier Neural Operators for partial differential equations

Conference · Sun Jun 30 20:00:00 EDT 2024 · OSTI ID:2429817

Related Subjects

58 GEOSCIENCES
97 MATHEMATICS AND COMPUTING
Carbon capture sequestration
Deep learning
Learned surrogates
Model parallelism
Multiphase flow
Operator learning

Model-parallel Fourier neural operators as learned surrogates for large-scale parametric PDEs

Citation Formats

References (18)

Similar Records

Related Subjects