[논문 리뷰] FeatGraph: A Flexible and Efficient Backend for Graph Neural Network Systems

Yuwei Hu, Zihao Ye외 7명|arXiv (Cornell University)|2020. 08. 26.

Advanced Graph Neural Networks참고 문헌 43인용 수 12

한 줄 요약

FeatGraph는 그래프 구조와 텐서 수준의 병렬성을 활용하여 그래프 탐색과 특성 차원 계산을 공동 최적화하는 고성능 그래프 신경망(GNN) 백엔드이다. 이는 일반화된 스퍼스 템플릿(GSpMM 및 GSDDMM)과 사용자 정의 함수(UDF)를 조합한 이중 수준 프로그래밍 인터페이스를 통해 작동하며, 최신 시스템 대비 CPU에서 최대 32배, GPU에서 최대 7배의 성능 향상을 달성한다.

ABSTRACT

Graph neural networks (GNNs) are gaining increasing popularity as a promising approach to machine learning on graphs. Unlike traditional graph workloads where each vertex/edge is associated with a scalar, GNNs attach a feature tensor to each vertex/edge. This additional feature dimension, along with consequently more complex vertex- and edge-wise computations, has enormous implications on locality and parallelism, which existing graph processing systems fail to exploit. This paper proposes FeatGraph to accelerate GNN workloads by co-optimizing graph traversal and feature dimension computation. FeatGraph provides a flexible programming interface to express diverse GNN models by composing coarse-grained sparse templates with fine-grained user-defined functions (UDFs) on each vertex/edge. FeatGraph incorporates optimizations for graph traversal into the sparse templates and allows users to specify optimizations for UDFs with a feature dimension schedule (FDS). FeatGraph speeds up end-to-end GNN training and inference by up to 32x on CPU and 7x on GPU.

연구 동기 및 목표

• 희소 연산과 특성 텐서 계산의 비효율적인 처리로 인한 GNN 워크로드의 성능 저하 문제를 해결한다.
• 기존 딥러닝 프레임워크의 한계를 극복하여 기본적인 희소 연산 외의 다양한 GNN 커널에 대해 탄력적인 지원을 제공한다.
• 특성 텐서 차원이 계산 패턴을 본질적으로 변화시킨다는 점을 고려하여, 전통적인 그래프 처리 시스템과 GNN 워크로드 사이의 격차를 메운다.
• 그래프 탐색과 특성 차원 계산에 대한 분리된, 조합 가능한 최적화를 통해 CPU 및 GPU에서 고성능의 GNN 학습 및 추론을 가능하게 한다.
• 기존의 GNN 프레임워크인 DGL 및 PyTorch Geometric에 통합될 수 있는 이식성 있고 확장 가능한 백엔드를 제공한다.

제안 방법

• 정점 기반 계산을 위한 일반화된 SpMM과 간선 기반 계산을 위한 일반화된 SDDMM을 포함한 두 가지 일반화된 희소 템플릿을 도입하여, 특성 텐서 상에서 임의의 사용자 정의 함수(UDF) 지원을 가능하게 한다.
• 커널 정의를 거시적 탐색 로직(템플릿를 통한)과 미시적 특성 계산(UDF를 통한)으로 분리하여 각 수준에서 독립적인 최적화를 가능하게 한다.
• CPU에서의 데이터 국소성 향상과 메모리 액세스 오버헤드 감소를 위해 그래프 분할 및 캐시 인식 최적화를 적용한다.
• 희소 패턴에 적합한 병렬화 전략을 적용하여 GPU의 막대한 병렬성을 최대한 활용하며, 특히 특성 차원 계산에서의 성능 향상을 도모한다.
• CPU 및 GPU에서 최적의 스케줄링과 코드 생성 전략을 자동으로 탐색하기 위해 텐서 컴파일러 기법을 활용한다.
• 백엔드 가속을 위한 깔끔하고 확장 가능한 인터페이스를 제공함으로써 기존 GNN 프레임워크와의 탄력적인 통합을 지원한다.

실험 결과

연구 질문

• RQ1 희소 그래프 구조와 밀집 특성 텐서를 동시에 포함하는 GNN 워크로드를 효율적으로 표현하고 최적화할 수 있는 방법은 무엇인가?
• RQ2 기존 딥러닝 및 그래프 처리 시스템에서 GNN 전용 커널을 처리할 때 발생하는 주요 성능 저하 요인은 무엇인가?
• RQ3 그래프 탐색과 특성 차원 계산의 공동 최적화가 종단 간 GNN 학습 및 추론 성능을 크게 향상시킬 수 있는가?
• RQ4 템플릿 + UDF의 이중 수준 프로그래밍 추상화가 다양한 GNN 모델에 대해 탄력성과 고성능을 동시에 달성하는 데 얼마나 기여하는가?
• RQ5 특성 차원 계산과 그래프 탐색에 대한 대상 최적화가 CPU 및 GPU 플랫폼에서 성능에 미치는 영향은 어떠한가?

주요 결과

• FeatGraph는 종단 간 GNN 학습 및 추론 워크로드에서 최신 시스템 대비 CPU에서 최대 32배, GPU에서 최대 7배의 성능 향상을 달성한다.
• 일반화된 SpMM 및 SDDMM 템플릿과 사용자 정의 가능한 UDF의 조합은 MLP 집계 및 어텐션 메커니즘을 포함한 다양한 GNN 모델에 대한 표현적 지원을 가능하게 한다.
• 캐시 인식 그래프 분할 및 특성 차원 인식 로드 밸런싱 최적화는 메모리 액세스 오버헤드 감소로 인해 CPU 성능을 크게 향상시킨다.
• GPU 성능은 이전 GPU 기반 그래프 시스템이 대부분 忽시한 그래프 탐색 및 특성 텐서 계산 양쪽에서의 병렬성 활용으로 크게 향상된다.
• FeatGraph는 그래프 탐색과 특성 계산를 공동 최적화할 수 있는 능력 덕분에 cuSPARSE 및 MKL과 같은 벤더 최적화 라이브러리보다 GNN 전용 커널에서 뛰어난 성능을 보인다.
• 이식성과 확장성이 뛰어나 DGL 및 PyTorch Geometric와 같은 기존 GNN 프레임워크에 고성능 백엔드로 통합될 수 있다.