韓-美, "가정용 GPU로 기존 104배 넘는 초고효율 AI학습 가속기술 개발"

PC방이나 가정용 GPU를 이용해 고속 네트워크 없이도 AI학습 성능을 최대 104배까지 끌어 올릴 수 있는 획기적인 기술이 개발됐다.

KAIST는 전기및전자공학부 한동수 교수 연구팀이 일반 소비자용 GPU로 네트워크 대역폭이 제한된 분산 환경에서 AI 모델 학습을 혁신적으로 가속할 수 있는 기술을 개발했다고 19일 밝혔다.

기업과 연구자들이 고가의 데이터센터급 GPU(엔비디아 H100)나 고속 네트워크 없이도 AI 모델을 효율적으로 학습할 수 있는 길이 열린 셈이다.

기존에는 AI 모델 학습을 위해 개당 수천만 원에 달하는 고성능 서버용 GPU(엔비디아 H100) 여러 대와 이들을 연결하기 위한 400Gbps급 고속 네트워크 등 고가 인프라가 필요했다.

연구팀은 이 같은 문제 해결을 위해 '스텔라트레인(StellaTrain)'이라는 분산 학습 프레임워크를 개발했다. 고성능 H100 대비 10~20배 저렴한 GPU를 활용하고 고속 전용 네트워크 대신 대역폭이 수백에서 수천 배 낮은 일반 인터넷 환경에서도 효율적인 분산 학습이 가능하도록 알고리즘을 짰다.

CPU와 GPU를 병렬로 연결해 학습 속도도 높였다. 또 네트워크 속도에 맞춰 데이터를 효율적으로 압축 및 전송하는 알고리즘을 적용, 고속 네트워크 없이도 빠른 학습이 가능하도록 설계했다.

특히, 학습을 작업 단계별로 CPU와 GPU가 나눠 병렬 처리하도록 새로운 파이프라인 기술도 도입했다.

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

연구 결과, 스텔라트레인 기술을 사용하면 기존의 데이터 병렬 학습에 비해 최대 104배 빠른 성능을 낼 수 있는 것으로 나타났다.

임휘준 박사는 "스텔라트레인 기술을 사용하면 기존 데이터 병렬 학습 속도 대비 최대 104배 빠른 성능을 낼 수 있다"고 말했다.

KAIST 연구팀이 개발한 스텔라트레인 프레임워크. 네트워크 상황에 따라 압축률을 실시간 조절하는 동적 최적화 기술을 적용했다. 그림은 프레임워크 흐름도.(그림=KAIST)

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

연구는 KAIST 임휘준 박사, 예준철 박사과정 학생, UC 어바인 산기타 압두 조시(Sangeetha Abdu Jyothi) 교수와 공동으로 진행됐다.

연구 성과는 지난 8월 호주 시드니에서 열린 'ACM SIGCOMM 2024'에서 발표됐다.

한편, 한동수 교수 연구팀은 지난 7월 GPU 메모리 한계를 극복한 소수의 GPU로 거대 언어 모델을 학습하는 새로운 기술도 발표했다.

이 연구는 최신 거대 언어 모델의 기반이 되는 전문가 혼합형(Mixture of Expert) 모델을 제한된 메모리 환경에서도 효율적인 학습을 가능하게 한다.

기존에 32~64개 GPU가 필요한 150억 파라미터 규모의 언어 모델을 단 4개의 GPU만으로도 학습할 수 있다.