PaperlayerView exact revision on arXiv

Source-linked AI summary

DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs

Da Zheng, Chao Ma, Minjie Wang, Jinjing Zhou, Qidong Su, Xiang Song, Quan Gan, Zheng Zhang, George Karypis

arXiv:2010.05337v3cs.LGcs.DC

TL;DR

Large graphs make GNN training difficult because vertex dependencies expand mini-batches and drive distributed neighbor-data traffic. DistDGL provides distributed mini-batch training through partitioned graph data, synchronous execution, load balancing, halo replication, and sparse updates. It reports linear speedup without compromising model accuracy, including 13 seconds per epoch for a 100-million-node, 3-billion-edge graph on 16 machines.

  • Problem

    GNN training must scale to graphs with hundreds of millions or billions of nodes while handling dependent neighbor data and substantial distributed memory and communication demands.

  • Method

    DistDGL distributes graph data and computation across machines, using synchronous mini-batch training, min-cut partitioning, balancing constraints, halo vertices, and sparse embedding updates.

  • Results

    DistDGL achieves linear training speedup without compromising model accuracy and trains GraphSage for 13 seconds per epoch on 100 million nodes and 3 billion edges using 16 machines.

  • Takeaways & Limitations

    DistDGL demonstrates that reducing network communication together with balancing graph partitions and mini-batches supports efficient distributed GNN training on CPU clusters.

Abstract

from arXiv · show

Graph neural networks (GNN) have shown great success in learning from graph-structured data. They are widely used in various applications, such as recommendation, fraud detection, and search. In these domains, the graphs are typically large, containing hundreds of millions of nodes and several billions of edges. To tackle this challenge, we develop DistDGL, a system for training GNNs in a mini-batch fashion on a cluster of machines. DistDGL is based on the Deep Graph Library (DGL), a popular GNN development framework. DistDGL distributes the graph and its associated data (initial features and embeddings) across the machines and uses this distribution to derive a computational decomposition by following an owner-compute rule. DistDGL follows a synchronous training approach and allows ego-networks forming the mini-batches to include non-local nodes. To minimize the overheads associated with distributed computations, DistDGL uses a high-quality and light-weight min-cut graph partitioning algorithm along with multiple balancing constraints. This allows it to reduce communication overheads and statically balance the computations. It further reduces the communication by replicating halo nodes and by using sparse embedding updates. The combination of these design choices allows DistDGL to train high-quality models while achieving high parallel efficiency and memory scalability. We demonstrate our optimizations on both inductive and transductive GNN models. Our results show that DistDGL achieves linear speedup without compromising model accuracy and requires only 13 seconds to complete a training epoch for a graph with 100 million nodes and 3 billion edges on a cluster with 16 machines. DistDGL is now publicly available as part of DGL:https://github.com/dmlc/dgl/tree/master/python/dgl/distributed.

I. INTRODUCTION

GNNs target increasingly massive graphs, but vertex dependencies make mini-batch and distributed training communication-intensive. DistDGL addresses these challenges with distributed mini-batch training, graph partitioning, load balancing, and communication optimizations.

  • Graphs in applications such as social networks, recommendation, and knowledge graphs can contain billions of nodes or edges.Examples include Facebook’s social graph, Amazon’s user-item graph, and Freebase with 1.9 billion triples.
  • GNN mini-batches must include dependent neighboring samples, whose number can grow exponentially as more neighbor hops are explored.This dependency distinguishes GNN training from training with independent samples.
  • Distributed GNN training is constrained by terabyte-scale graph data, neighbor-data traffic, and synchronized updates needed for model accuracy.Neighbor vertex data accounts for most distributed traffic, unlike conventional distributed neural-network training, where parameter gradients dominate communication.
  • Existing systems either target full-batch computation, remain limited to single machines, or incur heavy network traffic when fetching neighbor data.Architectures designed around exchanging model gradients do not directly address GNNs’ vertex-dependency bottleneck.
  • DistDGL distributes graph data and training components across machines, using synchronous training, min-cut partitioning, multi-constraint balancing, and other communication optimizations.Its design includes non-local nodes in ego-networks and co-locates data with computation.
  • 13 seconds per epoch is reported for GraphSage on a graph with 100 million nodes and 3 billion edges using 16 machines.The experiments also report linear speedup without compromising model accuracy as machines increase.

A. Graph Neural Networks

GNNs learn representations from graph structure and vertex or edge features through message passing. Their models may contain shared dense parameters and selectively updated vertex-specific sparse embeddings.

  • GNNs learn joint representations from graph structure together with vertex and edge features.Message passing provides the formulation used for these models.
  • Each vertex broadcasts messages to neighbors and aggregates received messages to compute its representation.This process is applied iteratively across multiple GNN layers.
  • A GNN layer uses input vertex features and edge features to calculate messages, aggregate them, and update vertex representations.The paper denotes the relevant functions as f, L, and g.
  • The functions f, L, and g can be customizable or parameterized modules for messaging, aggregation, and representation updates.Repeated application generates representations for multiple layers.
  • Dense parameters are shared across vertices and updated every mini-batch, whereas sparse parameters represent vertex embeddings updated only for participating vertices.Sparse embeddings are additional model parameters used by some GNN models.

B. Mini-batch training

Mini-batch GNN training samples target vertices and recursively gathers neighbors to capture graph dependencies. DistDGL distributes this workflow through samplers, a distributed KVStore, trainers, and synchronized model updates.

  • B. Mini-batch training: Because vertex dependencies distinguish GNNs from conventional neural networks, sampled subgraphs must capture dependencies in the original graph.This creates a central scalability requirement for GNN frameworks.
  • B. Mini-batch training: GNN mini-batch training samples target vertices, recursively selects bounded-fan-out neighbors, and computes target representations from sampled messages.The recursion depth follows the number of GNN layers.
  • A. Distributed Training Architecture: DistDGL uses synchronous SGD in which machines compute gradients on local mini-batches, synchronize them, and update local model replicas.This preserves the distributed training workflow across machines.
  • A. Distributed Training Architecture: Samplers generate mini-batch graph structures, while a distributed KVStore serves vertex and edge data and manages vertex embeddings.Sampler access uses an interface compatible with DGL’s neighbor-sampling API.
  • A. Distributed Training Architecture: Trainers fetch sampled graphs and features, run forward and backward computation, and send dense and sparse updates through separate components.Dense gradients are synchronized, while sparse embedding gradients are returned to the KVStore.
  • A. Distributed Training Architecture: DistDGL co-locates graph partitions with computation so samplers can access neighbors locally without inter-sampler communication.METIS assigns edges uniquely while allowing duplicated HALO vertices across partitions.

B. Graph Partitioning

DistDGL partitions graphs to minimize cross-partition edges while balancing training workloads, then assigns graph data and IDs to support efficient distributed access.

  • Graph Partitioning: Graph partitioning minimizes the number of edges crossing partitions and is amortized across multiple distributed training runs.Partitioning is performed once before training.
  • Graph Partitioning: METIS groups densely connected vertices, assigns incident edges to the same partition, and keeps neighboring data locally accessible for sampling.Each edge has a unique assignment, while duplicated vertices are designated HALO vertices.
  • Graph Partitioning: DistDGL uses multi-constraint partitioning to balance vertices, edges, and training samples for synchronous mini-batch workloads.Balancing only vertex counts is insufficient because trainers need similarly sized batches and comparable batch counts per epoch.
  • Graph Partitioning: DistDGL extends METIS with edge sparsification, out-of-core processing, and limited refinement to reduce partitioning memory and computation costs.Coarse graphs retain high-weight edges so their solutions remain useful for finer graphs.
  • Graph Partitioning: Vertex and edge features are partitioned with graph data, while local and global IDs are maintained for efficient internal lookup and model-facing identification.Only core-vertex and edge features are assigned to partitions, avoiding feature duplication.

C. Distributed Key-Value Store

DistDGL’s distributed KVStore manages partitioned features and sparse embeddings, emphasizing co-location, fast local access, and efficient sparse updates.

  • Distributed Key-Value Store: DistDGL uses a distributed in-memory KVStore for vertex features, edge features, and vertex embeddings across graph partitions.The design targets better data co-location, faster high-speed-network access, and efficient sparse embedding updates.
  • Distributed Key-Value Store: Separate vertex-data and edge-data partition policies align KVStore placement with the graph partitions on each machine.Vertex and edge data may be mapped differently because their graph partitioning patterns differ.
  • Distributed Key-Value Store: Shared memory lets trainers access most locally co-located KVStore data directly without inter-process communication overhead.Local KVStore placement is enabled by co-location of data and computation.
  • Distributed Key-Value Store: The KVStore supports sparse vertex embeddings for transductive mini-batch models, where only a small subset is involved and updated per iteration.This support addresses distributed sparse-embedding updates.

D. Distributed Sampler

DistDGL combines distributed sampling, balanced trainer workloads, asynchronous sparse updates, and hybrid CPU parallelism to support mini-batch training.

  • Distributed Sampler: Trainers issue sampling requests to sampler servers according to core-vertex assignments, while multiple sampling workers generate mini-batches in parallel.Sampling requests can be overlapped with mini-batch training.
  • Distributed Sampler: Shared memory accelerates local sampling, and asynchronous remote requests overlap network RPCs with local sampling computation.This design hides network latency when local sampling takes substantial time.
  • Mini-batch Trainer: DistDGL splits training samples evenly across trainers using balanced partitions and a two-level workload assignment strategy.The strategy trades a small amount of data locality for computation balance.
  • Mini-batch Trainer: Uniform independent sampling combined with balanced partitions preserves uniform sampling across the dataset under synchronous SGD.The passage states that distributed training therefore theoretically does not affect convergence rate or model accuracy.
  • Mini-batch Trainer: Dense parameters use synchronous SGD, while sparse vertex embeddings use asynchronous Hogwild updates to overlap communication and computation.Concurrent conflicts are described as rare because trainers typically update different embeddings.
  • Mini-batch Trainer: Distributed CPU training combines multiprocessing across NUMA-oriented trainer processes with OpenMP multithreading inside each process.More trainer processes also increase parameter-update communication overhead.

IV. EVALUATION

DistDGL’s evaluation tests scalability, data locality, and workload balancing for node classification, using GraphSAGE on OGB datasets and multi-machine clusters.

  • IV. EVALUATION: The evaluation asks whether DistDGL scales GNN training, improves data locality, and balances workloads across machines.The study focuses on node classification rather than link prediction.
  • IV. EVALUATION: The benchmark uses three-layer GraphSAGE with hidden size 256 and fan-outs 15, 10, and 5 on two OGB datasets.Experiments use AWS EC2 m5n.24xlarge instances with 96 VCPUs and 384GB RAM each.
  • IV. EVALUATION: DistDGL is compared with Euler on four m5n.24xlarge instances using the same global mini-batch size.Euler uses a different multiprocessing-based parallelization strategy.

A. DistDGL vs. other distributed GNN frameworks

DistDGL outperforms Euler in distributed mini-batch training by improving data copying and sampling through locality-aware partitioning and co-located computation.

  • 2.2× speedup over Euler across different batch sizes demonstrates DistDGL’s overall training advantage.
  • More than 5× speedup in data copying is DistDGL’s main performance advantage over Euler.METIS partitions minimize edge cuts and co-locate trainers with partition data, bringing copy speed close to local memory performance.
  • 2× speedup in sampling results from sampling most vertices and edges locally within each partition.
  • DistDGL is slightly faster than Euler in mini-batch computation and gradient synchronization.DistDGL uses DGL and PyTorch for mini-batch computation and PyTorch for gradient synchronization.

B. DistDGL’s sparse embedding vs. Pytorch’s sparse embedding

DistDGL’s distributed embeddings substantially accelerate GraphSage training, while its broader system scales across machines and preserves convergence accuracy.

  • Sparse embedding comparison: Almost 70× speedup is achieved with DistDGL’s sparse embeddings over PyTorch’s sparse embeddings on OGBN-PRODUCT.DistDGL updates sparse embeddings through its KVStore, whereas PyTorch’s AllReduce requires padding sparse gradients to equal shapes.
  • Scalability: Linear speedup is achieved as the number of machines increases for both OGB datasets.The result indicates that network communication does not become the bottleneck and that the system remains balanced as machines increase.
  • Scalability: 13 seconds per epoch is required to train GraphSage on OGBN-PAPERS100M using 16 m5.24xlarge machines.
  • Single-machine comparison: DistDGL running on one machine with two trainers outperforms DGL’s multiprocessing training.The comparison attributes the difference to DistDGL’s dedicated sampler processes versus PyTorch dataloader multiprocessing.
  • Convergence: DistDGL quickly converges to almost the same peak accuracy as single-machine training.Single-machine training takes much longer to converge.

D. Ablation Study

The ablation study shows that effective distributed GNN performance requires both reducing network communication and balancing computation across graph partitions.

  • 2.14× speedup over random partitioning is achieved by default METIS on OGBN-PRODUCT through superior network-communication reduction.
  • 4% additional improvement over default METIS partitioning comes from adding multiple balancing constraints on OGBN-PRODUCT.
  • Default METIS performs much worse than random partitioning on OGBN-PAPERS100M because its partitions are highly imbalanced.This occurs despite METIS effectively reducing edge cuts between partitions.
  • METIS partitioning with multi-constraints achieves good performance on both datasets.The study compares multi-constraint METIS with random partitioning and default METIS on a four-machine cluster.

B. Distributed GNN Training

Distributed GNN training must address graph-dependent mini-batch computation and neighbor-data communication at scale. DistDGL combines locality-aware partitioning, co-located computation, and balancing strategies to achieve scalable training without compromising accuracy.

  • Full-graph distributed training requires aggregated device memory to fit the graph and can make each model update computationally expensive.
  • Distributed mini-batch GNN frameworks face substantial network traffic because vertex dependencies require fetching neighbor data.
  • DistDGL uses METIS partitioning with minimum edge cuts and co-locates data and computation to reduce network communication.
  • Multiple strategies balance graph partitions and mini-batches generated from each partition.
  • Linear training speedup on CPU clusters is achieved without compromising model accuracy.
Loading 2010.05337v3…