[논문 리뷰] AGL: a Scalable System for Industrial-purpose Graph Machine Learning

Dalong Zhang, Xin Huang외 9명|arXiv (Cornell University)|2020. 03. 05.

Advanced Graph Neural Networks참고 문헌 30인용 수 32

한 줄 요약

AGL은 산업용 그래프 머신 러닝을 위한 확장 가능하고 장애 내성 있는 시스템으로, k-hop 이웃과 MapReduce를 사용하여 일반 클러스터에서 GNN의 학습과 추론을 완전하게 제공합니다.

ABSTRACT

Machine learning over graphs have been emerging as powerful learning tools for graph data. However, it is challenging for industrial communities to leverage the techniques, such as graph neural networks (GNNs), and solve real-world problems at scale because of inherent data dependency in the graphs. As such, we cannot simply train a GNN with classic learning systems, for instance parameter server that assumes data parallel. Existing systems store the graph data in-memory for fast accesses either in a single machine or graph stores from remote. The major drawbacks are in three-fold. First, they cannot scale because of the limitations on the volume of the memory, or the bandwidth between graph stores and workers. Second, they require extra development of graph stores without well exploiting mature infrastructures such as MapReduce that guarantee good system properties. Third, they focus on training but ignore the optimization of inference over graphs, thus makes them an unintegrated system. In this paper, we design AGL, a scalable, fault-tolerance and integrated system, with fully-functional training and inference for GNNs. Our system design follows the message passing scheme underlying the computations of GNNs. We design to generate the $k$-hop neighborhood, an information-complete subgraph for each node, as well as do the inference simply by merging values from in-edge neighbors and propagating values to out-edge neighbors via MapReduce. In addition, the $k$-hop neighborhood contains information-complete subgraphs for each node, thus we simply do the training on parameter servers due to data independency. Our system AGL, implemented on mature infrastructures, can finish the training of a 2-layer graph attention network on a graph with billions of nodes and hundred billions of edges in 14 hours, and complete the inference in 1.2 hour.

연구 동기 및 목표

• Address the challenge of scaling graph neural networks (GNNs) for industrial-scale graphs with billions of nodes and hundreds of billions of edges.
• Propose an integrated system design (training and inference) based on mature infrastructures (MapReduce, parameter servers) to ensure fault tolerance and scalability.
• Introduce the concept of k-hop neighborhoods to enable data-independence and simplify training on subgraphs.
• Develop three main modules (GraphFlat, GraphTrainer, GraphInfer) to realize scalable generation of neighborhoods, training, and inference.
• Improve efficiency and throughput of both training and inference on large-scale attributed graphs.

제안 방법

• Use the message passing paradigm of GNNs to define and generate information-complete k-hop neighborhoods for each target node.
• Implement GraphFlat to distributedly generate and flatten k-hop neighborhoods into protobufs stored on a distributed filesystem.
• Apply re-indexing, sampling, and inverted indexing to handle hub nodes and balance load in GraphFlat.
• Develop GraphTrainer as a distributed parameter-server-like framework with pipeline, pruning, and edge-partitioning optimizations to train on k-hop neighborhoods on commodity hardware.
• Design GraphInfer to perform distributed inference by hierarchical model segmentation and a MapReduce-based message-passing pipeline that propagates embeddings layer by layer.

실험 결과

주요 결과

• AGL can train a 2-layer Graph Attention Network (GAT) on a graph with 6.23e9 nodes and 3.38e11 edges in 14 hours using 100 workers.
• Inference on the same graph completes in 1.2 hours.
• The system achieves near-linear speedup on CPU clusters compared with single-machine training (vs. DGL/PyG on a monster machine).
• GraphFlat with GraphTrainer and GraphInfer leverages MapReduce and parameter-server infrastructures to provide fault-tolerant, scalable training and inference on industrial-scale graphs.
• GraphInfer benefits from embedding reuse across nodes, maximizing inference efficiency compared to architecture like DGL and AliGraph.
• The approach represents one of the largest-scale deployments of graph embeddings in real industrial scenarios.