[논문 리뷰] DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs

Da Zheng, Chao Ma외 7명|arXiv (Cornell University)|2020. 10. 11.

Advanced Graph Neural Networks참고 문헌 34인용 수 23

한 줄 요약

DistDGL은 DGL 기반의 분산 GNN 학습 시스템으로, 국소성 인식 그래프 분할, 데이터 및 계산의 공재원, 다중 제약 조건 기반 로드 밸런싱을 통해 빌리언 스케일 그래프에서 효율적이고 확장 가능한 미니배치 학습을 가능하게 한다. 16台의 머신에서 선형 속도 향상을 달성하며, 1억 노드, 30억 간선 그래프의 학습 에포크를 13초 만에 완료하면서 모델 정확도를 유지한다.

ABSTRACT

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.

연구 동기 및 목표

• 단일 머신의 메모리 및 계산 용량을 초월하는 빌리언 스케일 그래프에서 GNN을 효율적으로 학습시키는 데 도전하는 것.
• 정점 간 의존성으로 인해 이웃 노드 특징을 가져오는 데 의해 지배되는 분산 GNN 학습의 통신 오버헤드를 줄이는 것.
• 클러스터 전반에 걸쳐 그래프 데이터, 특징, 계산을 공재원함으로써 높은 병렬 효율성과 메모리 확장성을 달성하는 것.
• 다중 제약 조건 기반 그래프 분할 및 균형 잡힌 미니배치 생성을 통해 머신 간 계산 로드를 균형 잡는 것.
• 비국소 노드를 포함한 에고 네트워크에서 동기식, 미니배치 학습을 통해 대규모 그래프로 확장하면서도 모델 정확도를 유지하는 것.

제안 방법

• 계산 분해를 위해 일관된 오너-계산 규칙을 사용하여 클러스터의 여러 머신에 그래프 구조, 노드 특징, 임베딩을 분산 배포한다.
• 상호 분할 통신을 최소화하고 로드 밸런싱을 보장하기 위해 다중 균형 제약 조건을 적용한 METIS 기반 최소 컷 그래프 분할을 사용한다.
• 메시지 전달 중 반복적인 원격 데이터 가져오기를 줄이기 위해 할로 노드를 파artitions 간에 복제한다.
• 동기식 SGD 학습 중 그래디언트 통신 오버헤드를 줄이기 위해 희소 임베딩 업데이트를 사용한다.
• 샘플링 서버와 메모리 기반 KVStore 서버를 트레이너가 있는 동일한 머신에 통합하여 데이터와 계산을 공재원한다.
• DGL 호환 API를 통해 최소한의 코드 변경으로 전도적 및 인도적 GNN 모델을 모두 지원한다.

실험 결과

연구 질문

• RQ1 이웃 데이터 가져오기로 인한 통신 오버헤드가 지배적인 상황에서, 대규모 분산 GNN 학습을 어떻게 효율적으로 만들 수 있는가?
• RQ2 다중 제약 조건 기반 로드 밸런싱을 적용한 최소 컷 그래프 분할이 빌리언 스케일 그래프에서 학습 성능을 크게 향상시킬 수 있는가?
• RQ3 통신 감소와 로드 밸런싱이 함께 작용할 경우, 분산 GNN 학습에서 선형 속도 향상을 얼마나 효과적으로 달성할 수 있는가?
• RQ4 DistDGL은 16台의 머신을 통해 1억 노드, 30억 간선 그래프로 확장하면서도 모델 정확도를 어떻게 유지하는가?
• RQ5 DGL 기반으로 구축된 시스템은 최소한의 코드 변경으로 단일 머신 학습에서 분산 학습으로의 원활한 마이그레이션을 가능하게 할 수 있는가?

주요 결과

• DistDGL은 16대의 머신에서 선형 속도 향상을 달성하며, 1억 노드, 30억 간선 그래프의 학습 에포크를 정확도 손실 없이 단 13초 만에 완료한다.
• ogbn-product 그래프에서 다중 제약 조건 기반 로드 밸런싱을 적용한 METIS 분할은 기본 METIS 대비 4% 성능 향상을 기록했으며, 무작위 분할 대비 2.14배 빠른 성능을 보였다.
• ogbn-papers100M 그래프에서 기본 METIS는 높은 불균형으로 인해 무작위 분할보다 성능이 열 劣하므로, 다중 제약 조건 기반 로드 밸런싱의 필요성을 입증했다.
• 최소 컷 분할에 의한 통신 감소와 다중 제약 조건 기반 로드 밸런싱의 조합이 고성능을 달성하는 데 필수적이다.
• 할로 노드 복제와 희소 임베딩 업데이트를 통해 학습 중 통신 트래픽을 크게 감소시켰다.
• DistDGL는 DGL의 API와 호환되어 단일 머신에서 분산 학습으로의 코드 마이그레이션을 최소한의 노력으로 가능하게 한다.