[논문 리뷰] Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining

Tim Kaler, Nickolas Stathas외 5명|arXiv (Cornell University)|2021. 10. 16.

Advanced Graph Neural NetworksComputer Science참고 문헌 35인용 수 21

한 줄 요약

이 논문은 이웃 샘플링 최적화, 공유 메모리 병렬 배치 준비, 그리고 GPU 계산과 데이터 전송을 파ipelin-ing함으로써 GNN 학습 및 추론을 가속화하는 SALIENT 시스템을 제안한다. 단일 GPU에서 표준 PyTorch-Geometric 대비 3배의 성능 향상을 달성했으며, ogbn-papers100M에서 16개 GPU로 병렬 처리 시 8배의 성능 향상을 기록했으며, 3층 GraphSAGE 모델을 1에포크당 2.0초만에 학습하고 테스트 정확도 64.58%를 달성한다.

ABSTRACT

Improving the training and inference performance of graph neural networks (GNNs) is faced with a challenge uncommon in general neural networks: creating mini-batches requires a lot of computation and data movement due to the exponential growth of multi-hop graph neighborhoods along network layers. Such a unique challenge gives rise to a diverse set of system design choices. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment, under which we identify major performance bottlenecks hitherto under-explored by developers: mini-batch preparation and transfer. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler, a shared-memory parallelization strategy, and the pipelining of batch transfer with GPU computation. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised. Such an observation unifies training and inference, simplifying model implementation. We report comprehensive experimental results with several benchmark data sets and GNN architectures, including a demonstration that, for the ogbn-papers100M data set, our system SALIENT achieves a speedup of 3x over a standard PyTorch-Geometric implementation with a single GPU and a further 8x parallel speedup with 16 GPUs. Therein, training a 3-layer GraphSAGE model with sampling fanout (15, 10, 5) takes 2.0 seconds per epoch and inference with fanout (20, 20, 20) takes 2.4 seconds, attaining test accuracy 64.58%.

연구 동기 및 목표

• 소규모 배치 준비 및 데이터 전송으로 인한 고비용으로 인해 발생하는 GNN 학습 및 추론의 성능 저하 문제를 해결하기 위해, 특히 분산 다중 GPU 환경에서의 성능 저하를 해결한다.
• PyTorch-Geometric 및 DGL과 같은 기존 GNN 프레임워크에서 아직 탐색되지 않은 성능 저하 요인을 규명하고 해결한다. 특히 배치 처리 중 CPU 활용도 및 메모리 대역폭 문제를 중심으로 분석한다.
• 정확도에 미치는 영향이 최소한임을 경험적으로 검증함으로써, 이웃 샘플링을 통한 학습과 추론의 통합을 통해 모델 구현을 단순화한다.
• 데이터 전송과 계산을 파이프라인 처리하여 GPU 활용도를 극대화함으로써, 종단 간 학습 시간이 오직 가장 느린 구성 요소에 의해 제한되는 것을 보장한다.
• PyTorch와 PyG와의 호환성을 유지하는 일반 목적의 시스템을 구축하여, 다른 프레임워크 및 대규모 분산 환경으로의 쉽게 확장 가능하도록 설계한다.

제안 방법

• CPU 아키텍처 최적화를 고려한 성능 최적화된 이웃 샘플러 설계로, 다단계 이웃 확장 과정에서의 오버헤드를 감소시킨다.
• 기존 프레임워크에서 발생하는 CPU 및 메모리 대역폭 저하 문제를 해결하기 위해, 배치 준비 과정에 공유 메모리 병렬 처리를 구현한다.
• CPU와 GPU 간의 파이프라인 데이터 전송을 도입하여, 데이터 이동과 GNN 계산을 겹치도록 하여 GPU 활용도를 향상시킨다.
• 다중 머신 및 GPU에서의 분산 학습을 위해 PyTorch의 DDP 모듈을 활용하여, 최소한의 코드 변경으로 수평 확장이 가능하도록 한다.
• 기존 GNN 아키텍처 및 프레임워크(예: GraphSAGE)를 그대로 활용하여 아키텍처 변경 없이도 호환성과 도입 용이성을 확보한다.
• 샘플링이 추론 정확도에 미치는 영향을 경험적으로 평가하여, 학습 및 추론 전반에 걸쳐 샘플링을 균일하게 적용해도 정확도 손실가 최소화됨을 검증한다.

실험 결과

연구 질문

• RQ1 다중 GPU 환경에서 이웃 샘플링을 사용할 경우 GNN 학습 및 추론의 주요 성능 저하 요인은 무엇인가?
• RQ2 배치 준비 및 데이터 전송을 어떻게 최적화할 수 있으며, 이로 인해 전체 학습 스루풋에 미치는 영향을 최소화할 수 있는가?
• RQ3 이웃 샘플링은 추론 정확도에 어느 정도 영향을 미치며, 학습 및 추론 전반에 걸쳐 균일하게 적용해도 정확도 손실이 크지 않은가?
• RQ4 GPU 계산과 데이터 전송을 파이프라인 처리하면 GNN 워크로드에서 거의 이상적인 GPU 활용도를 달성할 수 있는가?
• RQ5 제안된 시스템은 다중 GPU 및 다중 머신 환경에서 어떻게 스케일업되며, 기존 구현 대비 달성 가능한 성능 향상은 어느 정도인가?

주요 결과

• SALIENT는 ogbn-papers100M에서 3층 GraphSAGE 모델 학습 시 단일 GPU에서 표준 PyTorch-Geometric 대비 3배의 성능 향상을 달성한다.
• 16개 GPU를 사용할 경우 SALIENT는 8배의 병렬 성능 향상을 기록하여 동일한 모델과 데이터셋에서 1에포크당 학습 시간을 2.0초로 단축시킨다.
• fanout (20, 20, 20)로 추론을 수행할 경우 배치당 2.4초가 소요되며, 테스트 정확도 64.58%를 달성하여 샘플링으로 인한 정확도 저하가 최소한임을 입증한다.
• 배치 준비, 데이터 전송, GPU 계산 간의 거의 완벽한 겹침을 달성하여 종단 간 학습 시간이 거의 오직 가장 느린 개별 구성 요소에 의해 결정된다.
• 최적화된 샘플링과 공유 메모리 병렬 처리 덕분에, 특히 CPU 및 메모리 대역폭에서의 성능 저하 요인을 크게 감소시켰다.
• 제안된 최적화 기법은 일반적이며 다른 GNN 프레임워크 및 분산 환경에 적용 가능하며, GPU 기반 슬라이싱 또는 데이터 캐싱과 같은 기술과의 통합 가능성도 보유하고 있다.