Breaking Boundaries: Distributed Domain Decomposition with Scalable Physics-Informed Neural PDE Solvers

In this work, we develop SciML models to solve boundary value problems (BVP) of the form

The vectors $\bm{x}$ are in the domain $\Omega$ or on the domain boundary $\partial\Omega$. The function $u$ is the solution of the differential equation. The differential operator is denoted by $D$, while $B$ is the boundary operator. The forcing function is $f$, and $g$ is the boundary function. A classic example of a BVP is the 2D Laplace equation with a Dirichlet boundary condition:

where the vector $\bm{x}=[x,y]$ and $\Delta=(\partial^{2}/\partial x^{2}+\partial^{2}/\partial y^{2})$ is the Laplacian operator (Evans, 2010).

Neural PDE solvers (or neural solvers for short) (Li et al., 2020a; Kovachki et al., 2021; Lu et al., 2021a) are a type of model that learns to approximate the PDE solution operator and solve various instances of a BVP with different boundary conditions. They are trained using a dataset that consists of solved boundary value problems for a specific PDE. A neural solver may take a discretized boundary function as input, denoted by $\hat{{\bm{g}}}=\{g(\bm{x}_{bc}^{1}),\dots,g(\bm{x}_{bc}^{N})\}$, where $\bm{x}_{bc}^{i}$ are $N$ points sampled on the boundary. $\hat{{\bm{g}}}$ specifies the particular instance of the BVP to solve. Therefore, an neural solver, represented by the function $\mathcal{N}(\bm{x},\hat{{\bm{g}}};\theta)$, approximates the solution $u({\bm{x}})$ for the instance of the BVP determined by the boundary function $g$.

This study employs a special type of network called a physics-informed neural PDE solver (Wang et al., 2021b; Raissi et al., 2019). While neural solvers trained on labeled input-output pairs can learn the solution operator of a PDE, their ability to generalize to out-of-distribution data, such as boundary or initial conditions outside the training set, significantly increases the demand on the training dataset size. In SciML, data is often scarce and computationally expensive to generate. To address this challenge, physics-informed neural solvers adopt a similar strategy to PINNs by incorporating an additional PDE loss as a form of regularization. This physics loss effectively constrains the space of possible solutions, softly enforcing the PDE constraints. By incorporating domain knowledge into the training process, the model is more robust to noise and uncertainties present in the dataset. As an example, for the Laplace equation, the PDE loss for a batch of $n$ collocation points ${\bm{X}}=\{{\bm{x}}_{1},\dots,{\bm{x}}_{n}\}\subseteq\Omega$ is defined as

Intuitively, as the loss function is minimized during training, the PDE residual will approach zero, $\Delta\mathcal{N}(\bm{x},\hat{{\bm{g}}};\theta)\rightarrow 0$, indicating that the network approximately satisfies the Laplace equation $\Delta\mathcal{N}(\bm{x},\hat{{\bm{g}}};\theta)\approx\Delta u(\bm{x})=0$.

Domain decomposition methods are widely used in solving boundary value problems (Balay et al., 2019; Hecht, 2012; Dolean et al., 2015). These methods partition the domain of a BVP into smaller subdomains, and then iteratively combine solutions of the subdomains to develop the global solution. Domain decomposition methods enable scaling across multiple nodes, making them a powerful tool for scaling PDE solvers.

The classic example of domain decomposition is the alternating Schwarz method (ASM) (Schwarz, 1869; Gander, 2008). ASM relies on overlapping subdomains to ensure continuity and information propagation between subdomains. While Schwarz methods are commonly used as preconditioners for Krylov methods (Dolean et al., 2016), in this work we use a variant of ASM as the solver.

Continuing with the Laplace example, the domain $\Omega$ can be partitioned into two subdomains $\Omega_{1}$ and $\Omega_{2}$, such that $\Omega_{1}\cap\Omega_{2}\not=\emptyset$. The subdomain interfaces are $\Gamma_{1}=\partial\Omega_{1}\cap\Omega_{2}$ and $\Gamma_{2}=\Omega_{1}\cap\partial\Omega_{2}$. To solve the global domain using ASM, the following routine is applied iteratively:

The superscript denotes the iteration of the solution: $u^{n}_{i}$ is the solution of subdomain $i$ at iteration $n$. The process alternates between solving for $u_{1}$ and $u_{2}$. The solution for $u_{1}$ is used to set the interface condition on $\Gamma_{2}$. Then, with this condition, we solve for $u_{2}$. Subsequently, the solution for $u_{2}$ is used to set the interface condition on $\Gamma_{1}$. Schwarz proved the convergence of this iterative scheme for general elliptic PDEs (Schwarz, 1869). Lions proved that ASM can be used to solve systems with an arbitrary number of subdomains, and a parallel version of ASM exhibits similar convergence properties to the original Schwarz method (Lions et al., 1988). However, it is worth noting that the convergence rate of Schwarz methods is signifiantly influenced by the mesh parameter and overlap. A system consisting of many subdomains with little overlap require more iterations to converge compared to a system with fewer subdomains and greater overlap. To address these issues, several improvements to the alternating Schwarz method have been proposed (Gander, 2006; Dubois et al., 2012). We leave exploring these improvements to future work.

Mosaic Flow (Wang et al., 2022a) is a novel approach for solving BVPs on diverse domains with arbitrary boundary conditions. It consists of two primary components:

1. (1)

   The subdomain solver (SDNet) is a physics-informed neural PDE solver trained to solve a BVP on a small domain with arbitrary boundary conditions. Although this paper focuses on Dirichlet boundary conditions, SDNet can also be used with Neumann or Robin boundary conditions. SDNet’s effectiveness arises from its ability to rapidly generate predictions for any point within the domain. Note that while the boundary function input to SDNet is discretized, the $xy$-coordinate can be in continuous space. This is unlike finite difference or finite elements methods, which require discretizing the interior of the domain (Larrson and Thomée, 2003).

2. (2)

   The Mosaic Flow predictor (MFP) illustrated in Figure 2 is an iterative algorithm that leverages pre-trained SDNet’s inferences to solve BVPs on large domains—much larger than those that can be solved with a SDNet directly. The iterative algorithm decomposes the domain into smaller atomic subdomains, and updates the solution within each subdomain using SDNet’s predictions. It also ensures the spatial regularity of the solution along the subdomain boundaries using overlapping subdomains, inspired by the alternating Schwarz method. By utilizing SDNet as the sub-domain solver, MFP inherits its ability to make predictions for arbitrary points within the domain. This feature results in a significant performance advantage, as Mosaic Flow can compute the solution for only a small fraction of grid points, specifically the interfaces of the subdomains, as opposed to all grid points in the entire domain, as done in classical ASM.

Mosaic Flow combines the efficiency of SDNet on small domains with the scalability of MFP on larger domains, enabling the efficient solution of complex BVPs.

In general, our approach is agnostic of the choice of model for SDNet. For instance, one could use pure MLPs, a flavor of DeepONet (Lu et al., 2021a), or Fourier layers (Li et al., 2020a). The architecture of the neural solver used in this work, shown in Figure 3, is a variant of DeepONet that we inherit and improve. We first apply 1D convolutions to the input boundary conditions to create a high-dimensional embedding. The reason for using 1D convolutions is to take advantage of the inherent spatial structure of the boundary conditions, which can be seen as a 1D curve along the boundary of the domain. By convolving this signal, we capture local patterns and relevant features for predicting the solution within the domain. We anticipate this treatment of the boundaries will enhance convergence performance without affecting per-iteration performance, as convolutions are computationally efficient. We choose not to use Fourier layers due to the non-periodic nature of Dirichlet boundary conditions (Li et al., 2020a). Although FNO can handle non-periodic boundaries, the combination of convolutions and fully-connected layers proves sufficient for capturing global phenomena.

Next, we apply the input-split optimization discussed in Section 3.2 to the high-dimensional boundary embedding coupled with the input ${\bm{x}}$. The rest of the architecture is composed of a stack of linear layers, each followed by a nonlinear activation function. We use the GELU activation function (Hendrycks and Gimpel, 2016) because PINN training tends to have better convergence properties when using a smooth activation function (Shin et al., 2020).

A common input embedding in physics-informed neural PDE solvers is to concatenate the spatial coordinates ${\bm{x}}$ with the discretized boundary function $\hat{{\bm{g}}}$ into a single input vector (Lu et al., 2021a; Wang et al., 2022a). However, this input-concat approach is highly inefficient. For example, in a square 2D domain discretized into an $N\times N$ resolution, inferring the solution, $u({\bm{x}})$ at a single point ${\bm{x}}$ in the domain requires an input vector of dimension $4N+2$. The discretized boundary function $\hat{{\bm{g}}}$ is a vector of dimension $4N$ and the additional $2$ dimensions are for the $xy$-coordinates.

When inferring the solution of $q$ points in a batch, the input becomes a $q\times(4N+2)$ matrix ${\bm{\mathbf{{I}}}}$:

where the matrix ${\bm{\mathbf{{G}}}}\in\mathbb{R}^{q\times 4N}$ is formed by replicating the vector $\hat{{\bm{g}}}$ for each point in the batch and the rows of ${\bm{\mathbf{{X}}}}\in\mathbb{R}^{q\times 2}$ are the coordinates of the $q$ points. Denoting the weights of the first layer of the neural network as ${\bm{\mathbf{{W}}}}\in\mathbb{R}^{d\times(4N+2)}$, the output of this layer is given by the matrix multiplication of ${\bm{\mathbf{{I}}}}$ with ${\bm{\mathbf{{W}}}}$, followed by the application of the activation function $\phi$. Mathematically, we can express this as:

To reduce the computational cost and remove the redundancy in ${\bm{\mathbf{{G}}}}$ introduced by the input-concat approach, we split weight matrix ${\bm{\mathbf{{W}}}}\in\mathbb{R}^{d\times(4N+2)}$ into two column blocks, denoted as ${\bm{\mathbf{{W}}}}_{1}\in\mathbb{R}^{d\times 4N}$ and ${\bm{\mathbf{{W}}}}_{2}\in\mathbb{R}^{d\times 2}$. Using eq. 5, we can rewrite eq. 6 to arrive at our optimized approach:

where $\oplus$ is a broadcasted sum along the second axis of ${\bm{X}}{\bm{W}}_{2}^{T}$. Note that the discretized boundary condition $\hat{{\bm{g}}}$ is no longer replicated for each point in the input, but is computed only once and broadcasted along the batch dimension. This reduces the overall number of computations required by the network. With eq. 6, the cost of the first layer is $O(qNd)$. In comparison, with our optimized input-split approach in eq. 8, the cost is reduced to $O(Nd+qd)$. More importantly, this reduces the memory requirement of input tensor from $q(4N+2)$ words to $4N+2q$ words; when $q$ and $N$ are large, this saving can be substantial. The reduction in memory usage achieved by the optimized approach makes it possible to scale training to larger batch sizes.

After optimizing the network architecture, we accelerate the training of neural PDE solvers with a physics-informed loss using data parallelism. Recall that when training a physics-informed model, we use a loss function with multiple terms: $\mathcal{L}(\theta)=\mathcal{L}_{data}(\theta)+\mathcal{L}_{pde}(\theta)$ where $\mathcal{L}_{data}(\theta)$ represents the data loss function. The loss function that enforces the PDE constraints, $\mathcal{L}_{pde}(\theta)$ requires the computation of higher-order derivatives with respect to the model inputs. In the case of the Laplace equation, this involves calculating the second derivatives $\mathcal{N}_{xx}$ and $\mathcal{N}_{yy}$. This results in a large autograd graph that consumes significant device memory. The size of this autograd graph limits the batch size on a single GPU, motivating the need to scale up to multiple GPUs to improve performance. It is worth noting that without the PDE loss, a purely data-driven model could be trained with a larger batch size on a single device. However, such a model may exhibit physical inaccuracies and require significantly larger dataset for training.

To efficiently train the network with multiple loss terms, we separate the data and collocation points into distinct forward passes. This approach simplifies the application of different losses to their respective coordinates, as the data loss can only be applied to points with known solutions. However, when using distributed data parallelism (DDP), it is important to preserve the standard semantics of stochastic gradient descent (SGD). In data-parallel training, the model is replicated across different compute nodes, and local gradients are computed on each process. To synchronize gradients, an allreduce (Chan et al., 2007) operation is commonly used, where the gradients are averaged across processes. To maintain the correct semantics of SGD, we must be mindful of when gradient synchronization occurs. If synchronization occurs after both forward passes, it will compute a sum of averages rather than a true global average. Although this approach may yield satisfactory results in practice, it does not offer the same convergence guarantees.

To maintain consistency with SGD and ensure reliable convergence, we propose the method outlined in Algorithm 1 for each training iteration. In step 1, the forward and backward passes are computed for the data points (lines 5-6) on each process without averaging gradients across processes. Then, in step 2, we apply the forward and backward passes for the collocation points (lines 8-9). The gradients for the collocation points are accumulated onto the gradients for the data points (line 9), and this sum is subsequently averaged across all processes using the allreduce operation in step 3. The averaged gradients are applied locally on each device, ensuring consistency. The proposed approach not only ensures accurate gradient accumulation but also offers the advantage of performing a single allreduce operation per training iteration, instead of two separate operations for data and collocation points. This optimization reduces communication overhead and enhances the scalability and efficiency of the training process.

The baseline MFP adopts a sequential approach to solve subdomains, which ensures that all predictions are based on updated boundary conditions. Empirical evidence suggests that relaxing this can have a negative effect on prediction accuracy. However, by observing Figure 2, it becomes evident that atomic subdomains within each iteration of the algorithm do not overlap. This creates an opportunity for concurrently predicting these non-overlapping subdomains. To leverage this, we implement a batching technique that combines the atomic subdomains as input for SDNet inference. This effectively increases GPU occupancy by exploiting device-level parallelism as demonstrated in Section 5.3.

The MFP takes the boundary conditions of subdomains as its input. During the development of the distributed algorithm, a key factor is both accurately and effectively managing updates to boundary conditions within overlapping regions. The boundary information of the subdomains can be organized into a Cartesian grid. In the example illustrated in Figure 2, the distance between neighboring grid points is $\frac{1}{2}m$, which is a tunable hyperparameter. Choosing a smaller distance allows for more subdomains to be placed in the domain, potentially resulting in more accurate results. However, this also leads to increased computation and communication costs, as shown in Section 4.3. In this study, we choose a value of $\frac{1}{2}m$ because it is the largest distance we can use without significantly sacrificing accuracy.

To design a parallel algorithm for the MFP, we first divide the global domain into a 2D grid, where each block is assigned to a processor. The processors are assigned to this 2D grid in a row-wise scan pattern. It is worth noting that a processor mapping strategy based on locality-preserving space-filling curves such as Morton order (Morton, 1966) or Z-order could provide better load balancing and reduced data movement (Pekkilä et al., 2022), although we leave this for future studies. We refer to the region owned by a processor as the processor subdomain. In order to iteratively approach the final solution using Schwarz methods, neighboring processor subdomains need to communicate boundary information. To enable this exchange of information, we give each processor additional halo boundary information from its neighboring processor subdomains. Figure 4 illustrates the distribution of processor subdomains and how the boundary information is exchanged between processors.

Our proposed distributed algorithm is outlined in algorithm 2, where $t$, $\epsilon$, $\bm{g}$, SDNet, and $\bm{n}$ respectively denote the maximum number of iterations, convergence threshold, global boundary condition, the pre-trained SDNet model, and the neighbors of the current processor. We use the hat notation to denote a local variable (for example, $\hat{\bm{g}}_{0}$ denotes the local part of $\bm{g}_{0}$). The algorithm can be conceptually divided into two phases. The first phase is the iteration loop from 2 to 9. In each iteration, the boundary information of the local atomic subdomains is input to the pre-trained SDNet, which predicts the values only along the center lines of these atomic subdomains (3). Since the center lines of one subdomain are the boundary of another, these predictions are subsequently input to SDNet for the next iteration. After the inference step, the boundary information in the region where processor subdomains overlaps is packed into a contiguous buffer and sent to the corresponding neighbors with communicate_new_boundaries in 4. Figure 4 provides a graphical illustration of which parts of the boundaries are being communicated. The second phase starts after $t$ iterations or upon reaching the convergence threshold. In this phase, each processor uses the most recent atomic subdomain boundaries as input for SDNet, to predict the values at every grid point within each atomic subdomain, forming the local solution $\hat{{\bm{\mathbf{{S}}}}}$ (10). Finally, an all_gather is performed to collect the distributed subdomains. In the region where processor subdomains overlap, the final solution is obtained by computing the average of the predictions.

To design a scalable algorithm, we partially relax the order of subdomain updates in the baseline algorithm to balance accuracy requirements with communication efficiency. In the algorithm illustrated in Figure 2, as subdomains are solved by SDNet, the update to the boundary information inside the subdomain is applied immediately. However, when processors solve subdomains on the overlapped region, the update to the boundary information cannot be reflected in the neighboring processor until the communication step. In our parallel algorithm, we relax this principle by communicating only once per iteration. This relaxation in synchronization not only reduces the communication frequency but also makes the communication pattern agnostic to subdomain placement schemes. It is worth highlighting that the baseline principle of immediate updates to boundary information still holds within individual processor subdomains. This relaxation is similar to the algorithm proposed by Lions (Lions et al., 1988), which was proven to converge to the global solution. Empirical results in Section 5.3 show that this modification does not prevent the distributed MFP from finding accurate solutions.

We now analyze the costs associated with the distributed MF Predictor. Suppose the global domain has a resolution of $N\times N$ and is distributed across $P$ processors arranged in a $\sqrt{P}\times\sqrt{P}$ grid. The resolution of each subdomain is $m\times m$. As a result, each processor is assigned a subdomain consisting of $\frac{N}{m\sqrt{P}}\times\frac{N}{m\sqrt{P}}$ non-overlapping subdomains. Assuming that the subdomain boundaries form a Cartesian grid with interval $\frac{m}{d}$, and the overlapping region is $\frac{2(d-1)}{d}m$ wide along each subdomain boundary, there are $\frac{(dN)^{2}}{m^{2}P}$ subdomains in each processor with all eight neighbors. Using the alpha-beta model and removing the trailing terms, the communication cost of each processor is $C_{comm}=8I\alpha+\frac{1}{\beta}I(16\frac{Nd}{\sqrt{P}})$, where $\alpha$ models the network latency and $\beta$ models the network bandwidth.

Since communication is performed in every iteration, both the bandwidth and latency cost scale linearly with the iteration count. As the communication is limited to each processor and its immediate neighbors, the latency cost is not influenced by $P$ or $N$. Bandwidth cost increases linearly with $\frac{N}{\sqrt{P}}$, which is the length of one side of each subdomain, and $d$, which controls how dense the subdomains are placed. Denoting the computation cost of making 1 SDNet inference as $c$, we can express the computation cost of each processor as $C_{comp}=c\frac{(dN)^{2}}{m^{2}P}$. From this, we expect the computation cost to scale linearly with the number of processors. Note that we assume communication can be carried out simultaneously with all neighbors in our analysis, which may not always hold in practice.

We generate two distinct datasets: one for training SDNet and another for evaluating the MFP. The training dataset consists of small domains of fixed size, while the test dataset includes larger domains of arbitrary sizes. To construct these datasets, we generate boundary conditions using Gaussian processes and follow a similar approach to the original Mosaic Flow paper (Wang et al., 2022a). First, we use a Sobol Sequence (Sobol, 1998) to sample the hyperparameters of an infinitely differentiable Gaussian kernel of a 1-dimensional Gaussian process. Then, from each Gaussian process, we draw a sample function (i.e., a 1-D curve). This function serves as the discretized boundary function $\hat{{\bm{g}}}$ described in Section 2. Each boundary value problem for the Laplace equation is solved using pyAMG. (Bell et al., 2022).

We implement the distributed MFP in Python. For GPU-to-GPU communication, we use mpi4py (Dalcin and Fang, 2021), which is built with a CUDA-aware MPI library to enhance communication performance. To generate boundary conditions and ground truth solutions of the Laplace equation on larger domains, we use the method described in Section 5.1. We evaluate both batched inference for device-level parallelism and distributed inference for node-level parallelism.