전기및전자공학부 한동수교수 연구팀, 일반 소비자용 GPU로 분산 환경에서 AI 모델 학습 가속화 기술 개발

<(좌측부터) 한동수 교수, 임휘준 박사, 예준철 박사과정 사진>

전기및전자공학부 한동수 교수 연구팀이 일반 소비자용 GPU를 활용해 네트워크 대역폭이 제한된 분산 환경에서 AI 모델 학습을 획기적으로 가속하는 기술을 개발했다.

최신의 AI 모델을 학습하기 위해서는 수천만 원대의 고성능 GPU와 고속 전용 네트워크 등 고비용 인프라가 필요하다.

이로 인해 학계 및 중소기업의 대다수 연구자는 저렴한 일반 소비자용 GPU에 의존해 모델을 학습하고 있지만, 네트워크 대역폭 제약으로 효율적인 모델 학습에 어려움을 겪고 있는 실정이다.

<그림 1. 기존 저비용 분산 딥러닝 환경에서의 문제>

한동수 교수 연구팀은 이러한 문제를 해결하기 위해 StellaTrain이라는 분산 학습 프레임워크를 개발하였다.

StellaTrain은 저비용 GPU에서 모델 학습을 가속하기 위해 CPU와 GPU를 함께 활용하는 파이프라인을 도입했고, 네트워크 환경에 따라 배치 크기 및 압축률을 동적으로 결정하는 알고리즘을 통합 적용해 고속 전용 네트워크 없이도 다중 클러스터 및 다중 노드 환경에서 빠른 모델 학습을 달성할 수 있게 했다.

StellaTrain은 학습 파이프라인을 최적화하여 GPU 활용률을 극대화하기 위해 그래디언트 압축과 최적화 과정을 CPU로 오프로딩하는 전략을 채택했다.

특히 CPU에서 효율적으로 동작하는 희소 최적화 기술과 캐시 인식 기반의 그래디언트 압축 기술을 새롭게 개발해 적용함으로써, CPU의 작업이 GPU의 연산과 중첩되는 끊임없는 학습 파이프라인을 구현했다.

또한 네트워크 상황에 따라 배치 크기와 압축률을 실시간으로 조절하는 동적 최적화 기술을 적용해 제한된 네트워크 환경에서도 높은 GPU 활용률을 달성했다.

<그림 2. StellaTrain 학습 파이프라인 개요도>

이를 통해 StellaTrain은 저비용 멀티 클라우드 환경에서 분산 모델 학습 속도를 크게 향상시켜, 기존 PyTorch DDP 대비 최대 104배의 성능 향상을 달성하였다.

한동수 교수 연구팀은 이번 연구를 통해 고가의 데이터센터급 GPU와 고속 네트워크 없이도 AI 모델을 효율적으로 학습시킬 수 있는 길을 열었으며, 이는 학계와 중소기업 등 자원이 제한된 환경에서 AI 연구 및 개발을 가속하는 데 큰 도움이 될 것으로 기대된다.

한동수 교수는 “KAIST가 AI 시스템 분야에서 우리나라의 리더십을 발휘하며 두각을 나타내고 있다”고 강조했다.

그는 “거대 IT 기업의 전유물로 여겨지던 대규모 언어 모델(LLM) 학습을 보다 저렴한 컴퓨팅 환경에서 구현할 수 있도록 적극적인 연구를 진행하겠다”며, “이번 연구가 그 목표를 향한 중요한 초석이 되기를 기대한다”고 덧붙였다.

이번 연구에는 KAIST 임휘준 박사와 예준철 박사과정, 그리고 UC Irvine의 Sangeetha Abdu Jyothi 교수가 참여했으며, 본 연구 결과는 8월 4일부터 8일까지 호주 시드니에서 열린 컴퓨터 네트워킹 분야 최고 권위의 국제 학회인 ACM SIGCOMM 2024에서 발표되었다 (논문명: Accelerating Model Training in Multi-cluster Environments with Consumer-grade GPUs).

한편, 한동수 교수 연구팀은 최근 오스트리아 빈에서 열린 ICML 2024에서 MoE (Mixture of Experts) 모델 학습을 가속할 수 있는 프레임워크인 ES-MoE를 발표하는 등 AI 시스템 분야에서 지속적인 연구 성과를 내고 있다.

한동수 교수 연구팀은 GPU 메모리 한계를 극복하고 대규모 MoE 모델 학습의 확장성과 효율성을 크게 향상시켜, 단 4개의 GPU로 150억 파라미터 규모의 언어 모델을 미세 조정할 수 있게 하는 성과를 거뒀다.

이는 한정된 컴퓨팅 자원으로도 대규모 AI 모델을 효과적으로 학습할 가능성을 열었다는 점에서 큰 의의가 있다.

<그림 3. ES-MoE 프레임워크 개요도>