GPU 4개만으로도 AI 학습 시킨다

KAIST, AI 학습 기술 '스텔라트레인' 개발
기존 방법으로는 최대 42개 GPU 필요해
기존 병렬 AI 학습 기술보다 104배 단축

한국과학기술원(KAIST) 전기및전자공학부 한동수 교수

[파이낸셜뉴스] 한국과학기술원(KAIST) 전기및전자공학부 한동수 교수팀이 일반 소비자용 GPU 4개 만으로도 150억 파라미터 규모의 거대언어모델(LLM)을 학습시킬 수 있는 기술을 개발했다. 일반적인 방법으로 이 정도의 LLM을 학습시키기 위해서는 32~42개 GPU가 필요하다. 이를 통해 AI 학습에 필요한 최소 GPU 대수를 8분의 1~16분의 1로 낮출 수 있게 됐다.

19일 KAIST에 따르면, 연구진은 고속 네트워크 없이도 AI 모델을 효율적으로 학습할 수 있는 '스텔라트레인(StellaTrain)'이라는 분산 학습 프레임워크를 개발했다.

이 기술은 개당 수천만원에 달하는 고성능 H100에 비해 10~20배 저렴한 소비자용 GPU를 사용한다.

또 고속의 전용 네트워크 대신 대역폭이 수백에서 수천 배 낮은 일반 인터넷 환경에서도 AI 학습을 가능하게 한다.

스텔라트레인 기술은 CPU와 GPU를 병렬로 활용해 학습 속도를 높인다. 또 네트워크 속도에 맞춰 데이터를 효율적으로 압축 및 전송하는 알고리즘을 적용해 고속 네트워크 없이도 여러 대의 저가 GPU를 이용해 빠른 학습을 가능하게 했다.

특히, 학습을 작업 단계별로 CPU와 GPU가 나눠 병렬적으로 처리할 수 있는 새로운 파이프라인 기술을 도입해 연산 자원의 효율을 극대화했다.

이와함께 원거리 분산 환경에서도 GPU 연산 효율을 높이기 위해, AI 모델별 GPU 활용률을 실시간으로 모니터링해 모델이 학습하는 샘플의 개수를 동적으로 결정하고, 변화하는 네트워크 대역폭에 맞춰 GPU 간의 데이터 전송을 효율화했다.

이 기술을 테스트한 결과, 기존의 데이터 병렬 학습에 비해 최대 104배 빠른 성능을 냈다.

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

[e-mail] 김만기 기자