Multi-process thermodynamic graph learning for 2D fluid simulation

Solving partial differential equations (PDEs) for fluid simulation is computationally expensive, especially when dealing with complex geometries and high-resolution meshes. Recent advances in physics-informed graph neural networks (PIGNNs) have demonstrated potential in approximating such simulations more efficiently. Particularly, thermodynamic informed graph neural networks (TIGNNs) offer a promising data-driven alternative to traditional PDE solvers for fluid simulations. However, existing TIGNN implementations suffer from significant training inefficiencies, requiring prolonged runtimes and high memory consumption due to the need to maintain large parameter matrices in GPU memory. Inspired by multi-processor strategies in deformable solid simulations, we propose a novel multi-processor thermodynamic-informed graph neural network (MP-TIGNN) architecture to significantly accelerate training without compromising accuracy. Our approach enables faster convergence and reduces memory usage with mixed-precision training by leveraging fully sharded data parallelism (FSDP) across multiple GPUs. Experimental results show that our approach reduces training time by approximately 70% compared to the original setup while maintaining similar prediction accuracy.