비싼 GPU 없이 AI 학습 성능 104배 높인 기술 개발

한동수 KAIST 전기및전자공학부 교수 연구팀
AI 학습 모델 가속화 기술‘스텔라트레인’등장

고가 데이터센터급 GPU나 고속 네트워크 없이도 인공지능(AI) 모델 학습 성능을 104배 높인 기술이 국내에서 개발됐다.

KAIST는 한동수 전기·전자공학부 교수 연구팀이 일반 소비자용 GPU를 활용해, 네트워크 대역폭이 제한된 분산 환경에서도 AI 모델 학습을 수십에서 수백 배 가속할 수 있는 기술을 개발했다고 19일 밝혔다.

현재 AI 모델을 학습하기 위해 개당 수천만 원에 달하는 고성능 서버용 GPU(엔비디아 H100) 여러 대와 이들을 연결하기 위한 400Gbps급 고속 네트워크를 가진 고가 인프라가 필요하다. 이에 소수 거대 IT 기업을 제외한 대부분의 기업과 연구자들은 비용 문제로 고가 인프라를 도입하기 어려웠다.

한동수 교수 연구팀은 엔비디아 고성능 H100 가격의 10분의 1에 불과한 소비자용 GPU를 활용해 고속 전용 네트워크의 수백에서 수천 배 낮은 저대역폭인 일반 인터넷 환경에서도 효율적인 분산 학습이 가능한 ‘스텔라트레인(StellaTrain)’ 분산 학습 프레임워크를 개발했다.

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

연구팀은 스텔라트레인 기술을 사용하면 기존의 데이터 병렬 학습에 비해 최대 104배 빠른 성능을 낼 수 있다고 밝혔다.

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

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

KAIST 임휘준 박사, 예준철 박사과정 학생, UC 어바인의 산기타 압두 조시(Sangeetha Abdu Jyothi) 교수가 공동 참여한 이번 연구 결과는 지난달 8월 호주 시드니에서 열린 국제 학술대회 ‘ACM SIGCOMM 2024’에서 발표됐다.