Distributed Multigrid Neural Solver on Megavoxel Domains

Authors: Aditya Balu (Iowa State University), Sergio Botelho (RocketML Inc), Biswajit Khara (Iowa State University), Vinay Rao (RocketML Inc), Soumik Sarkar (Iowa State University), Chinmay Hegde (New York University (NYU)), Adarsh Krishnamurthy (Iowa State University), Santi Adavani (RocketML Inc), and Baskar Ganapathysubramanian (Iowa State University)

Abstract: We consider the distributed training of neural networks that serve as PDE solvers producing full field outputs. We specifically consider neural solvers for the generalized 3D Poisson equation over megavoxel domains. A scalable framework is presented that integrates two distinct advances. First, we accelerate training a large model via a method analogous to the multigrid technique used in numerical linear algebra. Here, the network is trained using a hierarchy of increasing resolution inputs in sequence, analogous to the 'V’, 'W’, F', and 'Half-V' cycles used in multigrid approaches. In conjunction with the multi-grid approach, we implement a distributed deep learning framework which significantly reduces the time to solve. We show the scalability of this approach on both GPU and CPU clusters. This approach is deployed to train a generalized 3D Poisson solver that scales well to predict output full-field solutions up to the resolution of 512 x 512 x 512.




Back to Technical Papers Archive Listing