Cooperative Minibatching in Graph Neural Networks

Thank you very much for your feedback. Let us clarify some of the weaknesses pointed out and reply to your questions below.

Weakness 1: We will work on increasing the formality of our paper to make it read better. We will also consider moving some of the background into the appendix to make more space for other content in the final revision. To reply to your question, the “work” refers to the “computation" that needs to be carried out for sequential execution, as it is usually referred to in computational complexity. Hence, in this context, for a given minibatch size, it is a quantity that monotonically increases as the number of sampled vertices and edges increases, and the details of this dependency are characterized by the underlying GNN model. Appendix A.4 and Table 5 go more into detail about what we mean by work. However, simply looking at how many vertices and edges will be sampled for different batch sizes is a good enough approximation to reason about the final runtime for any GNN model. We know that there will be redundant edges for independent minibatching meaning it will be slower in theory.

Weakness 2 and Questions: Appendix A.4 along with Table 5 characterizes the work more formally. There, we take into account the GNN model complexity and the different bandwidths available in a computer system (PCI-e, inter-GPU bandwidth, GPU memory bandwidth). To reply to your question, if communication speed (inter-GPU bandwidth) is infinite, then cooperative minibatching will be guaranteed to be faster than independent minibatching. Given $P$ PEs, the main complexity difference comes from the fact that Independent Minibatching has $\frac{B}{P}$ as the local batch size for each PE but Cooperative Minibatching has $B$ as the global batch size, which PEs cooperate to process together (note that the global batch size is $B$ for both methods). If $W$ denotes the work for a given minibatch size, we know that $P W(\frac{B}{P}) \geq W(B)$ due to the theorems we prove in our paper. The paragraph right after Table 3 in our paper already makes the connection between the empirical curves in Figure 2 and the observed runtime results in Tables 2 and 3. We will make the necessary modifications to our paper to make this more clear.

Furthermore, our proposed approaches are agnostic to the underlying GNN model, as only the layer outputs in between the layers need to be communicated when using Cooperative Minibatching, because every GNN model has vertex embeddings as the input to the layer and output of the layer. The GCN and R-GCN models we use are the computationally least expensive GNN models for homogenous and heterogenous graphs. If we were to use GAT or R-GAT models in our experimental evaluation, because these models are computationally slower, the communication runtime would be even less significant increasing the advantage of Cooperative Minibatching compared to Independent.

Weakness 3: Note that we already cite [1] in our paper. There, they propose to fetch different channels of vertex features from other PEs (We always partition across vertices), while duplicating the mini-batch graph structure on each processing element for the very first layer (or last layer with the sampling-focused notation in our paper), performing GNN aggregation operation in an intra-layer model parallelism fashion for only a single layer.

Note that, the approach used in [1] is not as easily generalizable to different GNN models such as Graph Attention Networks, which they point out require special handling. On the contrary, our proposed Cooperative Minibatching approach is agnostic to the GNN model and requires no special handling. In the same first layer, [1] switches to data parallelism (Independent Minibatching) as can be seen in Figure 5 in [1]. In our work, we propose to process all stages of GNN training including graph sampling, feature loading, and all layers of forward/backward stages in a cooperative manner as the redundancy argument applies at all stages of GNN training. Moreover, it is possible to even combine the intra-layer parallelism approach in [1] with our proposed approach.

[1] Gandhi, Swapnil, and Anand Padmanabha Iyer. "P3: Distributed deep graph learning at scale." 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2021.

Thank you very much for your feedback. Let us reply to your questions below:

1. In our experiments, we used GCN and R-GCN, which are the simplest and fastest GNN models for homograph datasets (a single edge type), and heterographs (multiple edge types). Our proposed methods are agnostic to the underlying GNN model. The more complicated and higher complexity the GNN model is, the more benefit we will see from Cooperative Minibatching as the reduction in work due to redundant computations will outweigh the cost of the communication even further. Moreover, the more layers there are, the more overlap there will be between multiple PEs, meaning that Cooperative Minibatching will be even more beneficial. We chose 3 layers as a realistic number that people use in industrial deployments of GNNs. Our proposed methods will apply to any GNNs that utilize minibatch training with graph sampling (Cooperative and Dependent) or with full neighborhoods (Cooperative only).

2. As you have already pointed out, Figure 4b shows that the cache miss rate difference between 64 and 256 is almost indistinguishable. So, we could have very well chosen $\kappa=64$, and our runtime results would have ended up being the same. The lower $\kappa$ is, the less the chance that it will negatively affect convergence. We agree that using $\kappa=64$ is a good option as well. A good rule of thumb would be to use a $\kappa$ value that gives almost the same cache miss rate as $\kappa=\infty$, and the smallest such $\kappa$ should be chosen.

3. Table 1 reports the cache sizes, and its unit is in the number of vertices whose features are cached. Thus, the size of the cache would be cache-size times #feats times 4 bytes for the float32 datatype. The optimal cache size is the largest size you can pick without running out of memory on your system and PEs. We ran our code in different scenarios with different-sized datasets. For example, we could have cached all of the Reddit dataset but then we would not have been able to demonstrate the benefits of dependent minibatching on that dataset as the cache miss rate would have been 0% for all $\kappa$ values.

4. The communication costs are provided in Appendix A.4. In particular, Table 5 gives the complexities of different GNN stages taking into account the PCI-e bandwidth, the NVLink bandwidth and the GPU memory bandwidth along with the number of sampled edges and vertices in each layer while also taking into account the underlying GNN model computational complexity.

5. Thank you for pointing us to a concurrent work that we were not aware of. Indeed the split parallelism and cooperative training method proposed in [1] are very similar to our Cooperative Minibatching method. The differences we see are that while we propose cooperative minibatching for all stages of GNN training (sampling, feature loading, and forward/backward), [1] proposes it only for feature loading and forward/backward stages. Furthermore, they are not proving the theorems we prove in our paper and they are not proposing the dependent minibatching approach. We have cited this work as concurrent work in our rebuttal revision. However, the fact that similar approaches are showing up in the literature provides further proof that our proposed methods are credible and useful.

We started this work in June 2022, the first version of our code has been publicly available since August 2022 and our first submission of this paper was in January 2023, which we can prove after the double-blind review period.

6. We are working on more illustrative figures for the camera-ready revision of our paper.

[1]. GSplit: Scaling Graph Neural Network Training on Large Graphs via Split-Parallelism

We added Figure 6 to Appendix A.4 to illustrate the difference between Independent and Cooperative Minibatching methods and accompany algorithm 1 as per your suggestion. We will continue to work on reorganizing our paper for the camera-ready version to see if we can swap some of the more relevant content from the Appendix with content from the main paper according to your and the other reviewers' suggestions.

Thank you very much for your feedback on our paper. Let us first address some of the weaknesses you pointed out and then provide answers to your questions.

Weakness 1 and Question 1: Appendix A.4 and Table 5 go into more detail about the complexities of Independent vs Cooperative Minibatching. There, the complexity is characterized by the number of sampled edges and vertices in each layer (tied directly to the theorems and empirical figures in our paper). We also talk about when cooperative minibatching is a better choice than independent minibatching taking into account the available bandwidths of the computer system.

But to summarize, the important thing is the GNN model's computational complexity and relative speed of communication in the computer system. In our experiments, we use two different GNN models, GCN, and R-GCN, however, our approach works with any GNN model as it only requires routing the layer outputs in between each layer. GCN and R-GCN are computationally faster models compared to for example GAT and R-GAT models (Graph Attention Networks) due to their use of the computationally more expensive SDDMM kernels (GCN uses SPMM only). Thus, our choice of GCN and R-GCN models is the worst-case scenario for Cooperative Minibatching. In our experimental evaluation, we selected single-node multi-GPU systems as these systems have NVLink providing relatively fast communication, the bandwidths are provided in Table 2, $\gamma$ and $\alpha$ give us the memory bandwidth of the GPU and the communication bandwidth across GPUs respectively, which are used in Appendix A.4 to characterize the complexities of communication and computation of each GNN training stage.

New computer systems have increasingly more efficient communication. The new hardware on the market by NVIDIA provides NVLink connections that work in the multi-node setting for up to 256 GPUs [1] providing intra-node-like bandwidth across nodes. However, we do not have access to such systems. As our method strongly depends on the existence of a high bandwidth communication between PEs, it will not bring benefits in any arbitrary setting. In our case, NVLink bandwidth is $\alpha=600$GB/s vs $\gamma=2$TB/s of the GPU global memory bandwidth for the 8 GPU DGX-A100 system in Table 2. Note that multi-GPU systems by AMD and Intel have equivalent high-bandwidth intra-node GPU interconnects.

Weakness 2 and Question 2: We randomly permute the vertices of the original graph, then logically assign equal-sized contiguous ranges of vertices to each PE in our main experiments. Even with such naive partitioning, we show favorable runtime results. The problem with graph partitioning is that it is hard to ensure load balance for the GNN setting. Even if you partition the graph while load-balancing the number of vertices or edges in each partition, GNN computation specifically requires the L-hop neighborhoods of training vertices to be load-balanced.

Our results in Table 6 show the load balance issues with METIS partitioning, even though the runtime on the papers dataset decreases from 13ms to 12ms, we also see that the runtime goes from 183ms to 185ms on the mag dataset with partitioning. This is due to the fact that on papers, the faster GCN model is used, where the communication overhead is relatively higher, so reducing communication improves runtime. However, for the R-GCN model used for the mag dataset, we see a slowdown, because the computation runtime is dominant, and ensuring load balance takes priority over reducing communication. We note that our experimentation with graph partitioning is preliminary and follow-up research is required to ensure that Cooperative Minibatching can be made even faster than it already is, which we will leave as future work. Our goal in this paper is to show that Cooperative Minibatching is a viable alternative to Independent Minibatching and to motivate further research on it.

Weakness 3 and Question 3: The only memory overhead of cooperative minibatching compared to independent minibatching is the allocated buffers for the all-to-all communication calls. Otherwise, because there are no duplicate vertices or edges on the PEs in Cooperative Minibatching, Independent Minibatching has higher memory overhead, which is exacerbated when the number of GPUs is increased. Moreover, as the number of layers is increased, the overlap between the 4-hop neighborhoods will be even greater than the overlap for 3-hop neighborhoods. Thus, if there were 4 layers, the 4-layer equivalent of Table 2 would show even more favorable results for Cooperative Minibatching. We do not readily have access to the computer systems we used in our experiments so we can not easily provide you with the empirical verification.

[1] https://www.nvidia.com/en-us/data-center/nvlink/