Distributed Training of Large Graph Neural Networks with Variable Communication Rates

Juan Cervino, Md Asadullah Turja, Hesham Mostafa, Nageen Himayat, Alejandro Ribeiro

Keywords: Graph Neural Networks, Distributed Training, Variable Compression, Large Scale Graphs

TL;DR: We train GNNs on large scale graphs in a distributed way using variable compression for the communications between machines.

Abstract:

Training Graph Neural Networks (GNNs) on large graphs presents unique challenges due to the large memory and compute requirements.
Distributed GNN training, where the graph is partitioned across multiple machines, is a common approach to train GNNs on large graphs.
However, as the graph cannot generally be decomposed into small non-interacting components, data communication between the training machines quickly limits training speeds.
Compressing the communicated node activations by a fixed amount improves the training speeds, but lowers the accuracy of the trained GNN.
In this paper, we introduce a variable compression scheme for reducing the communication volume in distributed GNN training without compromising the accuracy of the learned model.
Based on our theoretical analysis, we derive a variable compression method that converges to a solution that is equivalent to the full communication case.
Our empirical results show that our method attains a comparable performance to the one obtained with full communication and that for any communication budget, we outperform full communication and any fixed compression ratio.

Primary Area: general machine learning (i.e., none of the above)

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 5873