Armada: Memory-Efficient Distributed Training of Large-Scale Graph Neural Networks

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Waleffe, Roger, Sarda, Devesh, Mohoney, Jason, Vlatakis-Gkaragkounis, Emmanouil-Vasileios, Rekatsinas, Theodoros, Venkataraman, Shivaram
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866913706592436224
author Waleffe, Roger
Sarda, Devesh
Mohoney, Jason
Vlatakis-Gkaragkounis, Emmanouil-Vasileios
Rekatsinas, Theodoros
Venkataraman, Shivaram
author_facet Waleffe, Roger
Sarda, Devesh
Mohoney, Jason
Vlatakis-Gkaragkounis, Emmanouil-Vasileios
Rekatsinas, Theodoros
Venkataraman, Shivaram
contents We study distributed training of Graph Neural Networks (GNNs) on billion-scale graphs that are partitioned across machines. Efficient training in this setting relies on min-edge-cut partitioning algorithms, which minimize cross-machine communication due to GNN neighborhood sampling. Yet, min-edge-cut partitioning over large graphs remains a challenge: State-of-the-art (SoTA) offline methods (e.g., METIS) are effective, but they require orders of magnitude more memory and runtime than GNN training itself, while computationally efficient algorithms (e.g., streaming greedy approaches) suffer from increased edge cuts. Thus, in this work we introduce Armada, a new end-to-end system for distributed GNN training whose key contribution is GREM, a novel min-edge-cut partitioning algorithm that can efficiently scale to large graphs. GREM builds on streaming greedy approaches with one key addition: prior vertex assignments are continuously refined during streaming, rather than frozen after an initial greedy selection. Our theoretical analysis and experimental results show that this refinement is critical to minimizing edge cuts and enables GREM to reach partition quality comparable to METIS but with 8-65x less memory and 8-46x faster. Given a partitioned graph, Armada leverages a new disaggregated architecture for distributed GNN training to further improve efficiency; we find that on common cloud machines, even with zero communication, GNN neighborhood sampling and feature loading bottleneck training. Disaggregation allows Armada to independently allocate resources for these operations and ensure that expensive GPUs remain saturated with computation. We evaluate Armada against SoTA systems for distributed GNN training and find that the disaggregated architecture leads to runtime improvements up to 4.5x and cost reductions up to 3.1x.
format Preprint
id arxiv_https___arxiv_org_abs_2502_17846
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Armada: Memory-Efficient Distributed Training of Large-Scale Graph Neural Networks
Waleffe, Roger
Sarda, Devesh
Mohoney, Jason
Vlatakis-Gkaragkounis, Emmanouil-Vasileios
Rekatsinas, Theodoros
Venkataraman, Shivaram
Machine Learning
Distributed, Parallel, and Cluster Computing
We study distributed training of Graph Neural Networks (GNNs) on billion-scale graphs that are partitioned across machines. Efficient training in this setting relies on min-edge-cut partitioning algorithms, which minimize cross-machine communication due to GNN neighborhood sampling. Yet, min-edge-cut partitioning over large graphs remains a challenge: State-of-the-art (SoTA) offline methods (e.g., METIS) are effective, but they require orders of magnitude more memory and runtime than GNN training itself, while computationally efficient algorithms (e.g., streaming greedy approaches) suffer from increased edge cuts. Thus, in this work we introduce Armada, a new end-to-end system for distributed GNN training whose key contribution is GREM, a novel min-edge-cut partitioning algorithm that can efficiently scale to large graphs. GREM builds on streaming greedy approaches with one key addition: prior vertex assignments are continuously refined during streaming, rather than frozen after an initial greedy selection. Our theoretical analysis and experimental results show that this refinement is critical to minimizing edge cuts and enables GREM to reach partition quality comparable to METIS but with 8-65x less memory and 8-46x faster. Given a partitioned graph, Armada leverages a new disaggregated architecture for distributed GNN training to further improve efficiency; we find that on common cloud machines, even with zero communication, GNN neighborhood sampling and feature loading bottleneck training. Disaggregation allows Armada to independently allocate resources for these operations and ensure that expensive GPUs remain saturated with computation. We evaluate Armada against SoTA systems for distributed GNN training and find that the disaggregated architecture leads to runtime improvements up to 4.5x and cost reductions up to 3.1x.
title Armada: Memory-Efficient Distributed Training of Large-Scale Graph Neural Networks
topic Machine Learning
Distributed, Parallel, and Cluster Computing
url https://arxiv.org/abs/2502.17846