[论文解读] Highly-scalable, physics-informed GANs for learning solutions of stochastic PDEs

Liu Yang, Sean Treichler+9 more|arXiv (Cornell University)|Oct 29, 2019

Model Reduction and Neural Networks参考文献 29被引用 16

一句话总结

该论文提出了一种高度可扩展的、物理信息驱动的生成对抗网络（GAN）框架，用于在汉福德场地的地下水流建模中求解高维随机偏微分方程（PDE）。通过利用领域分解、多GPU并行计算以及物理约束训练，该方法在27,500块V100 GPU上实现了93.1%的加速效率，实现了每秒1207 PFlops的持续半精度性能——这是首个适用于混合随机PDE问题、随机维度高达1,000的百亿亿次计算（exascale-ready）GAN。

ABSTRACT

Uncertainty quantification for forward and inverse problems is a central challenge across physical and biomedical disciplines. We address this challenge for the problem of modeling subsurface flow at the Hanford Site by combining stochastic computational models with observational data using physics-informed GAN models. The geographic extent, spatial heterogeneity, and multiple correlation length scales of the Hanford Site require training a computationally intensive GAN model to thousands of dimensions. We develop a hierarchical scheme for exploiting domain parallelism, map discriminators and generators to multiple GPUs, and employ efficient communication schemes to ensure training stability and convergence. We developed a highly optimized implementation of this scheme that scales to 27,500 NVIDIA Volta GPUs and 4584 nodes on the Summit supercomputer with a 93.1% scaling efficiency, achieving peak and sustained half-precision rates of 1228 PF/s and 1207 PF/s.

研究动机与目标

• 解决汉福德场地地下水流建模中的不确定性量化问题，该问题复杂且高维，存在广泛污染与数据稀缺。
• 克服由空间异质性和多尺度相关长度引起的随机PDE（SPDE）求解中的维度灾难问题。
• 开发一种可扩展的、物理信息驱动的GAN（PI-GAN）架构，能够处理高达1,000维的随机维度。
• 通过领域分解和优化通信机制，实现在大规模HPC系统上高效、分布式的GAN训练。
• 将来自1,000至10,000口井的观测数据整合到统一的混合问题公式中，以推断压力水头和水力传导率场作为随机过程。

提出的方法

• 提出一种分层领域分解方案，将单个生成器和数百个判别器映射到多个GPU上，以实现可扩展训练。
• 实现一种物理信息损失函数，将随机PDE约束直接嵌入GAN训练过程，确保生成解的物理一致性。
• 采用WGAN-GP结合梯度惩罚以稳定训练，并提升高维随机过程中的分布保真度。
• 设计一种通信优化流水线，采用异步和重叠调度策略，最大限度减少空闲时间并提升GPU利用率。
• 在TensorFlow中对生成器和判别器架构进行优化，针对V100 GPU实现每节点59.3 TF/s和38.6 TF/s的峰值性能。
• 采用自由运行训练调度并结合动态负载均衡，以在4,584个Summit节点上保持高并行效率。

实验结果

研究问题

• RQ1 物理信息驱动的GAN框架能否有效建模来自地下水流、具有1,000多个随机维度的高维随机PDE？
• RQ2 如何有效将领域分解应用于具有单一全局生成器和多个本地判别器的GAN，以实现百亿亿次计算规模的训练？
• RQ3 在使用半精度算术于27,500块GPU上训练PI-GAN时，可实现的加速效率和性能水平如何？
• RQ4 将数千口井的观测数据纳入模型，能否显著提升推断的水力传导率场和压力水头分布的准确性与可靠性？
• RQ5 在大规模分布式环境中，尽管存在量化效应和GPU工作负载不平衡，该提出的GAN架构能否保持训练稳定性和收敛性？

主要发现

• 所提出的PI-GAN框架在4,584个Summit节点上实现了93.1%的加速效率，表明在极端规模下GAN训练具有近乎最优的强可扩展性。
• 该实现的持续半精度性能达到每秒1,207 PFlops，在27,500块V100 GPU上峰值性能达1,228 PFlops，标志着首个适用于SPDE的百亿亿次计算就绪GAN。
• 该方法成功求解了随机维度高达1,000的混合SPDE问题，远超以往基于GAN的SPDE求解器的能力。
• 自由运行训练调度实现了最高的持续性能（1,207 PF/s）和并行效率（93.1%），优于顺序调度和重叠调度。
• 该框架通过从有限传感器数据中学习分布，而非暴力采样，有效处理了高维、多尺度的随机过程（例如相关长度达到百万分之一量级）。
• 观察到量化和工作负载不平衡效应在特定GPU数量（如10,500块GPU）下导致轻微效率下降，但整体在所有测试配置下均保持稳健的可扩展性。