Optimizing Task Placement and Online Scheduling for Distributed GNN Training Acceleration ^††thanks: This work was supported in part by grants from Hong Kong RGC under the contracts HKU 17204619, 17208920 and 17207621.

GNNs learn effective graph embeddings by iteratively aggregating neighborhood features (Fig. 1) [23][24]. The derived embeddings can be further processed (e.g., using DNN layer, softmax operation), to produce decisions for downstream tasks (e.g., node classification, link prediction).

To construct a mini-batch for GNN training, a set of training nodes are sampled from the input graph, and their $L$-hop neighbors are used for embedding generation by a $L$-layer GNN. Using features of all $L$-hop neighbors of the selected training nodes may lead to GPU/CPU memory overflow or high computation complexity. A common practice is to recursively sample neighbors of each training node with a sampling algorithm (e.g., [13][14]), and a sub-graph is formed among the training node and its sampled $L$-hop neighbors. Each sub-graph with its features renders one sample in the mini-batch.

Using mini-batches of graph samples, GNN training is similar to DNN training: forward propagation is carried out to compute a loss, and then backward propagation to derive gradients of the GNN model parameters based on the loss, using an optimization algorithm (e.g., stochastic gradient descent); a gradient update operation follows, which involves gradient aggregation among workers in distributed training and application of updated parameters to the GNN model.

Deep Graph Library (DGL) [22] is a package built for easy implementation of GNN models on top of DL frameworks (e.g., PyTorch [25], MXNet [26]). The recent release of DGL supports distributed GNN training on relatively large graphs. It uses random sampling, collocates one worker with one graph store, and does not pipeline GNN training across iterations, leaving a large room for further performance improvement. Euler [27] is integrated with TensorFlow [28] for GNN training, which partitions a large graph in a round-robin manner and splits feature retrieving requests to allow concurrent transmissions; large data transfers still exist due to its locality-oblivious graph partition. AliGraph [29] adopts distributed graph storage, optimized sampling operators and runtime to efficiently support GNNs. PyTorch Geometric [30] is a deep learning library on irregularly structured input data such as graphs, supporting multi-GPU training on a single machine only. Dorylus [16] distributes GNN training over serverless cloud function threads on CPU servers, requiring specialized functions provided by AWS [31]. Large data traffic exists in these systems, and careful transfer scheduling and task deployment can enhance them for training time minimization.

NeuGraph [32] and PaGraph [18], which train GNN models on a single machine, adopt full-graph training by loading entire graphs into GPU memory, and are hence only feasible for training over small graphs. Considering multi-server clusters, ROC [33] splits the input graph over multiple GPUs or machines to achieve workload balance, and adopts a memory management scheme to reduce CPU-GPU data transfer. DistDGL [15] alleviates network transfer in distributed GNN training by co-locating each worker with its samplers on the same server, and partitioning the input graph with a minimum edge cut method. Further, various graph partition, sampling and caching methods have been proposed for enhancing distributed GNN training [2][34][35]. These studies focus on minimizing data transfer volumes across devices/machines. Optimization of task placement and execution scheduling is orthogonal to the existing efforts, and our solution can complement them to fully accelerate distributed GNN training. DGCL [36] is a recently proposed communication library for distributed GNN training, which decides data routing strategy for every graph node to the requiring worker(s), considering the detailed interconnection topology among workers. Its detailed communication plan is re-computed before every training epoch, which may incur substantial overhead for large graphs. Our design performs efficient, polynomial-time online scheduling on both task execution and data flows between tasks, effectively reducing the overall training time.

Task placement, computation and communication scheduling have been studied for DNN training on non-graph data [37][38][39]. The communication scheduling deals with arranging transmission time and order of gradient/parameter tensors for parameter synchronization [38][39]. Placement studies focus on worker placement to minimize interference [40] instead of proximity to data, and DNN operator placement to achieve model parallelism [41]. Computation scheduling deals with fine-grained operator execution ordering, in case of model- or pipeline-parallel DNN training [42][43][44]. Compared to distributed DNN training, GNNs are largely trained with data parallelism, incurring large graph data communication that blocks the computation and occupies a majority of the training time (up to 80% [17]). Instead of operator-level placement and scheduling of a GNN model, we study placement of tasks (samplers, workers and parameter servers), overlap both graph data transfer and tensor communication with computation (the graph data traffic is magnitudes larger than tensor transfers), and pipeline mini-batch training across training iterations, which are all dedicated for GNN training acceleration.

We train a GNN model (with $L$ embedding layers) in a cluster of $M$ physical machines. Partitions of a large graph used for GNN training are stored on the $M$ machines. Each machine $m\in[M]$¹¹1$[X]$ denotes set $\{1,2,\ldots,X\}$ is equipped with $R$ types of computational resources (e.g., GPU, CPU and memory), with type-$r$ resource available at the amount of $C^{r}_{m}$. Let $B^{m}_{in}$ ($B^{m}_{out}$) represent the available incoming (outgoing) NIC bandwidth on machine $m$.

There are four types of tasks in our GNN training job: (1) Graph store server: Each machine hosts a graph store server, to maintain one graph partition (including graph structure and node/edge features). (2) Sampler: Each sampler selects training nodes, retrieves sampled node/edge features from graph store servers and forms sub-graphs. (3) Worker: Each worker carries out forward and backward computation, pushes gradients to and pulls parameters from parameter servers for parameter synchronization. A worker is typically associated with one or multiple samplers, which supply mini-batches dedicatedly to the worker. (4) Parameter server (PS): PSs aggregate gradients from all workers, update the GNN model parameters and distribute updated parameters to all workers.

We use $J_{g},J_{s},J_{w}$ and $J_{ps}$ to represent the sets of graph store servers, samplers, workers and PSs, respectively, in the training job. We suppose the number of each type of tasks is specified by the ML developer: the number of graph stores is $M$ (as each machine hosts exactly one graph partition), the number of workers can be larger or smaller than $M$ (considering a machine may host multiple GPUs and CPUs, and a worker typically consumes one GPU or CPU), the number of samplers to serve each worker is usually fixed (e.g., 2 samplers per worker). Let $J=J_{g}\cup J_{s}\cup J_{w}\cup J_{ps}$ denote the set of all tasks. Each task $j\in J$ occupies a $w_{j}^{r}$ amount of type-$r$ resource, $\forall r\in[R]$. For example, graph store servers, samplers and PSs are commonly run on CPUs, while workers can run on GPUs [18] or CPUs [15], and consume the respective memory. Tasks of the same kind (e.g., all samplers) occupy the same amount of resources. Let $p_{j}$ denote the execution time of task $j$ in each iteration.

In a training iteration, each sampler selects a number of training nodes from the input graph and signals the graph store servers to acquire neighbor information. Upon requests from a sampler, a graph store server samples among $L$-hop neighbors of the training nodes that it hosts (using a given sampling algorithm), and sends the node/edge features back to the sampler. The sampler then sends sub-graph samples to its associated worker, which form a mini-batch from samples supplied by its sampler(s), for forward and backward computation. Computed gradients are sent from workers to the PSs and then updated parameters are dispatched from PSs to workers. The workflow is illustrated in Fig. 2.

We target overall training time minimization in our distributed GNN training job. Our design space includes two subproblems.

Given placements $\{y_{j}^{m}\}$, we design an online algorithm that decides start time of each task ($x^{t}_{j,n}$) and flow transmission ($k_{(j,n)\rightarrow(j^{\prime},n^{\prime})}^{t}$) over time.

We maintain two flow sets: (i) $F_{act}$, that stores every active flow $(j,n)\rightarrow(j^{\prime},n^{\prime})$ which currently has started but not finished transmission yet; (ii) $F_{pend}$, to store every pending flow $(j,n)\rightarrow(j^{\prime},n^{\prime})$ whose predecessor task $(j,n)$ has been done, and that has not started because its predecessor flow $(j,n-1)\rightarrow(j^{\prime},n^{\prime}-1)$ in the previous iteration has not completed transmission yet. For each task $(j,n)$, we use $\mathcal{F}(j,n)$ to represent the set of flows that originate from $(j,n)$ to tasks that reside on other machines (than where $j$ is).

Our online scheduling algorithm is in Alg. 1. We start by running graph store server processing for the first training iteration at $t=1$ (line 2). Then at each time $t$, we run every task that has received all required data and hence is available to execute (line 7). For each task $(j,n)$ completed at $t-1$, consider every flow $(j,n)\rightarrow(j^{\prime},n^{\prime})\in\mathcal{F}(j,n)$ in $t$: if the flow’s predecessor flow $(j,n-1)\rightarrow(j^{\prime},n^{\prime}-1)$ is in $F_{act}$ or $F_{pend}$ (indicating it not done yet), we add $(j,n)\rightarrow(j^{\prime},n^{\prime})$ to $F_{pend}$; otherwise, it is scheduled to transmit in $t$ and added to $F_{act}$ (lines 8-13). In addition, for every flow $(j,n)\rightarrow(j^{\prime},n^{\prime})$ ended at $t-1$, we check if its successor flow $(j,n+1)\rightarrow(j^{\prime},n^{\prime}+1)$ is in $F_{pend}$: if so, we move it from $F_{pend}$ to $F_{act}$ and start the flow transmission (lines 14-17). For every flow $(j,n)\rightarrow(j^{\prime},n^{\prime})$ which transfers in $t$, supposing $j$ placed on $m$ and $j^{\prime}$ on $m^{\prime}$, we set its traffic volume $k_{(j,n)\rightarrow(j^{\prime},n^{\prime})}^{t}$ at $t$ to $\min\{B^{m^{\prime}}_{in}/\Delta_{in}^{m^{\prime}},B^{m}_{out}/\Delta_{out}^{m}\}$ (lines 18-21). $\Delta_{in}^{m}$ ($\Delta_{out}^{m}$) is the ingress flow degree (egress flow degree) on machine $m$, counting the number of active flows entering and exiting from $m$, respectively:

Let $F_{one\_iter}$ denote the set of all inter-machine flows in one training iteration, including the transfer of updated parameters computed in this iteration from PS to workers. We define the one-iteration ingress flow degree $\widehat{\Delta^{m}_{in}}$ and one-iteration egress flow degree $\widehat{\Delta^{m}_{out}}$:

The detailed proof is given in Appendix A.

The detail proof is given in Appendix B.

A feasible task placement solution should respect resource capacity constraints in (2). We first randomly order the $M$ machines into $\{m_{1},m_{2},\ldots,m_{M}\}$. Note that placements of graph store servers are given (one on a machine). Let $[q_{s},q_{w},q_{ps},i]$ indicate that we can pack $q_{s}$ samplers, $q_{w}$ workers and $q_{ps}$ PSs within the first $i$ machine ($m_{1}$ to $m_{i}$) without violating resource capacities, and $(q_{s},q_{w},q_{ps},i)$ be a particular partial placement of putting $q_{s}$ samplers, $q_{w}$ workers and $q_{ps}$ PSs on machine $m_{i}$. We use $\mathcal{A}_{q_{s},q_{w},q_{ps},i}$ to denote an exact placement associated with $[q_{s},q_{w},q_{ps},i]$, specifying how many samplers, workers and PSs are placed in each of the $i$ machines, to make up for the total numbers of $q_{s}$, $q_{w}$ and $q_{ps}$. Let $\Omega(i)$ be the set of all $[q_{s},q_{w},q_{ps},i]$’s with $i$ fixed and $q_{s}\in[|J_{s}|],q_{w}\in[|J_{w}|],q_{ps}\in[|J_{ps}|]$.

We use dynamic programming to construct a feasible placement solution. We first consider all feasible placements $(q_{s},q_{w},q_{ps},1)$ on $m_{1}$. Let $\eta_{s}$ denote the maximal number of samplers that can be hosted by any machine, i.e. $\eta_{s}=\max_{m\in[M]}\min_{r\in[R]:w^{r}_{j}>0}\lfloor C^{\prime r}_{m}/w^{r}_{j}\rfloor,$ any $j\in J_{s}$ ($C^{\prime r}_{m}$ is available type-$r$ resource on $m$ excluding that occupied by the graph store server). Similarly, we can define an upper bound on the number of workers and PSs per machine, $\eta_{w}$ and $\eta_{ps}$. For every possible combination of $q_{s}\in\{0\}\cup[\min\{|J_{s}|,\eta_{s}\}]$, $q_{w}\in\{0\}\cup[\min\{|J_{w}|,\eta_{w}\}]$ and $q_{ps}\in\{0\}\cup[\min\{|J_{ps}|,\eta_{ps}\}]$, we check if the capacity constraints on $m_{1}$ are satisfied. For every feasible solution found, we add $[q_{s},q_{w},q_{ps},1]$ to $\Omega(1)$, and set $\mathcal{A}_{q_{s},q_{w},q_{ps},1}=\{(q_{s},q_{w},q_{ps},1)\}$.

Next, we iteratively construct $\Omega(i)$ based on $\Omega(i-1)$ until finding a complete feasible solution of placing all $|J_{s}|$ samplers, $|J_{w}|$ workers and $|J_{ps}|$ PSs onto the machines. For each $[q_{s},q_{w},q_{ps},i-1]\in\Omega(i-1)$, we examine whether $|J_{s}|-q_{s}$ samplers, $|J_{w}|-q_{w}$ workers and $|J_{ps}|-q_{ps}$ PSs can be fit into machine $m_{i}$. If so, we have identified a complete feasible placement solution that packs all tasks within the first $i$ machines: $\mathcal{A}_{solution}=\mathcal{A}(q_{s},q_{w},q_{ps},i-1)\cup\{(|J_{s}|-q_{s},|J_{w}|-q_{w},|J_{ps}|-q_{ps},i)\}$. Otherwise, we find every feasible placement $(q^{\prime}_{s},q^{\prime}_{w},q^{\prime}_{ps},i)$ with $q^{\prime}_{s}\in\{0\}\cup[|J_{s}|-q_{s}]$, $q^{\prime}_{w}\in\{0\}\cup[|J_{w}|-q_{w}]$, and $q^{\prime}_{ps}\in\{0\}\cup[|J_{ps}|-q_{ps}]$ that satisfies capacity constraint on machine $m_{i}$; and if $[q_{s}+q^{\prime}_{s},q_{w}+q^{\prime}_{w},q_{ps}+q^{\prime}_{ps},i]$ is not in $\Omega(i)$ yet, we add it into $\Omega(i)$, and set $A(q_{s}+q^{\prime}_{s},q_{w}+q^{\prime}_{w},q_{ps}+q^{\prime}_{ps},i)$ to be the union of $A(q_{s},q_{w},q_{ps},i-1)$ and $\{(q^{\prime}_{s},q^{\prime}_{w},q^{\prime}_{ps},i)\}$. We build from $\Omega(2)$ to $\Omega(M)$ and return the first complete feasible placement solution. We summarize our initial placement solution construction algorithm in Alg. 2, referred to as IFS.

The detailed proof is given in Appendix C.

Starting from the initial feasible placement, we iteratively search for better placement solutions, according to a $cost(\mathcal{Y})$ defined on the overall training makespan of placement $\mathcal{Y}$.

Practically, task placements should be decided before training starts and remain fixed during training (to avoid substantial overhead of VM/container migration and flow redirection). The online nature of execution scheduling is due to size variation of sampled graph data; we should identify a placement that works best in expectation of the traffic variation. To this end, we profile task execution time and inter-task traffic volumes by running the GNN training for some iterations ($50$ as in our evaluation), and produce their distributions. We simulate the training process under placement $\mathcal{Y}$ using Alg. 1, with time and traffic volume drawn from the distributions, and derive the training makespan $T^{\prime}_{\mathcal{Y}}$. The cost of placement $\mathcal{Y}$ is:

Our search explores the solution space by transferring from one placement $\mathcal{Y}_{z}$ to another $\mathcal{Y}_{z+1}$, for a total of $I$ transfers (the time budget). Give $\mathcal{Y}_{z}$, we uniformly randomly sample a task $j\in J\setminus J_{g}$. Let $M_{avail}$ denote the set of machines other than the one where $j$ is placed in $\mathcal{Y}_{z}$, which can host $j$ adhering to relaxed resource capacity constraints: