Fast, Scalable Graph Neural Networks Used for Subsurface in Porous-Media Simulations

_

This article, written by JPT Technology Editor Chris Carpenter, contains highlights of paper SPE 229244, “Fast, Scalable, and Memory-Efficient High-Performance Graph Neural Networks for Subsurface in Porous-Media Simulations Using Multi-GPU Parallelism,” by Zeeshan Tariq and Mohsin Shaikh, King Abdullah University of Science and Technology. The paper has not been peer-reviewed.

_

Numerical simulation of underground CO2 storage requires accurate, fast, field-scale calculations of saturation and reservoir pressures, yet high-fidelity physics-based simulators are computationally expensive for rapid analysis. Recent advances in deep learning, especially those such as graph neural networks (GNNs), offer promising alternatives by approximating simulation outputs while reducing computational costs dramatically by minimizing the total reliance on traditional simulators. However, training GNN models on large-scale, spatiotemporal data sets remains a significant bottleneck when constrained to a single graphics processing unit (GPU). To address the computational bottleneck of training GNN on large-scale and high-resolution data sets, the authors propose multi-GPU scalable training of GNNs.

Advantages of GNNs

GNNs are particularly well-suited for subsurface flow because they preserve the relational structure that governs connectivity and transmissibility. Through message passing, GNNs aggregate local information to form multiscale representations that capture both sharp saturation fronts and broad pressure support. Coupled with temporal encoders or autoregressive unrolling, GNNs can approximate spatiotemporal evolution at inference costs that are orders of magnitude lower than those of full-order simulators, enabling near-real-time scenario evaluation and screening. However, building GNN surrogates that are sufficiently accurate and stable for subsurface storage applications presents a distinct systems challenge: training at scale on large spatiotemporal data sets. In a single-GPU setting, these characteristics quickly induce two hard constraints. First, wall-clock training time becomes prohibitive, hindering rapid iteration with regard to architectures, loss functions, and data curation. Second, device-memory limits restrict batch size and temporal context, curtailing model depth and representational capacity precisely when long-horizon stability is most needed. Fig. 1 shows the pros and cons of training using the multi-GPU parallelism approach.

Addressing issues of training GNN models requires codesign of the model architecture, the data pipeline, and the distributed training stack. The authors consider a U-Net-enhanced graph convolutional network (U-GCN). The spatial operator employs message-passing layers to encode local connectivity, while U-Net-style skip connections preserve high-frequency information and facilitate optimization in deeper networks. On the systems side, the training regimen emphasizes data-parallel optimization to transform throughput and memory efficiency. First, domain-aware graph minibatching combines cluster-based prepartitioning with controlled layer-wise neighbor sampling, reducing cut edges and moderating neighbor explosion to stabilize per-step computation. Second, dynamic bucketing groups subgraphs by approximate node or edge counts so that each rank processes similarly shaped batches, mitigating stragglers and improving kernel efficiency. Third, an asynchronous input pipeline performs subgraph extraction, feature packing, and pinned-memory staging on the host, overlapping input/output (I/O) with device-compute and tuning prefetch depth to saturate interconnect bandwidth without exhausting host resources. The novelty of the proposed approach lies in the integration of reservoir physics, high-resolution spatiotemporal data, and advanced GNN-based surrogate modeling into a unified, scalable workflow for underground subsurface simulation. The authors demonstrate that the proposed U-GCN, trained on more than 4,000 scenarios, achieves near-physics accuracy while providing over three orders-of-magnitude faster inference compared with traditional simulators.