Asynchronous Heterogeneous Distributed Tensor Communication

Large-scale data-intensive linear algebra operations are hard to scale on distributed, heterogeneous systems. Further, the data distribution across machines makes it very difficult to compose implementations as each has their own requirements. A linear algebra abstraction is needed to compose implementations. There are many use cases for such an abstraction. We will focus on simulation of differential equations in the space of carbon capture sequestration (CCS). CCS is one of the most important technological problems needed to be solved in the energy industry today. One key component of this is doing many subsurface CO2 flow simulations to understand how plumes develop over time in sequestration reservoirs. Doing this with a traditional numerical solver in 3D would quickly become computationally and economically intractable at the scale and frequency these simulations would need to be run. In this proposal, we build the linear algebra abstraction library (LAAL) and implement a large scale distributed, heterogeneous application (Fourier Neural Operator) that can simulate differential equations orders of magnitude faster than conventional solvers (e.g. 100X). Using FNOs to solve subsurface CO2 simulations is ideal, as they require only an initial large upfront computational investment to train the model, which can then be used to make thousands of subsequent predictions. However, these currently do not scale beyond small problems. We will scale FNOs to industry-scale problem sizes. In Phase I, we focused on a key component of the FNO, the Fast Fourier Trans- form (FFT). A key to improving performance over state-of-the-art DOE implementations was overlapping the communication and computation. We show 2-3X improvement over NVIDIA’s CUDA implementation for problem sizes larger than the GPU memory. Further, while it is still a work in progress, we show that we are comparable if not faster than a state-of-the-art DOE FFT implementation (heFFTE). We used the learnings from the FFT work to build a distributed, heterogeneous FNO and showed weak scaling up to 512 GPUs. In Phase II we will build the LAAL, implement FFT and other linear operations using the LAAL, compose the operations into a large-scale FNO, and train the FNO to simulate CO2 plume evolution. We will work with our commercial partners (Chevron, Oxy, Microsoft) to implement a simulation-as-a-service model in the cloud. Companies across many industries run simulations to support and inform management decisions with large financial impacts. A simulation-as-a-service that can run 100X faster and at much lower costs enables many industries to obtain higher fidelity simulation results at lower cost enabling better decision-making.