CPU·GPU 연동 캐시와 자원 인식 그래프 분할로 최적화한 병렬 GNN 학습

초록

CaPGNN은 단일 서버의 다중 GPU 환경에서 전체 배치 GNN 학습 시 발생하는 대규모 통신 비용을 감소시키기 위해 CPU‑GPU 이중 캐시와 GPU 성능을 고려한 그래프 파티셔닝을 결합한 프레임워크이다. 적응형 캐시와 자원 인식 파티셔닝을 통해 최대 18.98배의 학습 가속과 99% 수준의 통신 절감 효과를 달성한다.

상세 분석

본 논문은 대규모 그래프에 대한 전체 배치(full‑batch) GNN 학습이 GPU 메모리 한계와 파티션 간 경계(halo) 정점의 중복 전송으로 인한 통신 병목에 직면한다는 점을 정확히 짚어낸다. 이를 해결하기 제안된 두 핵심 기술은 (1) Joint Adaptive Caching Algorithm (JACA)과 (2) Resource‑Aware Partitioning Algorithm (RAPA)이다.

JACA는 CPU와 GPU 메모리를 2단계 캐시로 활용한다. 먼저 그래프 파티션 단계에서 각 GPU가 필요로 하는 정점 피처를 추정하고, “중요도 점수”(접근 빈도, 차원 수, 최근 사용 시점 등)를 기반으로 고우선순위 정점을 GPU 캐시에, 나머지는 CPU 캐시에 배치한다. 캐시 업데이트는 경량화된 “증분 교체” 방식으로 수행되며, pinned memory와 비동기 스트림을 이용해 CPU‑GPU 간 전송을 겹치게 만든다. 또한, 사전 페치(prefetch)와 스테일러스 허용 파이프라인을 도입해 연산이 진행되는 동안 필요한 정점 데이터를 미리 로드함으로써 통신 지연을 최소화한다. 이러한 설계는 특히 halo 정점이 전체 정점 수의 30~70%를 차지하는 경우에 중복 전송을 크게 억제한다.

RAPA는 기존 METIS·Fennel 등 균형 파티셔닝 기법이 GPU 간 연산·통신 성능 차이를 무시한다는 한계를 보완한다. 논문은 각 GPU의 FLOPS, 메모리 용량, NVLink 대역폭 등을 정량화한 “자원 프로파일”을 정의하고, 이를 기반으로 파티션 크기와 halo 정점 비율을 동적으로 조정한다. 그래프 희소화(sparsification) 아이디어를 차용해, 낮은 차수 정점은 파티션 경계에서 제거하거나 복제 비율을 낮추어 전체 통신량을 감소시킨다. 결과적으로 고성능 GPU는 더 큰 서브그래프를 할당받아 연산 효율을 높이고, 저성능 GPU는 작은 파티션과 최소한의 halo를 받아 로드 밸런스를 맞춘다.

파이프라인 설계는 CUDA 스트림을 이용해 “계산‑통신 겹침”을 구현한다. 각 GPU는 로컬 정점 연산을 수행하면서 동시에 CPU‑GPU 캐시 업데이트와 halo 정점 교환을 비동기적으로 진행한다. 이때 스테일러스 허용 정책을 적용해 최신 정점 피처가 아니더라도 일정 횟수까지는 재사용하도록 함으로써, 파이프라인이 정체되지 않게 한다.

실험에서는 Reddit, ogbn‑products, Amazon 등 7개의 벤치마크 그래프와 GCN, GraphSAGE, GAT 등 3가지 모델을 대상으로 단일 서버 4‑GPU (RTX 3090, RTX 4090 등 이기종) 환경에서 평가했다. CaPGNN은 기존 DGL‑Dist, NeuGraph, HongTu 대비 평균 9.3배(최대 18.98배) 빠른 epoch 시간을 기록했으며, 통신량은 99% 이상 감소했다. 정확도 측면에서는 대부분 0.1% 이내의 차이 또는 소폭 향상을 보였으며, 이는 캐시 스테일러스가 모델 수렴에 미치는 영향을 최소화했기 때문이다. 또한, 다중 서버 다중 GPU 환경으로 확장했을 때도 동일한 캐시·파티셔닝 전략이 적용 가능함을 시연하였다.

한계점으로는 (1) 캐시 관리 오버헤드가 매우 큰 그래프(수억 정점, 고차원 피처)에서는 메모리 압박이 발생할 수 있고, (2) 현재 구현이 CUDA 기반 단일 머신에 최적화돼 있어, 비동기 네트워크 환경(예: Ethernet 기반 클러스터)에서는 추가적인 비용 모델이 필요하다는 점을 언급한다. 향후 연구에서는 계층형 캐시(SSD·NVMe)와 동적 파티션 재조정 기법을 결합해 더욱 큰 규모의 그래프에도 적용 가능한 프레임워크를 목표로 제시한다.