앨런 AI 연구소(Ai2)가 대규모 언어 모델 학습용 오픈소스 프레임워크 ‘Olmo-core 3’를 공개했습니다. 특히 전문가 혼합(MoE) 모델의 학습 과정에서 GPU 간 데이터 이동과 작업 조정에 드는 부담을 줄이는 데 초점을 맞췄습니다.
MoE는 입력에 따라 여러 신경망 전문가 가운데 일부만 골라 처리하는 방식입니다. 전체 모델이 커도 매번 모든 매개변수를 계산하지 않는 장점이 있지만, 전문가를 GPU 여러 대에 나눠 놓으면 데이터 전달과 통신이 병목이 될 수 있습니다.
GPU에 전문가를 상주시킨다
Ai2는 기존 완전 샤딩 데이터 병렬 처리(FSDP) 기반 시스템에서 분산 데이터 병렬 처리(DDP) 기반 시스템으로 바꿨습니다. 기존 방식은 작은 학습 배치마다 모델 가중치를 모았다가 다시 나누는 과정을 반복했습니다. Olmo-core 3는 전문가를 GPU에 두고 필요한 데이터를 해당 GPU로 전달해 이 반복 부담을 줄입니다.
Ai2가 공개한 사전 벤치마크에서는 NVIDIA B300 GPU 8개로 47B 매개변수 규모의 MoE 모델을 학습했을 때 GPU당 초당 5만2,000개 토큰을 처리했습니다. Ai2는 이를 기존 FSDP 기반 시스템보다 약 2.7배 높은 처리량이라고 설명했습니다. 다만 기사에는 벤치마크의 세부 설정이나 독립적인 검증 결과가 제시되지 않았습니다.
확장성과 공개 범위
Ai2는 전문가 수를 늘려 모델 전체 규모를 47B로 확장한 실험에서 처리량 감소 폭을 5% 미만으로 유지했다고 밝혔습니다. 또 저정밀도 연산 형식인 MXFP8을 적용해 처리량을 약 21% 높이고 최대 활성 메모리 사용량을 줄였다고 설명했습니다.
프레임워크의 코드와 기술 보고서는 GitHub에 공개됐습니다. 연구자와 개발자는 이를 바탕으로 자체 MoE 모델을 학습하거나 하드웨어에 맞춰 시스템을 조정할 수 있습니다. 다만 실제 비용 절감 규모나 다른 GPU 환경에서의 성능은 제공된 내용만으로 확인되지 않습니다.
왜 중요한가
MoE 모델을 직접 학습하는 연구자나 개발자는 Olmo-core 3의 코드와 기술 보고서를 활용해 GPU 배치와 데이터 전달 방식을 시험해 볼 수 있습니다. 다만 2.7배 개선은 Ai2가 제시한 특정 B300 벤치마크 결과이며, 다른 하드웨어에서 같은 성능이나 비용 절감이 나타나는지는 확인되지 않았습니다.
