[논문 리뷰] Marius: Learning Massive Graph Embeddings on a Single Machine

Jason Mohoney, Roger Waleffe외 3명|arXiv (Cornell University)|2021. 01. 20.

Advanced Graph Neural Networks참고 문헌 33인용 수 10

한 줄 요약

Marius는 단일 머신에서 데이터 이동을 최적화하기 위해 파artition 캐싱과 버퍼 인지 데이터 순서 정렬(BETA)을 활용한 프레임워크를 제안하여, 최신 기술 대비 최대 10배 빠른 훈련 속도를 달성하며, 단지 16GB GPU와 64GB CPU 메모리로도 10억 개 이상의 간선을 가진 그래프와 550GB의 파라미터를 처리할 수 있도록 한다.

ABSTRACT

We propose a new framework for computing the embeddings of large-scale graphs on a single machine. A graph embedding is a fixed length vector representation for each node (and/or edge-type) in a graph and has emerged as the de-facto approach to apply modern machine learning on graphs. We identify that current systems for learning the embeddings of large-scale graphs are bottlenecked by data movement, which results in poor resource utilization and inefficient training. These limitations require state-of-the-art systems to distribute training across multiple machines. We propose Marius, a system for efficient training of graph embeddings that leverages partition caching and buffer-aware data orderings to minimize disk access and interleaves data movement with computation to maximize utilization. We compare Marius against two state-of-the-art industrial systems on a diverse array of benchmarks. We demonstrate that Marius achieves the same level of accuracy but is up to one order of magnitude faster. We also show that Marius can scale training to datasets an order of magnitude beyond a single machine's GPU and CPU memory capacity, enabling training of configurations with more than a billion edges and 550 GB of total parameters on a single machine with 16 GB of GPU memory and 64 GB of CPU memory. Marius is open-sourced at www.marius-project.org.

연구 동기 및 목표

• 단일 머신에서 대규모 그래프 임베딩 훈련의 핵심 문제인 데이터 이동을 해결한다.
• GPU 메모리 외부에서의 비정규적인 메모리 접근 패턴으로 인한 낮은 GPU 활용도 문제를 해결한다.
• I/O와 계산의 겹침을 최적화하여 단일 GPU와 CPU의 메모리 용량을 초월하는 그래프 임베딩 훈련을 가능하게 한다.
• 기존 산업용 시스템 대비 훈련 처리량을 크게 향상시키면서도 높은 정확도를 유지하는 시스템을 설계한다.
• I/O를 최소화하고 하드웨어 활용도를 극대화하기 위해 새로운 데이터 순서 정렬 방식(BETA)과 파이프라인 아키텍처를 개발한다.

제안 방법

• I/O 지연을 숨기기 위해 데이터 이동, 전송, 계산을 교차 실행하는 파이프라인형 훈련 아키텍처를 도입한다.
• 자주 액세스되는 노드 임베딩 파티션을 재사용함으로써 중복 디스크 읽기를 줄이기 위해 파티션 캐싱을 구현한다.
• 근접한 노드 파티션을 가진 간선을 그룹화함으로써 I/O를 최소화하는 버퍼 인지 데이터 순서 정렬(BETA)을 제안한다.
• 노드 임베딩은 유한한 지연을 허용하는 비동기 훈련을, 간선 임베딩은 동기 훈련을 사용하여 다양한 간선 유형을 지원한다.
• 효율적인 내적 곱 연산을 활용한 미니배치 확률적 경사 하강법(SGD)을 적용하여 계산 오버헤드를 감소시킨다.
• 버퍼 크기, 파티션 수, 지연 한계 등을 최적화하여 처리량과 임베딩 품질의 균형을 이룬다.

실험 결과

연구 질문

• RQ1 분산 시스템 없이도 단일 머신에서 10억 개 이상의 간선을 가진 그래프 임베딩 훈련을 스케일링할 수 있는가?
• RQ2 GPU 메모리 외부 훈련에서 I/O 오버헤드를 최소화해 고성능 GPU 활용도를 달성할 수 있는가?
• RQ3 대규모 그래프 임베딩 훈련에서 I/O 비용을 줄이고도 효율성을 유지할 수 있는 데이터 순서 정렬 전략은 무엇인가?
• RQ4 파티션 캐싱과 파이프라인화와 같은 시스템 수준 최적화가 정확도를 희생시키지 않고 훈련 처리량을 얼마나 향상시킬 수 있는가?
• RQ5 배치 크기, 버퍼 크기, 지연 한계와 같은 하이퍼파라미터가 훈련 성능과 임베딩 품질에 미치는 영향은 어떠한가?

주요 결과

• Marius는 DGL-KE의 10%와 PBG의 30% 미만에 비해 평균 70%의 GPU 활용도를 달성하여 하드웨어 활용도를 크게 향상시켰다.
• Freebase86m 데이터셋에서 Marius는 동일한 모델 정확도를 유지하면서 DGL-KE와 PBG 대비 최대 10배 빠른 속도를 기록했다.
• Marius는 단지 16GB GPU와 64GB CPU 메모리로도 10억 개 이상의 간선을 가진 그래프와 최대 550GB의 총 모델 파라미터를 처리할 수 있도록 스케일링된다.
• BETA 순서 정렬은 노드 임베딩에 대한 디스크 액세스를 최소화하여, 특히 파티션 수가 많은 경우에 매우 효과적이다.
• 대규모 그래프에서는 큰 배치 크기(약 10,000)가 정확도 손실 없이 처리량을 향상시키지만, 임계점을 초과하면 이점이 감소한다.
• 버퍼 용량을 제곱 관계로 증가시킬 경우 디스크 스왑 수가 급격히 감소하므로, 버퍼 크기 설정이 성능에 매우 중요한 요소가 된다.