Deep Generative Models that Solve PDEs: Distributed Computing for Training Large Data-Free Models

Abstract

Recent progress in scientific machine learning (SciML) has opened up the possibility of training novel neural network architectures that solve complex partial differential equations (PDEs). Several (nearly data free) approaches have been recently reported that successfully solve PDEs, with examples including deep feed forward networks, generative networks, and deep encoder-decoder networks. However, practical adoption of these approaches is limited by the difficulty in training these models, especially to make predictions at large output resolutions (≥ 1024 × 1024).Here we report on a software framework for data parallel distributed deep learning that resolves the twin challenges of training these large SciML models training in reasonable time as well as distributing the storage requirements. Our framework provides several out of the box functionality including (a) loss integrity independent of number of processes, (b) synchronized batch normalization, and (c) distributed higher-order optimization methods.We show excellent scalability of this framework on both cloud as well as HPC clusters, and report on the interplay between bandwidth, network topology and bare metal vs cloud. We deploy this approach to train generative models of sizes hitherto not possible, showing that neural PDE solvers can be viably trained for practical applications. We also demonstrate that distributed higher-order optimization methods are 2-3× faster than stochastic gradient-based methods and provide minimal convergence drift with higher batch-size.

Distributed multigrid neural solvers on megavoxel domains

“…One of the key outcomes of our experiments was to demonstrate a practical approach to train MGDiffNet on domain sizes up to 512 3 . We applied our framework to train MGDiffNet for resolutions up to 256 3 on GPU-based HPC clusters using on-demand multi-GPU virtual machines on Microsoft Azure.…”

“…We applied our framework to train MGDiffNet for resolutions up to 256 3 on GPU-based HPC clusters using on-demand multi-GPU virtual machines on Microsoft Azure. To train DiffNet for resolutions > 256 3 we used PSC Bridges2 HPC cluster with baremetal access to CPU nodes. We first talk about our experiments to study the multigrid approach and then the scaling studies using distributed deep learning.…”

“…Table 3 provides evidence of this conclusion by showing the dollar cost on Microsoft Azure multi-GPU virtual machines. Here, a 256 3 problem is trained in two ways: first, a "standalone" training on a fixed mesh of 256 3 size; and second, through a MG cycle, using a sequence of meshes of sizes 32 3 , 64 3 , 128 3 and 256 3 . This experiment reaffirms that the MG approach outperforms the standalone training.…”

“…Khoo et al [23] extend efforts for solving parametric PDEs. In additiona to these mathematical developments, recent work such as Botelho et al [3], and Yang et al [50] enable the scalable training of models used for solving PDEs. Specifically, Yang et al [50] demonstrates the scalability of the framework to 27,500 GPUs.…”

“…As the spatial domain increases, traditional PINN (and its variants) need a vast number of collocation points. Similarly, in the parametric setting, using a convolutional neural network [3,23], the voxel resolution creates computational and memory requirement challenges. For example, in Figure 2 we see that the computational time per epoch increases quadratically with the increase in the resolution of the spatial domain.…”