통신 불필요 샘플링 및 4차원 하이브리드 병렬성을 이용한 확장 가능한 미니 배치 그래프 신경망 학습

Communication-free Sampling and 4D Hybrid Parallelism for Scalable Mini-batch GNN Training

그래프 신경망(GNN)은 다양한 실제 시나리오에서 파생된 그래프 데이터셋을 학습하는 데 널리 사용됩니다. 매우 큰 그래프에서 학습하려면 분산 학습이 필요하며, 미니 배치와 함께 샘플링을 사용하는 것이 GNN 학습을 병렬화하는 인기 있는 방법입니다. 기존의 분산 미니 배치 방식은 비용이 많이 드는 샘플링 방법과 데이터 병렬 처리 시 제한적인 확장성으로 인해 상당한 성능 병목 현상을 가지고 있습니다. 본 논문에서는 통신 불필요 분산 샘플링, 3차원 병렬 행렬 곱셈(PMM) 및 데이터 병렬 처리를 결합하여 확장 가능한 미니 배치 GNN 학습을 위한 4차원 병렬 프레임워크인 ScaleGNN을 제안합니다. ScaleGNN은 각 프로세스(GPU 장치)가 프로세스 간 통신 없이 로컬 미니 배치, 즉 서브 그래프 파티션을 구성할 수 있도록 하는 균일한 정점 샘플링 알고리즘을 도입합니다. 3D PMM을 통해 일반적인 데이터 병렬 처리보다 훨씬 많은 GPU 수로 미니 배치 학습을 확장할 수 있으며, 통신 오버헤드가 크게 감소합니다. 또한 샘플링과 학습을 겹쳐서 수행하고, 더 낮은 정밀도로 데이터를 전송하여 통신 오버헤드를 줄이고, 커널 융합 및 통신-계산 중첩을 통해 추가적인 최적화를 수행합니다. ScaleGNN을 다섯 가지 그래프 데이터셋으로 평가하여 Perlmutter에서 2048개의 GPU, Frontier에서 2048개의 GCD, Tuolumne에서 1024개의 GPU까지 강력한 확장성을 보여줍니다. Perlmutter에서 ScaleGNN은 ogbn-products 데이터셋에 대해 최첨단(SOTA) 기준보다 3.5배 더 빠른 전체 학습 속도를 달성했습니다.

Graph neural networks (GNNs) are widely used for learning on graph datasets derived from various real-world scenarios. Learning from extremely large graphs requires distributed training, and mini-batching with sampling is a popular approach for parallelizing GNN training. Existing distributed mini-batch approaches have significant performance bottlenecks due to expensive sampling methods and limited scaling when using data parallelism. In this work, we present ScaleGNN, a 4D parallel framework for scalable mini-batch GNN training that combines communication-free distributed sampling, 3D parallel matrix multiplication (PMM), and data parallelism. ScaleGNN introduces a uniform vertex sampling algorithm, enabling each process (GPU device) to construct its local mini-batch, i.e., subgraph partitions without any inter-process communication. 3D PMM enables scaling mini-batch training to much larger GPU counts than vanilla data parallelism with significantly lower communication overheads. We also present additional optimizations to overlap sampling with training, reduce communication overhead by sending data in lower precision, kernel fusion, and communication-computation overlap. We evaluate ScaleGNN on five graph datasets and demonstrate strong scaling up to 2048 GPUs on Perlmutter, 2048 GCDs on Frontier, and 1024 GPUs on Tuolumne. On Perlmutter, ScaleGNN achieves 3.5x end-to-end training speedup over the SOTA baseline on ogbn-products.