Breaking Boundaries: Distributed Domain Decomposition with Scalable Physics-Informed Neural PDE Solvers

TL;DR

Mosaic Flow is a scalable, distributed domain decomposition method that uses pre-trained physics-informed neural networks to efficiently solve large-scale PDEs, significantly reducing training time and enabling inference on much larger domains.

Contribution

The paper introduces Mosaic Flow, a novel distributed domain decomposition approach that leverages pre-trained neural networks for scalable PDE solving on large domains.

Findings

Achieved training of Laplacian operator in minutes on 32 GPUs.

Enabled inference on domains 4096 times larger than training domains.

Demonstrated strong scaling and maintained accuracy on large-scale PDE problems.

Abstract

Mosaic Flow is a novel domain decomposition method designed to scale physics-informed neural PDE solvers to large domains. Its unique approach leverages pre-trained networks on small domains to solve partial differential equations on large domains purely through inference, resulting in high reusability. This paper presents an end-to-end parallelization of Mosaic Flow, combining data parallel training and domain parallelism for inference on large-scale problems. By optimizing the network architecture and data parallel training, we significantly reduce the training time for learning the Laplacian operator to minutes on 32 GPUs. Moreover, our distributed domain decomposition algorithm enables scalable inferences for solving the Laplace equation on domains 4096 times larger than the training domain, demonstrating strong scaling while maintaining accuracy on 32 GPUs. The reusability of Mosaic Flow, combined with the improved performance achieved through the distributed-memory algorithms, makes it a promising tool for modeling complex physical phenomena and accelerating scientific discovery.