[论文解读] Model-Parallel Fourier Neural Operators as Learned Surrogates for Large-Scale Parametric PDEs

Thomas J. Grady, Rishi Khan+6 more|arXiv (Cornell University)|Apr 4, 2022

Seismic Imaging and Inversion Techniques被引用 9

一句话总结

本文提出一种基于域分解的模型并行傅里叶神经算子（FNO）框架，可将FNO扩展至大规模3D和4D参数化PDE问题，实现超过26亿变量的推理与训练。在Perlmutter系统上，该方法实现了512张A100 GPU的弱扩展性，创下迄今为止最大的FNO推理与梯度计算规模。

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.

研究动机与目标

• 解决傅里叶神经算子（FNO）在处理高维输入与权重约束下的大规模3D和4D参数化PDE问题时的可扩展性限制。
• 利用分布式GPU资源，实现对超过26亿变量问题的FNO高效推理与训练。
• 开发一种基于域分解的模型并行FNO框架，实现跨多张GPU的扩展性，同时保持FNO的归纳偏置与性能。
• 验证分布式FNO在复杂现实应用（如地下CO2羽流模拟与反问题）中的可行性。
• 为此前常规数值求解器难以处理的大规模优化与不确定性量化任务，提供快速、可微分的PDE求解器。

提出的方法

• 实施一种域分解策略，将输入数据与网络权重在多张GPU间分区，以实现模型并行。
• 基于FNO的线性代数形式，推导出保持分区间数学一致性的分布式谱卷积操作。
• 利用CUDA感知的MPI实现前向与反向传播中的高效GPU间通信，最小化通信开销。
• 通过PyTorch中的DistDL库，将原始FNO架构（4个谱块，嵌入维度20）适配至分布式设置。
• 将空间维度扩展至512×512×512，时间维度扩展至10,240个时间步，实现在Perlmutter系统上的弱扩展性实验。
• 使用自动微分进行分布式训练与推理，支持反问题工作流中的梯度计算。

实验结果

研究问题

• RQ1 能否通过模型并行性将FNO扩展至超过64³空间限制的大规模3D与4D参数化PDE问题？
• RQ2 所提出的域分解FNO在最多512张A100 GPU上的弱扩展性实验中表现如何？
• RQ3 与传统求解器相比，分布式FNO在大规模PDE问题的推理与梯度计算方面能实现多大程度的加速？
• RQ4 该模型并行FNO框架能否有效应用于具有高维输出的现实世界地下CO2流动模拟？
• RQ5 在现代HPC系统（如Perlmutter）上，分布式FNO在问题规模与可扩展性方面的性能上限是什么？

主要发现

• 模型并行FNO在Perlmutter系统上实现了512张A100 GPU的弱扩展性，验证了其在512×512×512×20空间-时间维度问题上的可扩展性。
• 该框架支持超过26亿变量问题的推理与梯度计算，代表迄今为止最大的FNO应用规模。
• 由于FNO的全连接通信特性，实现略低于理想的扩展性，但在所有测试规模下均保持了强劲性能。
• 该方法成功在Azure云环境中对分布式FNO进行训练，用于模拟地下多相CO2动力学，验证了其在现实应用中的可行性。
• 该方法可实现快速、可微分的PDE求解，适用于反问题与优化任务，其推理时间比传统求解器快数个数量级。
• 本工作为在真实体积数据上利用分布式算子学习求解耦合反问题与统计问题奠定了基础。