Breaking Boundaries: Distributed Domain Decomposition with Scalable
  Physics-Informed Neural PDE Solvers
Arthur Feeney Zitong Li Ramin Bostanabad Aparna Chandramowlishwaran
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.