- 전기및전자공학부 한동수 교수 연구팀, 고가의 GPU와 고속의 전용 네트워크 없이 제한된 인프라에서 AI 모델 학습을 가속하는 ‘스텔라트레인(StellaTrain)‘ 프레임 워크 개발

- 고가의 데이터센터급 GPU와 고속 네트워크 없이도 효율적인 AI 모델 학습이 가능해져 기존 데이터 병렬학습에 비해 최대 104배 빠른 성능으로, 학계와 중소기업의 AI 연구 및 개발 가속화 기대

▲(왼쪽부터) KAIST 전기및전자공학부 한동수 교수, 임휘준 박사, 예준철 박사과정.(제공=KAIST)

KAIST 연구진이 고가의 데이터센터급 GPU나 고속 네트워크 없이도 AI 모델을 효율적으로 학습할 수 있는 기술을 개발했다. 이 기술을 통해 자원이 제한된 기업이나 연구자들이 AI 연구를 보다 효과적으로 수행할 수 있을 것으로 기대된다.

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

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

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

기존의 저가 GPU를 사용할 경우, 작은 GPU 메모리와 네트워크 속도 제한으로 인해 대규모 AI 모델 학습 시 속도가 수백 배 느려지는 한계가 있었다. 하지만 연구팀이 개발한 스텔라트레인 기술은 CPU와 GPU를 병렬로 활용해 학습 속도를 높이고, 네트워크 속도에 맞춰 데이터를 효율적으로 압축 및 전송하는 알고리즘을 적용해 고속 네트워크 없이도 여러 대의 저가 GPU를 이용해 빠른 학습을 가능하게 했다. 

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

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

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

이번 연구는 KAIST 임휘준 박사, 예준철 박사과정 학생, UC 어바인의 산기타 압두 조시(Sangeetha Abdu Jyothi) 교수와 공동으로 진행됐으며, 연구 성과는 지난 8월 호주 시드니에서 열린 ACM SIGCOMM 2024에서 발표됐다.

한편, 한동수 교수 연구팀은 2024년 7월 GPU 메모리 한계를 극복해 소수의 GPU로 거대 언어 모델을 학습하는 새로운 기술도 발표했다. 해당 연구는 최신 거대 언어 모델의 기반이 되는 전문가 혼합형(Mixture of Expert) 모델을 제한된 메모리 환경에서도 효율적인 학습을 가능하게 한다. 

이 결과 기존에 32~64개 GPU가 필요한 150억 파라미터 규모의 언어 모델을 단 4개의 GPU만으로도 학습할 수 있게 됐다. 이를 통해 학습의 필요한 최소 GPU 대수를 8배~16배 낮출 수 있게 됐다. 해당 논문은 KAIST 임휘준 박사와 김예찬 연구원이 참여했으며, 오스트리아 빈에서 열린 AI 분야 최고 권위 학회인 ICML에 발표됐다. 이러한 일련의 연구 결과는 자원이 제한된 환경에서도 대규모 AI 모델 학습이 가능하다는 점에서 중요한 의미를 가진다. 

해당 연구는 과학기술정보통신부 한국연구재단이 주관하는 중견연구사업 (RS-2024-00340099), 정보통신기획평가원(IITP)이 주관하는 정보통신·방송 기술개발사업 및 표준개발지원사업 (RS-2024-00418784), 차세대통신클라우드리더십구축사업 (RS-2024-00123456), 삼성전자의 지원을 받아 수행됐다.

그림 1.  기존 저비용 분산 딥러닝 환경에서는 네트워크 제약으로 인해 멀티 클러스터 및 하이브리드 클라우드 환경에서 학습 시 GPU 이용률이 저하되고, 학습 시간과 클라우드 비용이 증가하는 문제가 있다.

그림 2. 한동수 교수 연구팀은 StellaTrain과 ES-MoE 연구를 통해 GPU 수가 제한적이거나 고성능 GPU가 없는 상황에서도 저렴한 컴퓨팅 환경으로 AI 모델 학습을 가능하게 했다.

그림 3. 연구팀이 개발한 StellaTrain 프레임워크는 다중 클러스터 및 다중 노드 환경에서 GPU 활용률을 극대화하고 학습 시간을 단축하기 위해, CPU 기반 그래디언트 압축 및 최적화로 통신 오버헤드를 최소화하고, 네트워크 상황에 따라 배치 크기와 압축률을 실시간으로 조절하는 동적 최적화 기술을 적용했다.

그림 4. 저비용 GPU에 의존하는 다수의 연구자들은 GPU 메모리 용량 초과로 인하여, 다량의 expert 모델을 이용한 학습이 불가능하다. 본 연구팀이 ICML 2024에서 발표한 ES-MoE 프레임워크에서는 CPU 메모리 및 SSD를 계층적으로 사용하여 부족한 GPU 메모리를 보완할 수 있는 Expert Switching 기법을 제안한다. 이는 저비용 GPU에서도 expert 개수를 크게 늘릴 수 있는 확장성을 제공한다.