AI 모델의 규모가 커질수록 학습에 필요한 GPU 자원과 비용이 늘어나는 가운데, 이를 해결하기 위한 오픈소스 기술 개발이 주목받고 있다. 모델의 성능뿐 아니라 제한된 컴퓨팅 자원으로 얼마나 효율적으로 학습할 수 있는지가 대규모 AI 개발의 핵심 과제로 떠오르고 있다. 

앨런 AI 연구소(Ai2)는 최근 LLM 학습 효율을 높인 오픈소스 프레임워크 '올모 코어 3(Olmo-core 3)'를 공개했다.

기존 학습 시스템을 재설계해 470억개 매개변수 규모의 모델에서 GPU당 처리량을 약 2.7배 높였다. 수조개 매개변수 규모의 모델까지 확장할 수 있도록 설계해 대규모 AI 모델 개발에 필요한 컴퓨팅 비용과 에너지 부담을 줄이는 데 초점을 맞췄다.

올모 코어 3의 핵심은 전문가 혼합(MoE) 모델을 위한 새로운 학습 시스템이다. MoE는 여러 전문 신경망을 두고 입력 토큰마다 일부 전문가만 선택해 처리하는 방식이다. 전체 매개변수 규모가 크더라도 실제 연산에 참여하는 매개변수를 제한할 수 있어, 모든 매개변수를 매번 사용하는 일반적인 모델보다 연산 효율을 높일 수 있다.

하지만 MoE 모델도 규모가 커지면 효율성이 떨어질 수 있다. 모델의 전문가를 여러 GPU에 나눠 배치하는 과정에서 토큰을 적절한 전문가에게 전달하고 작업을 조정하는 데 추가 통신 비용이 발생하기 때문이다. 학습 시스템이 비효율적으로 설계되면 일부 전문가만 연산에 참여하는 장점이 상쇄될 수 있다.

Ai2는 이를 해결하기 위해 기존의 '완전 샤딩 데이터 병렬 처리(FSDP·Fully Sharded Data Parallelism)' 방식에서 '분산 데이터 병렬 처리(DDP)' 기반 시스템으로 전환했다. 기존 시스템은 작은 학습 배치마다 모델 가중치를 모았다가 다시 분산하는 과정을 반복했다. 올모 코어 3는 전문가를 GPU에 상주시키고 필요한 데이터를 해당 GPU로 전달해 가중치를 반복적으로 수집하는 부담을 줄였다.

성능 개선은 사전 벤치마크를 통해 입증됐다. 엔비디아 'B300' GPU 8개를 활용한 테스트에서 470억개 매개변수를 가진 MoE 모델은 GPU당 초당 5만2000개 토큰을 처리해 기존 FSDP 기반 시스템보다 약 2.7배 높은 효율을 기록했다. 또 전문가 수를 늘려 전체 규모를 470억개로 확장하는 실험에서도 학습 처리량 감소 폭을 5% 미만으로 묶으며 안정적인 확장성을 증명했다.

아울러 대규모 모델을 여러 GPU에 효율적으로 분산하는 기술과 함께 행 단위 전문가 병렬화, GPU 상주 라우팅, 그룹형 행렬 곱셈(GEMM) 등을 결합해 데이터 이동 비용을 최소화했다. 나아가 저정밀도 연산 형식인 MXFP8을 적용해 학습 처리량을 약 21% 늘리고 최대 활성 메모리 사용량을 줄이는 성과도 거뒀다.

Ai2는 이번 프레임워크가 개별 기술의 성능뿐 아니라 학습 과정 전체의 균형을 고려해 설계되었으며, 모델 가중치뿐만 아니라 학습에 사용되는 시스템과 기술 전체를 공개한다는 점에서 의의가 크다고 설명했다.

연구자와 개발자는 올모 코어 3를 활용해 자체 MoE 모델을 학습하고, 하드웨어 환경에 맞춰 시스템을 조정하거나 전문가 라우팅과 병렬화 기법을 실험할 수 있다. Ai2는 관련 코드와 기술 보고서를 깃허브를 통해 공개해 대규모 AI 모델 학습 인프라에 대한 접근성을 높인다는 방침이다.

박찬 기자 cpark@aitimes.com