GSplit: 대규모 그래프 GNN 학습의 새로운 표준, 분할 병렬 처리

초록

기존 데이터 병렬 GNN 학습은 GPU 간 중복 샘플링과 계산으로 비효율적입니다. 본 연구는 ‘분할 병렬 처리(Split Parallelism)‘라는 새로운 하이브리드 병렬 훈련 패러다임을 제안합니다. 각 미니 배치를 여러 GPU에 걸쳐 중복 없이 분할하여 샘플링, 로딩, 학습을 수행함으로써 중복 작업을 제거합니다. 통신 오버헤드 문제를 해결하기 위한 경량 확률적 분할 알고리즘을 개발했으며, 이를 구현한 GSplit 시스템이 DGL, Quiver, P³ 대비 최대 4.4배 빠른 성능을 보입니다.

상세 분석

본 논문이 제안하는 ‘분할 병렬 처리(Split Parallelism)‘는 GNN 미니배치 학습의 근본적인 비효율성을 해결하는 패러다임 전환입니다. 기존 데이터 병렬 방식은 각 GPU가 독립적인 마이크로 배치를 처리하며, 샘플링된 k-hop 이웃 subgraph가 GPU 간에 중복되어 불필요한 특징값 로딩과 계산을 초래했습니다. GSplit은 하나의 미니 배치를 모든 GPU가 공유하고, 이를 ‘스플릿(split)‘이라는 겹치지 않는 파티션으로 실시간(on-the-fly) 분할하여 각 GPU에 할당합니다. 이로 인해 각 정점의 샘플링, 로딩, 순전파/역전파 계산이 정확히 한 GPU에서만 수행되어 이론적 중복 제거가 가능합니다.

핵심 기술적 도전과 해법은 다음과 같습니다:

1. 통신 오버헤드 최소화 및 부하 균형: 스플릿 할당은 GPU 간 특징값 셔플을 필요로 하며, 이 통신 비용이 중복 제거 이익을 상쇄할 수 있습니다. 저자들은 NP-hard인 온라인 그래프 분할 문제를 해결하기 위해 ‘확률적 분할 알고리즘’을 제안합니다. 이 알고리즘은 각 미니 배치를 샘플링할 때, 정점을 사전에 정의된 확률 분포에 따라 GPU에 할당합니다. 이 분포는 전체 그래프 구조를 기반으로 계산되며, 샘플된 서브그래프에서 ‘예상 통신 비용’을 최소화하고 ‘예상 작업량’을 균형 있게 분배함을 수학적으로 보장합니다. 이는 오프라인 분할(예: Metis)보다 현실적이며, 계산 오버헤드가 무시할 수 있을 정도로 가벼워 실시간 분할에 적합합니다.
2. 프로그래밍 추상화 및 커널 호환성: 시스템의 실용성을 위해 데이터 병렬 방식의 프로그래밍 모델을 최대한 유지했습니다. 이를 위해 ‘로컬 프론티어(Local Frontier)‘와 ‘믹스드 프론티어(Mixed Frontier)’ 개념을 도입하고, 효율적인 셔플을 위한 인덱스 구조를 제공합니다. 이 추상화는 기존에 최적화된 단일-GPU 샘플링/학습 커널(예: DGL, PyG)을 GSplit에서 재사용할 수 있게 합니다.
3. 캐싱과의 시너지: 분할 병렬 처리 자체로도 중복 로딩을 제거하지만, 특징값을 GPU 메모리에 캐싱하는 기법(Quiver 등)과 결합하면 추가적인 성능 향상을 기대할 수 있습니다. 각 GPU는 자신의 스플릿에 속한 정점의 특징값을 우선적으로 캐싱하므로, 캐시 효율성이 자연스럽게 향상됩니다.

실험 결과, GSplit은 Orkut, Papers100M 등 대규모 그래프에서 GAT, GraphSAGE 등 다양한 모델에 대해 기존 최신 시스템(DGL, Quiver, P³)을 큰 폭으로 앞섰습니다. 이 성능 향상의 주된 원인은 분할 알고리즘이 통신 비용을 효과적으로 통제하면서 중복 제거의 이점을 실현했기 때문으로 분석됩니다.