Paper: Cooperative Minibatching in Graph Neural Networks (ICLR-2024)

Reviewer comment (Writing):
The writing is informal and not well organized, which is reflected but not limited to the following aspects: 1) ambiguous terminology and inconsistent word choice, e.g., "work" (not clear what this really means); 2) overly extensive background (not sure how most content in the background is related to the target problem); 3) causal usage of theorem.

Author response:
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.

Reviewer comment (Theory):
The theoretical result states that the expected work increases as batch size increases, which clearly depends on the setting. For example, what if, in the extreme case, the communication speed is infinite? What if I am using different GNN models? The theoretical result does not provide any specification with this regard and it is not clear how one can prove such a result. In addition, it is not clear how the theoretical results motivate or connect with the proposed approach.

Author response:
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.

Reviewer comment (Novelty):
No discussion on related works, despite there exists an extensive line of research from both algorithm and system communities that try to optimize distributed/large-scale GNN training. For example, I found the idea is closely related to layer sampling that is commonly used in GNN, and the communication optimization proposed in Gandhi, Swapnil, and Anand Padmanabha Iyer. "P3: Distributed deep graph learning at scale." 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI), 2021.

Author response:
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.

Reviewer comment (Experiments):
The paper only evaluates the proposed method on one GNN model. How generalizable is this method to other types of GNNs, especially those deep GNNs with tens of layers?

Author response:
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).

Reviewer comment (Reproducibility):
Could you clarify the unit of measurement for the cache size? Is it quantified in bytes or in terms of the number of vertices/features? What are the key factors for determining the optimal cache size for different datasets, GNN models, and the hardware platforms?

Author response:
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.

Reviewer comment (Evaluation):
What are the communication cost with and without applying the proposed technique?

Author response:
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.

Reviewer comment (Novelty):
The work in [1] also solves the redundancy issue across the PEs using cooperative training. However, this related work is not cited or compared in this paper.

Author response:
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.

Reviewer comment (Presentation):
It might be beneficial to include more illustrative figures, especially, for the algorithm 1. This would help readers to follow the steps of the proposed method more easily.

Author response:
We are working on more illustrative figures for the camera-ready revision of our paper.  
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.

Reviewer comment (Evaluation):
The paper seems to be missing a discussion of the communication costs of its method, which seem like they could be significant, especially in a multi-node setup. This is in contrast to independent mini-batching, which has the nice property of avoiding communication. The appendix (A3) notes communication overhead is not an issue, but this could be measured in detail, and in any case only considers a small-scale, on-node scenario.

Author response:
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.

Reviewer comment (Reproducibility):
How is the 1D partitioning done, in detail? Appendix A3 notes that using METIS was beneficial, so why not always use this (or a similar graph partitioning algorithm)?

Author response:
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.

Reviewer comment (Experiments):
The paper does not discuss the memory overhead of cooperative mini-batching. It seems to me that all samples in a K-hop neighborhood of each vertex need to be present for the method, which seems like it would result in significant memory overheads, especially as the number of layers increases (the paper only considers a network with three layers).

Author response:
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.

Reviewer comment (Experiments):
The paper seems to be only considering neighborhood sampling as a way of performing independent mini-batching. However, there are other ways to do this, e.g., graph cut algorithms. (See, for example, Rozemberczki et al., "Little Ball of Fur: A Python Library for Graph Sampling", CIKM 2020.)

Author response:
Weakness 4: We have tried to incorporate the well-established GNN sampling methods available in the literature.
NS is the first work proposing sampling to train GNNs. Random walks were also proposed to sample from l-hop (l > 1)
neighborhood instead of the immediate 1-hop neighborhood as in NS. Finally, LABOR is a very recent GNN 
sampling method that is similar but ensures more overlapping vertices compared to NS. Sampling in the general
graph setting may be richer, however, there are complications for GNNs as one usually wants the sampling
method to approximate training without sampling in an unbiased manner. Thus, parts of the richer
graph-sampling literature may not be applicable to GNNs and our goal is not to explore new ways to perform sampling for GNNs.

Reviewer comment (Experiments):
What are the communication costs of cooperative mini-batching? How does the method perform in cross-node scenarios?

Author response:
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.

Reviewer comment (Experiments):
How is partitioning done? Why not always use a method that reduces cross-device edges?

Author response:
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.