"AI 학습에 엔비디아 GPU 필수 아냐"…카이스트, 대안기술 `스텔라트레인` 개발

실시간
뉴스

일반

"AI 학습에 엔비디아 GPU 필수 아냐"…카이스트, 대안기술 '스텔라트레인' 개발

디지털데일리 발행일 2024-09-19 09:07:42

이건한 기자

URL복사

(왼쪽부터) 카이스트 한동수 교수, 임휘준 박사, 예준철 박사과정. [ⓒ 카이스트]

[디지털데일리 이건한 기자] 한국과학기술원(KAIST·카이스트) 연구진이 일반 소비자용 그래픽처리장치(GPU)와 네트워크 대역폭이 제한된 분산 환경에서도 인공지능(AI) 모델 학습 속도를 100배 이상 가속할 수 있는 기술을 개발했다.

19일 카이스트에서 따르면 한동수 전기및전자공학부 교수 연구팀은 기존의 고가 GPU 다수가 필요했던 AI 연구 인프라 문제를 해결하기 위해 '스텔라트레인(StellaTrain)'으로 명명된 분산학습 프레임워크를 개발했다. 기존에는 AI 모델을 학습하기 위해 개당 수천만원에 달하는 고성능 서버용 GPU(엔비디아 H100) 여러 대와 이들을 연결하기 위한 400Gbps급 고속 네트워크를 가진 고가 인프라가 필요했다. 하지만 이 같은 비용은 소수의 거대 IT 기업을 제외하면 대부분의 기업과 연구자들은 감당하기 어려운 수준이다.

하지만 저가 GPU를 사용할 경우 작은 GPU 메모리와 네트워크 속도 제한으로 인해 대규모 AI 모델 학습 때 속도가 수백배 느려지는 한계가 있었다. 한 교수팀이 개발한 스텔라트레인은 중앙처리장치(CPU)와 GPU를 병렬로 활용해 학습 속도를 높이고 네트워크 속도에 맞춰 데이터를 효율적으로 압축 및 전송함으로써 이런 한계를 극복한다.

특히 학습을 작업 단계별로 CPU와 GPU가 나눠 병렬적으로 처리할 수 있는 새로운 파이프라인 기술을 도입해 연산 자원의 효율이 극대화했다. 또한 원거리 분산 환경에서도 GPU 연산 효율을 높이기 위해 AI 모델별GPU 활용률을 실시간으로 모니터링해 모델이 학습하는 샘플의 개수(배치 크기)를 동적으로 결정하고, 변화하는 네트워크 대역폭에 맞춰 GPU 간 데이터 전송을 효율화하는 기술을 개발했다. 연구 결과 스텔라트레인 기술을 사용하면 기존의 데이터 병렬 학습에 비해 최대 104배 빠른 성능을 낼 수 있는 것으로 나타났다.

한 교수는 "이번 연구가 대규모 AI 모델 학습을 누구나 쉽게 접근할 수 있게 하는 데 큰 기여를 할 것"이라며 "앞으로도 저비용 환경에서 대규모 AI 모델을 학습할 수 있는 기술 개발을 계속할 계획"이라고 말했다.

이번 연구는 카이스트 임휘준 박사, 예준철 박사과정 학생, 미국 UC어바인의 산기타 압두 조시 교수와 공동으로 진행됐으며, 연구성과는 지난 8월 호주 시드니에서 열린 연례 학숙대회 'ACM 시그콤(ACM SIGCOMM) 2024'에서 발표됐다.

한편, 한 교수 연구팀은 지난 7월에도 GPU 메모리 한계를 극복하는 기술을 개발한 바 있다. 기존에 32~64개 GPU가 필요했던 150억 파라미터 규모의 언어 모델을 단 4개의 GPU만으로도 학습할 수 있도록 한 기술이다. 임휘준 박사와 김예찬 연구원이 해당 논문에 참여했으며, 오스트리아 빈에서 열린 AI 분야 최고 권위 학회인 ICML에서 발표했다.