[논문 리뷰] PyTorch-Direct: Enabling GPU Centric Data Access for Very Large Graph Neural Network Training with Irregular Accesses

Seung Won Min, Wu Kun외 6명|arXiv (Cornell University)|2021. 01. 19.

Advanced Graph Neural Networks인용 수 11

한 줄 요약

PyTorch-Direct는 대규모 그래프 신경망(GNN) 학습을 위한 GPU 중심의 데이터 접근 패러다임을 도입하여, CPU 간섭 없이도 GPU가 주로 메모리에 산재해 있는 비정규적인 노드 특징에 직접 접근할 수 있도록 한다. 통합 텐서 유형을 도입하고 메모리 접근 패턴을 최적화함으로써, 평균적으로 데이터 전송 시간을 47.1% 감소시키고, GNN 학습 속도를 최대 1.6배로 향상시키며, 시스템 전력 소비량도 12.4%~17.5% 감소시킨다.

ABSTRACT

With the increasing adoption of graph neural networks (GNNs) in the machine learning community, GPUs have become an essential tool to accelerate GNN training. However, training GNNs on very large graphs that do not fit in GPU memory is still a challenging task. Unlike conventional neural networks, mini-batching input samples in GNNs requires complicated tasks such as traversing neighboring nodes and gathering their feature values. While this process accounts for a significant portion of the training time, we find existing GNN implementations using popular deep neural network (DNN) libraries such as PyTorch are limited to a CPU-centric approach for the entire data preparation step. This "all-in-CPU" approach has negative impact on the overall GNN training performance as it over-utilizes CPU resources and hinders GPU acceleration of GNN training. To overcome such limitations, we introduce PyTorch-Direct, which enables a GPU-centric data accessing paradigm for GNN training. In PyTorch-Direct, GPUs are capable of efficiently accessing complicated data structures in host memory directly without CPU intervention. Our microbenchmark and end-to-end GNN training results show that PyTorch-Direct reduces data transfer time by 47.1% on average and speeds up GNN training by up to 1.6x. Furthermore, by reducing CPU utilization, PyTorch-Direct also saves system power by 12.4% to 17.5% during training. To minimize programmer effort, we introduce a new "unified tensor" type along with necessary changes to the PyTorch memory allocator, dispatch logic, and placement rules. As a result, users need to change at most two lines of their PyTorch GNN training code for each tensor object to take advantage of PyTorch-Direct.

연구 동기 및 목표

• CPU 중심의 데이터 준비로 인해 발생하는 대규모 GNN 학습의 성능 저하 문제를 해결하여 CPU 부하를 줄이고 GPU 활용도를 높이기.
• 기존 DNN 라이브러리인 PyTorch와 같이 비정규적인 그래프 데이터 접근을 위해 CPU를 통해 데이터 수거가 필요한 비효율적인 절차를 극복하기.
• 주로 메모리에 산재해 있는 비정규적인 노드 특징에 GPU가 직접 접근하여 지연 시간과 데이터 이동 오버헤드를 줄이기.
• 통합 텐서 추상화와 PyTorch의 메모리 할당기 및 디스패치 로직에 대한 후행 호환성 있는 변경을 통해 프로그래머의 노력 최소화하기.
• 기존 PyTorch GNN 코드베이스와의 호환성을 유지하면서도, 상당한 엔드 투 엔드 학습 속도 향상과 시스템 전력 소비 절감을 달성하기.

제안 방법

• GPU 커널이 주로 메모리에 존재하는 데이터에 사전 CPU 측 데이터 수거 없이 직접 접근할 수 있도록 하는 새로운 '통합 텐서' 유형 도입.
• GPU 접근이 가능한 주로 메모리 지원을 위한 PyTorch 메모리 할당기 및 배치 규칙 수정으로 적절한 정렬 및 코alescing 최적화 지원.
• 주로 메모리의 비정규적인 정렬된 데이터 구조에 접근할 때 지연 시간을 줄이기 위해 정렬 인식 기반 GPU 커널 구현.
• GPU의 고도로 동시성 처리 능력을 활용해 주로 메모리의 산재한 데이터에 직접 메모리 접근을 수행함으로써 CPU의 데이터 수거 참여를 배제.
• 최소한의 코드 변경으로 기존 PyTorch 학습 워크플로우에 새로운 패러다임 통합 — 텐서당 최대 두 줄의 코드 변경으로 GPU 중심 접근 활성화 가능.
• 정렬 인식 기반 커널 디스패치 및 코ales스드 메모리 접근 전략을 통한 메모리 접근 패턴 최적화로 대역폭 활용도 향상.

실험 결과

연구 질문

• RQ1 GPU가 CPU 간섭 없이 주로 메모리에 산재해 있는 비정규적인 데이터에 효율적으로 접근할 수 있는가? 이는 GNN 학습에서 데이터 전송 오버헤드 감소에 기여하는가?
• RQ2 CPU 중심 접근 방식과 비교했을 때, GPU 중심의 데이터 접근 모델이 데이터 준비 시간을 얼마나 줄이고 엔드 투 엔드 GNN 학습 성능을 향상시키는가?
• RQ3 PyTorch-Direct는 GNN 학습 중 시스템 전력 소비에 어떤 영향을 미치는가? 특히 CPU 활용도 감소 시에 어떻게 되는가?
• RQ4 기존 PyTorch GNN 코드베이스에 GPU 중심 패러다임을 도입하기 위해 필요한 최소한의 프로그래머 노력은 무엇인가?
• RQ5 메모리 정렬은 GNN 워크로드에서 주로 메모리에 저장된 비정규적인 데이터에 대한 GPU 직접 접근 성능에 어떤 영향을 미치는가?

주요 결과

• PyTorch-Direct는 엔드 투 엔드 GNN 학습에서 기준 PyTorch 접근 방식 대비 평균적으로 데이터 전송 시간을 47.1% 감소시켰다.
• 이 프레임워크는 GNN 학습 성능을 최대 1.6배로 향상시켰으며, 다양한 데이터셋과 아키텍처에서 관측된 속도 향상은 1.01배에서 1.45배의 범위를 보였다.
• 직접 전기계량기 측정을 통해 확인한 결과, 낮아진 CPU 활용도로 인해 학습 중 시스템 전력 소비량이 12.4%에서 17.5% 감소했다.
• 마이크로 벤치마크 결과, PyTorch-Direct는 기준 PyTorch 접근 방식 대비 평균적으로 2.39배 높은 성능을 기록했으며, 다양한 시스템 구성에서 거의 이상적인 대역폭 활용도를 달성했다.
• PyTorch-Direct의 메모리 정렬 최적화는 2052바이트 특징에 대해 단순 GPU 접근 방식 대비 1.95배 성능 향상을 보이며, 코ales스드 접근 패턴의 효과를 입증했다.
• 사용자는 텐서당 최대 두 줄의 코드 변경으로 PyTorch-Direct의 전반적인 이점을 활성화할 수 있으며, 이는 후행 호환성과 낮은 도입 장벽을 보장한다.