Graph neural networks have become increasingly popular in recent years due to their ability to naturally encode relational input data and their ability to operate on large graphs by using a sparse representation of graph adjacency matrices. As we look to scale up these models using custom hardware, a natural assumption would be that we need hardware tailored to sparse operations and/or dynamic control flow. In this work, we question this assumption by scaling up sparse graph neural networks using a platform targeted at dense computation on fixed-size data. Drawing inspiration from optimization of numerical algorithms on sparse matrices, we develop techniques that enable training the sparse graph neural network model from Allamanis et al. (2018) in 13 minutes using a 512-core TPUv2 Pod, whereas the original training takes almost a day.
The authors present a framework to implement graph neural networks training
efficiently ---an inherently sparse task--- using "custom dense hardware", here
Tensor Processing Units (TPUs V2). The key steps are: (1) reordering the labels
to reduce the bandwidth of the adjacency matrix, (2) (Sometimes approximate)
Decomposition using block matrices for efficient storage and computations, and
(3) memory layout optimization. They evaluate their framework on the VarMisuse
dataset and compare the performance against a GPU implementation running on
Nvidia Tesla V100. In the best configuration, they were able to reach 78% test
accuracy in 13 minutes vs 19h for the baseline.
The paper is well structured and thorough. I was able to understand the
challenges and the solutions proposed, even without prior knowledge in the
architecture this work focuses on. I think it is sufficiently detailed to enable
an independent implementation without referring to the authors source code.
However, I feel that it might lack novelty. Indeed, as described in the Related
Work section, each component of the pipeline is well known and used very
frequently in the HPC community. Sometimes, knowing how arrange common
primitives is very powerful, and looking at the results of the experimental
section, this is enough to improve performance by almost two orders of
magnitude. However, I think it might be more due to some intrinsic properties of
the dataset than the method itself. As made clear by the title of section 5.1,
it is the data itself that has low bandwidth. If we consider a dataset that does
not satisfy this requirement, stage (1) has no effect, and it is impossible to
perform the decomposition done in (2). In that situation, the contributions of
this paper would be nullified.  It would be perfectly reasonable to think that
most datasets have low bandwidth, but this not a claim that made the authors.
This work would be considerably more impactful if it measured the bandwidth of
more well recognized datasets. My small exposition to these problem does not
allow me to make any meaningful suggestion.
Finally, I am puzzled by the "dense hardware"/GPU distinction. From my
experience, GPU devices *are* designed for and extremely efficient at dense linear
algebra. Sparse operations are historically performed on CPUs. While they are
possible on GPUs, they are usually much slower. For example CUSparse, the sparse
matrices library, part of NVidia CUDA toolkit, was only introduced in its
version 4. It's a clear indication that sparse operations are not a strength of
GPUs. According to my experience writing GPU code, I feel that this approach
would actually perform extremely well on GPUs as it does on TPUs. I think it is
thus important to compare this framework on GPUs too. Since these
optimizations are not TPU specific and have not been applied in the GPU based
GNN libraries referenced in this paper reinforce my concerns that they are
problem-specific.
Even though the performance gains demonstrated are sizeable, the fact the
approach does not seem TPU specific and is potentially problem specific makes me
lean towards rejection.