GPU 한 대로 95배 빠르게 그래프 분석 AI ‘FlexGNN’ 개발

KAIST 연구진이 단 한 대의 GPU 서버만으로도 대규모 GNNGraph Neural Network) 모델을 최고속 학습할 수 있는 세계 최고 성능의 소프트웨어 기술 개발에 성공했다.

KAIST는 전산학부 김민수 교수 연구팀이 여러 대의 GPU 서버를 활용하는 기존 방식과 달리 한 대의 GPU 서버에서 대규모 풀(full) 그래프 AI 모델을 빠르게 학습하고 추론할 수 있는 GNN 시스템 ‘FlexGNN(플렉스지엔엔)’을 개발했다고 13일 밝혔다. FlexGNN은 기존 기술 대비 학습 속도를 최대 95배 향상한다.

최근 기후, 금융, 의료, 제약, 제조, 유통 등 다양한 분야에서는 데이터를 정점과 간선으로 구성된 그래프 형태로 변환해 분석 및 예측하는 사례가 증가하고 있다.

전체 그래프를 모두 학습에 활용하는 풀 그래프 방식이 더욱 우수한 정확도를 보이지만 학습 과정에서 대규모의 중간 데이터(intermediate data)가 발생해 메모리 부족 현상이 빈번히 발생하고 여러 서버 간의 데이터 통신으로 인해 학습 시간이 길어지는 한계가 있었다.

연구팀이 개발한 FlexGNN은 이러한 문제를 극복하기 위해 여러 대의 GPU 서버 대신 단일 GPU 서버에서 SSD(솔리드 스테이트 드라이브)와 메인 메모리를 활용한 최적의 AI 모델 학습을 수행한다.

특히 데이터베이스 시스템의 질을 최적화시키는 AI 퀴리 최적화 학습을 통해 GPU-메인 메모리-SSD 계층 간 모델 파라미터, 학습 데이터, 중간 데이터를 최적의 시점과 방식으로 계산을 시키는 새로운 학습 최적화 기술을 개발했다.

이를 통해 FlexGNN은 데이터 크기, 모델 규모, GPU 메모리 등 가용 자원 상황에 따라 유연하게 최적의 학습 실행 계획을 생성해 높은 자원 효율성과 학습 속도를 구현한다.

그 결과, 메인 메모리 용량을 훨씬 초과하는 데이터에 대해서도 GNN 모델을 학습하며 단일 GPU 서버에서도 최대 95배 빠르게 학습이 가능해졌다. 특히 기후 예측 등에서 슈퍼컴퓨터보다 정밀한 분석이 가능한 풀 그래프 AI 구현이 현실화됐다.

KAIST 김민수 교수는 “날씨 예측과 신소재 발견 등 복잡한 문제를 해결하는데 풀 그래프 GNN 모델이 활발히 활용되면서 관련 기술의 중요성이 점점 높아지고 있다”며 “FlexGNN이 그동안 어려움으로 남아 있던 그래프 AI 모델의 학습 규모와 속도 문제를 획기적으로 해결한 만큼, 다양한 산업 분야에 널리 활용되기를 기대한다”고 밝혔다.

이번 연구는 KAIST 전산학부 배정민 박사과정이 제1 저자로, 김민수 교수창업기업인 그래파이의 한동형 CTO가 제2 저자로 참여했으며, 김 교수가 교신저자를 맡았다.

연구 결과는 세계적 권위의 데이터마이닝 학술대회인 ‘ACM KDD’에서 지난 8월 5일에 발표됐다. FlexGNN 기술은 향후 그래파이의 그래프 DB 솔루션인 그래프온(GraphOn)에도 적용될 예정이다.