서론

그래프는 다양한 실제 세계의 네트워크를 표현하는 데 사용되며, 객체와 그 관계는 각각 노드와 엣지로 모델링된다. 최근에는 *그래프 신경망(GNNs)*라는 클래스의 딥 뉴럴 네트워크가 이러한 그래프 구조 데이터를 학습하도록 특별히 설계되어 다양한 하류 작업에서 혁신을 이루었다.

style="width:47.0%" />

현재 연구들은 노드 특성과 그래프 위상을 모두 고려하여 그래프 구조 정보를 학습하도록 모델 아키텍처를 설계했지만, 그들의 기본 가정은 전체 입력 데이터가 GNN 학습 동안 GPU 또는 메인 메모리에 존재한다는 것이다. 그러나 실제 세계의 그래프 크기는 계속 증가하고 있어 이 가정은 더 이상 실용적이지 않다: 실제 세계의 그래프 크기는 종종 GPU 메모리(예: NVIDIA H100의 80GB)나 심지어 단일 머신의 메인 메모리 용량(예: 256GB)을 초과한다. 예를 들어, 3층 GAT을 학습하는 경우 yahoo-web 그래프는 약 1.4B 노드와 6.6B 엣지로 구성되어 있으며, 학습에 필요한 노드 특성, 그래프 위상 및 중간 결과를 포함하여 약 1.5TB가 필요하다.

이 도전 과제를 해결하기 위해 저장기반 접근 방식을 이용한 GNN 학습이 연구되고 있다. 이 접근방식은 최신 고성능 외부 저장장치(NVMe SSD 등)를 활용하여 전체 그래프 위상과 노드 특성을 외부 저장장치에 저장하고 필요한 경우만 메인 메모리로 로드한다. 저장기반 GNN 학습은 다음과 같은 두 가지 과정으로 구성된다(Figure 1 참조):

(1) 데이터 준비: (i) 저장 장치에 저장된 그래프를 순회하여 GNN 학습을 위해 필요한 대상 노드의 이웃 노드를 찾고, (ii) 그들의 특성을 저장장치에서 메인 메모리로 수집하고, (iii) 이를 GPU로 전송한다.

(2) 계산: (iv) 전방 전파(즉, 예측)와 (v) 후방 전파(즉, 손실 및 기울기 계산)를 GPU에서 수행한다.

최신 GPU의 고성능 덕분에 계산 단계는 크게 가속되었지만, 데이터 준비 단계는 저장장치와 메인 메모리 간 I/O로 인해 전체 과정에서 중요한 병목 현상이 될 수 있다. 기존 연구들은 데이터 준비 단계를 개선하는 데 주력하여 그 잠재력을 보여주었다.

그러나 이들 방법에도 여전히 개선할 여지가 있다. Ginex와 GNNDrive라는 최고 수준의 저장기반 GNN 학습 방법에 대한 사전 실험을 통해 데이터 준비 단계의 시간 비율을 분석했다. 구체적으로, GCN과 GraphSAGE(SAGE) 두 가지 GNN 모델을 twitter-2010(TW), ogbn-papers100M(PA), com-friendster(FR) 세 개의 실제 세계 그래프 데이터셋에서 학습시켰다. Figure 2(a)에 보여진 것처럼 데이터 준비 단계가 전체 학습 과정을 지배한다 (즉, 총 실행 시간의 최대 96%). 깊이 있는 분석을 위해 각각의 개별 I/O 크기 역시 측정했다. Figure 2(b)는 저장장치 I/O의 크기 분포를 보여주며 많은 수의 작은 I/O가 발생하고 있으며, 몇몇 큰 I/O만 존재한다. 이러한 많은 작은 I/O는 GPU 활용률을 포함한 컴퓨팅 자원의 효율성을 크게 저하시키는 것으로 나타났다(Figure 2(c) 참조).

이 현상은 실제 세계 그래프가 파워 법도 분포를 가진다는 점에서 비롯된다. 즉, 대부분의 노드는 몇 개의 엣지만(즉, 이웃)을 가지며 소수의 노드만 매우 많은 엣지를 가진다. 따라서 GNN 학습에 필요한 이웃 노드 수가 대부분의 경우 매우 작을 가능성이 높다. 그러나 기존 저장기반 GNN 학습 방법들은 이러한 중요한 특성을 간과하고 있다. 그들은 단지 메인 메모리에서 캐시된 데이터를 재사용할 확률을 높이는 것(즉, 캐시 히트 비율)에 초점을 맞추고 있으며 GNN 학습이 필요할 때마다 몇 개의 노드만 저장장치에서 읽어들이므로 많은 수의 작은 I/O를 생성한다. 예를 들어, 전체 그래프를 분할하고 같은 파티션 내에서 대상 노드를 선택함으로써 샘플링된 노드들의 국소성을 강화하여 캐시 히트 비율을 높이려는 방법이 있다. 하지만 이러한 접근법들은 여전히 많은 수의 작은 I/O 처리라는 근본적인 도전 과제를 해결하지 못한다.

이 도전 과제는 작은 I/O들을 합치고 함께 처리함으로써 해결할 수 있다. 그러나 단순히 I/O 유닛 크기를 늘리는 것은 최적의 해결책은 아니며, 필요 없는 데이터(즉, 대상 노드와 관련없는 노드/엣지)가 각각의 I/O에 포함되어 메인 메모리 공간을 낭비할 수 있다.

우리의 연구: 이 논문에서는 이러한 저장기반 GNN 학습에서의 근본적인 도전 과제를 해결하기 위해 새로운 프레임워크 AGNES를 제안한다. AGNES는 (i) 저장, (ii) 메모리, 그리고 (iii) 연산 계층으로 구성되며 각각이 서로 밀접하게 상호작용하여 효율적인 I/O 처리를 수행한다. 우리는 블록 단위의 저장장치 I/O 처리와 새로운 데이터 레이아웃을 통해 작은 크기의 I/O 수를 줄이는 방법을 제시하며, 고성능 저장 장치의 I/O 대역폭(I/O-BW)을 최대한 활용한다. 또한 각각의 I/O 효율성을 더욱 향상시키기 위해(즉, 캐시 히트 비율), 실제 세계 그래프의 특징에 기반한 간단하면서도 효과적인 전략인 하이퍼 배치 처리를 제안한다. 이 방법은 각 반복에서 블록 내에서 GNN 학습을 위해 필요한 데이터를 가능한 한 많이 수집하여 동시에 처리한다.

기여: 본 연구의 주요 기여는 다음과 같다.

• 관찰: 현존하는 연구들은 저장기반 GNN 학습에서 중요한 도전 과제인 많은 수의 작은 I/O 처리에 대해 간과하고 있다는 점을 발견했다.

• 프레임워크: 새로운 프레임워크 AGNES를 제안하여 블록 단위의 저장장치 I/O 처리와 하이퍼 배치 처리를 통해 이 도전 과제를 해결한다.

• 평가: 다섯 개의 실제 세계 그래프에 대한 종합적인 실험을 통해 AGNES가 최고 수준의 저장기반 GNN 학습 방법보다 크게 우수함을 보여주었다. 구체적으로, AGNES는 최대 4.1배 빠르게 학습을 완료하며 I/O-BW 활용률도 최대 4.5배 더 높은 것으로 나타났다.

관련 연구

이 섹션에서는 기존의 GNN 학습 접근 방식을 검토하고 그와 우리의 작업 간의 관계를 설명한다.

저장기반 접근법: 최근, 우리 주요 연구 대상인 저장기반 GNN 학습 방법들이 연구되고 있다. 이러한 접근방식은 외부 저장장치를 활용하여 대규모 그래프 학습을 수행하며 전체 그래프 위상과 노드 특성을 외부 저장장치에 저장하고 필요한 부분만 메인 메모리로 로드한다. Ginex는 상태의 최고 수준인 저장기반 GNN 학습 방법으로, 노드 특성 벡터를 위한 캐싱 메커니즘을 사용하여 I/O 혼잡 문제를 효과적으로 해결하며 단일 머신에서 억단위 그래프 데이터셋을 처리할 수 있다. MariusGNN은 그래프를 여러 파티션으로 분할하고, 이들을 메인 메모리에 버퍼링하여 샘플링 결과를 재사용함으로써 저장장치 I/O 병목 현상을 완화한다. 또한 다중 단계 샘플링의 중복성을 최소화하는 데이터 구조와 디스크 기반 학습을 위한 두 가지 수준의 미니배치 교체 정책을 채택한다. GNNDrive는 다양한 단계를 지원하여 샘플 단계에서 메모리 경쟁을 완화하기 위해 버퍼 관리를 수행하고 비동기 특성 추출을 통해 메모리 사용 문제를 해결하며 저장장치 I/O 병목 현상을 완화한다. OUTRE는 분할 기반 배치 구성과 히스토리 임베딩을 사용하여 샘플링 기반 GNN 학습에서 이웃 중복성 및 시대적 중복성을 감소시킨다.

기존의 저장기반 접근법들은 단일 머신에서 대규모 그래프 학습에 대한 큰 가능성을 보여주었지만, 여전히 데이터 준비 단계에서 중요한 병목 현상을 겪고 있다(Figure 2 참조). 이는 자연스럽지만 중요하게 간과된 도전 과제를 해결하지 못한 것이다. 우리의 연구는 이러한 문제에 대한 첫 번째 상세한 접근 방식이다.

다른 접근법: 저장기반 접근법 외에도 메모리 기반 및 분산 시스템 기반 접근법이 연구되고 있다. 메모리 기반 접근법은 그래프 데이터나 노드 특성을 메인 메모리에 저장하여 GPU 메모리 크기를 초과하는 대규모 그래프를 처리한다. 예를 들어, PyG는 CPU와 GPU를 동시에 사용하여 GNN 모델의 학습 속도를 향상시키는 방법을 채택하고 있다. DGL은 메인 메모리에서 GPU로 데이터 전송을 빠르게 하기 위해 제로-카피 접근법을 채택한다. PaGraph는 CPU와 GPU 간의 데이터 전송 병목 현상을 해결하기 위해 자주 액세스되는 고도 노드를 GPU 메모리에 캐싱한다. 반면, 분산 시스템 기반 접근법은 분산 시스템의 풍부한 컴퓨팅 능력과 메모리 용량을 활용하여 단일 머신이 처리할 수 없는 크기의 그래프에서 GNN 모델을 학습한다. AliGraph는 각 머신에 노드 데이터를 로컬로 캐싱함으로써 네트워크 통신 비용을 줄이는 방법을 채택하고 있다. DistDGL은 주어진 그래프를 최소 자르기 분할 알고리즘을 사용하여 분할하여 단지 네트워크 통신 비용만 줄이지 않고 그래프 분할과 각 분할에서 생성된 미니배치의 균형을 맞추는 방법을 채택한다. DistDGLv2는 다중 수준 분할 알고리즘과 비동기 미니배치 생성 파이프라인을 사용하여 DistDGL을 개선한다.

그러나 메모리 기반 접근법은 메인 메모리 용량(즉, 덜 스케일링)을 초과하는 대규모 그래프를 처리할 수 없으며 분산 시스템 기반 접근법은 여러 머신에서 결과를 집계하기 위한 상당한 머신 간 통신 오버헤드와 고성능 분산 시스템 유지에 필요한 비용과 노력(즉, 덜 효율적이고 비싸)을 필요로 한다.

제안된 프레임워크: AGNES

이 섹션에서는 저장기반 GNN 학습을 위한 새로운 프레임워크를 제안한다. 이 프레임워크는 Accelerating storage-based training for Graph NEural networkS(AGNES)로 명명된다.

사전 지식

이 논문에서 사용된 표기법은 Table 1에 설명되어 있다.

그래프 신경망(GNNs): GNNs는 그래프 구조를 기반으로 노드의 임베딩 벡터를 표현하기 위해 메시지 패싱 메커니즘을 활용한다. 더 구체적으로, GNN의 각 층은 집합과 업데이트 두 단계로 구성된다. 집합 단계에서는 각 노드에 대해 이웃 노드들의 임베딩 벡터를 대상 노드의 임베딩으로 집합한다. 업데이트 단계에서는 집합된 임베딩을 비선형 함수와 함께 완전 연결 층을 통해 통과시킨다. 임베딩 $`\textbf{h}_{v}\in\mathbb{R}^{d}`$의 업데이트는 (1)에서 표현된다:

\begin{equation}
    \textbf{h}^{(i+1)}_{v}=\psi(\phi(\textbf{h}^{(i)}_{v'} | v' \in N(v), \textbf{h}^{i}_{v}))
\end{equation}

여기서 $`N(v)`$는 노드 $`v`$의 이웃 노드 집합을 나타내며, $`\psi(\cdot)`$와 $`\phi(\cdot)`$는 각각 집합과 업데이트 함수이다. 여러 층을 중첩하면 GNN 모델은 대상 노드의 $`k`$-hop 이웃 노드들의 정보를 그 임베딩에 반영할 수 있다.

GNNs의 미니배치 학습: 한편, 실제 세계 그래프와 그 노드 특성을 저장하는 것은 종종 수백 GB 또는 심지어 수 테라바이트(TB)가 필요하며 메인 메모리 용량을 초과한다. 단일 머신에서 이러한 큰 그래프를 처리하기 위해 저장기반 GNN 학습 방법은 전체 그래프를 외부 저장장치(NVMe SSD 등)에 저장하고 각 반복마다 그 전체 그래프에서 GPU 메모리로 로드할 수 있는 노드의 부분집합(즉, 미니배치)만 처리한다.

그러나 GNN 학습에서도 목표 노드와 그 이웃 노드들의 특성 벡터를 수집하는 것은 여전히 많은 메모리를 필요로 할 수 있다. 이러한 메모리 문제를 해결하기 위해 기존의 저장기반 방법들은 (i) 대상 노드와 관련된 이웃 노드 중 일부만 무작위로 샘플링하고 (ii) 이를 GNN 학습에서 그 임베딩을 업데이트하는 데 사용한다.

📊 논문 시각자료 (Figures)

감사의 말씀