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.

翻译：Mosaic Flow是一种新颖的区域分解方法，旨在将物理信息神经PDE求解器扩展到大规模域。其独特方法利用在小型域上预训练的网络，仅通过推理求解大型域上的偏微分方程，从而实现了高复用性。本文提出了Mosaic Flow的端到端并行化方案，结合了数据并行训练和用于大规模问题推理的域并行化。通过优化网络架构和数据并行训练，我们将在32个GPU上学习拉普拉斯算子的训练时间显著缩短至分钟级。此外，我们的分布式区域分解算法能够对训练域4096倍大小的域进行可扩展的拉普拉斯方程求解，在32个GPU上保持了强扩展性，同时维持了精度。Mosaic Flow的复用性，结合通过分布式内存算法实现的性能提升，使其成为建模复杂物理现象和加速科学发现的 promising 工具。