Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs

We present the Sequential Aggregation and Rematerialization (SAR) scheme for\ndistributed full-batch training of Graph Neural Networks (GNNs) on large\ngraphs. Large-scale training of GNNs has recently been dominated by\nsampling-based methods and methods based on non-learnable message passing. SAR\non the other hand is a distributed technique that can train any GNN type\ndirectly on an entire large graph. The key innovation in SAR is the distributed\nsequential rematerialization scheme which sequentially re-constructs then frees\npieces of the prohibitively large GNN computational graph during the backward\npass. This results in excellent memory scaling behavior where the memory\nconsumption per worker goes down linearly with the number of workers, even for\ndensely connected graphs. Using SAR, we report the largest applications of\nfull-batch GNN training to-date, and demonstrate large memory savings as the\nnumber of workers increases. We also present a general technique based on\nkernel fusion and attention-matrix rematerialization to optimize both the\nruntime and memory efficiency of attention-based models. We show that, coupled\nwith SAR, our optimized attention kernels lead to significant speedups and\nmemory savings in attention-based GNNs.We made the SAR GNN training library\npublicy available: \\url{https://github.com/IntelLabs/SAR}.\n