[论文解读] SciAI4Industry -- Solving PDEs for industry-scale problems with deep learning

Philipp Witte, Russell J. Hewett+3 more|arXiv (Cornell University)|Nov 23, 2022

Reservoir Engineering and Simulation Methods被引用 4

一句话总结

该论文提出 SciAI4Industry，一个基于 Julia 的分布式 API，用于在云端生成大规模科学深度学习的训练数据，结合通过域分解实现的模型并行深度学习。该方法可实现对工业级偏微分方程（PDE）问题的 3D 神经网络训练——以 200 万个单元的 CO₂ 流动模拟为例，实现了比传统求解器快 5 个数量级的推理速度和 3,200 倍的成本降低，且并行效率超过 90%。

ABSTRACT

Solving partial differential equations with deep learning makes it possible to reduce simulation times by multiple orders of magnitude and unlock scientific methods that typically rely on large numbers of sequential simulations, such as optimization and uncertainty quantification. Two of the largest challenges of adopting scientific AI for industrial problem settings is that training datasets must be simulated in advance and that neural networks for solving large-scale PDEs exceed the memory capabilities of current GPUs. We introduce a distributed programming API in the Julia language for simulating training data in parallel on the cloud and without requiring users to manage the underlying HPC infrastructure. In addition, we show that model-parallel deep learning based on domain decomposition allows us to scale neural networks for solving PDEs to commercial-scale problem settings and achieve above 90% parallel efficiency. Combining our cloud API for training data generation and model-parallel deep learning, we train large-scale neural networks for solving the 3D Navier-Stokes equation and simulating 3D CO2 flow in porous media. For the CO2 example, we simulate a training dataset based on a commercial carbon capture and storage (CCS) project and train a neural network for CO2 flow simulation on a 3D grid with over 2 million cells that is 5 orders of magnitudes faster than a conventional numerical simulator and 3,200 times cheaper.

研究动机与目标

• 解决在工业环境中为科学人工智能生成大规模训练数据集的挑战，因为传统模拟器过于缓慢且需要高性能计算（HPC）资源。
• 通过域分解实现模型并行，克服在单个 GPU 内存限制下训练大规模 3D PDE 深度神经网络的限制。
• 证明为现实工业问题（如碳捕集与封存（CCS）中的 3D CO₂ 流动）训练大规模深度学习代理模型的可行性。
• 为需要大量模拟的应用（如优化和不确定性量化）提供经济高效、高性能的推理能力。
• 通过 Julia 中的高级、云原生编程接口，使 HPC 级别的模拟能力实现普惠化。

提出的方法

• 开发一个高级 Julia API，抽象化云 HPC 基础设施，使用户无需管理集群即可并行生成大规模训练数据集。
• 通过域分解实现模型并行深度学习，将空间和时间域划分为多个 GPU 上的子域，以在内存限制内处理大规模 PDE 问题。
• 采用傅里叶神经算子（FNO）架构求解时间依赖型 PDE，通过在 GPU 间划分张量，实现 4D FFT 和 iFFT 操作中的高效通信。
• 利用单节点内 GPU 之间的高带宽 NVLink 互连，确保训练和推理过程中保持超过 90% 的高并行效率。
• 将分布式训练管道与云虚拟机（如 Azure ND96amsr）集成，实现训练和推理的可扩展性，同时最小化成本和管理开销。
• 应用混合并行化原则，为未来支持更大批量和多节点扩展做好准备，尽管当前实现仅使用模型并行，批量大小为 2。

实验结果

研究问题

• RQ1 高级、云原生的编程 API 是否能简化科学 AI 的大规模训练数据集生成，而无需用户管理 HPC 基础设施？
• RQ2 基于域分解的模型并行深度学习是否适用于解决单个 GPU 内存容量无法容纳的大规模 3D PDE 问题？
• RQ3 在单节点内跨多个 GPU 扩展时，该方法是否能实现高并行效率（例如 >90%）？
• RQ4 该方法生成的深度学习代理模型是否能在真实工业级 PDE 问题中实现比传统数值求解器更快的速度和更低的成本？
• RQ5 该框架在多大程度上能支持工业应用中的高吞吐量科学工作流（如不确定性量化和优化）？

主要发现

• 所提出的基于云的 API 实现了大规模 PDE 训练数据的可扩展、自动化生成，且用户无需管理 HPC 集群或云基础设施。
• 采用域分解的模型并行 FNO 训练在最多 8 块 A100 GPU 上实现了超过 90% 的并行效率，支持在超过 200 万个单元的 3D 网格上进行训练。
• 训练后的 FNO 模型对多孔介质中 3D CO₂ 流动的模拟速度比传统 OPM 求解器快 5 个数量级（单次模拟 0.12 秒 vs. 6.8 小时）。
• 在考虑云虚拟机成本的情况下，FNO 模型的单次模拟成本仅为 OPM 求解器的 1/3,200（单次成本 0.11 美元 vs. 3.4 美元）。
• FNO 代理模型在仅 1,848 次模拟后即可实现成本盈亏平衡，使其在高吞吐量应用（如井位优化）中具备经济可行性。
• 该框架支持在 3D 中对高达 512³ 网格或在 2D 中对高达 12,000² 网格的静态 PDE 问题进行训练，展示了对商业规模问题的可扩展性。