Distributed Multigrid Neural Solvers on Megavoxel Domains
Distributed Multigrid Neural Solvers on Megavoxel Domains
复制标题
DOI：
10.1145/3458817.3476218
复制
发表时间：
2021-04
期刊：
影响因子：
--
通讯作者：
Aditya Balu;Sergio Botelho;Biswajit Khara;Vinay Rao;C. Hegde;S. Sarkar;Santi S. Adavani;A. Krishnamurthy;B. Ganapathysubramanian
中科院分区：
文献类型：
--
作者：
Aditya Balu;Sergio Botelho;Biswajit Khara;Vinay Rao;C. Hegde;S. Sarkar;Santi S. Adavani;A. Krishnamurthy;B. Ganapathysubramanian
我们考虑了大规模神经网络的分布式训练，这些神经网络用作产生全场输出的偏微分方程(PDE)解算器。我们特别考虑了广义三维Poisson方程的神经求解器。提出了一个可扩展的框架，它集成了两个截然不同的进步。首先，我们通过一种类似于数值线性代数中使用的多重网格技术的方法来加速训练大型模型。在这里，网络使用递增分辨率输入的层次化顺序来训练，类似于多网格方法中使用的“V”、“W”、“F”和“半V”循环。结合多重网格方法，我们实现了一个分布式深度学习框架，极大地减少了求解时间。我们展示了这种方法在GPU(云上的Azure VM)和CPU集群(PSC桥2)上的可扩展性。这种方法被用来训练一个通用的三维泊松解算器，该解算器可以很好地预测高维输入族的输出全场解，分辨率高达512×512×512。这一策略为在异质集群上快速、可扩展地训练神经偏微分方程组解算器提供了可能。
We consider the distributed training of large scale neural networks that serve as PDE (partial differential equation) 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 scalability of this approach on both GPU (Azure VMs on Cloud) and CPU clusters (PSC Bridges2). 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 × 512 × 512 for a high dimensional family of inputs. This strategy opens up the possibility of fast and scalable training of neural PDE solvers on heterogeneous clusters.