조 단위 MoE 모델을 향한 오픈 트레이닝 스택의 진화
대규모 언어 모델(LLM) 개발 비용과 전력 소비가 급증하면서, 학계 연구진과 중소 연구소가 첨단 모델 개발에 접근하기는 점차 어려워지고 있습니다. 전문가 혼합(MoE, Mixture of Experts) 구조는 토큰당 모델 전체가 아닌 일부 전문가만 활성화하여 연산 효율을 극대화하는 대안으로 부상했으나, 여전히 전체 가중치를 GPU 클러스터 메모리에 분산 유지하고 데이터를 정확한 전문가에게 라우팅하는 과정에서 막대한 통신 비용이 발생합니다. 모델 규모가 커질수록 이러한 통신 오버헤드가 MoE의 연산 효율 이점을 상쇄하는 문제가 있었습니다.
앨런 인공지능 연구소(Ai2)는 이러한 병목을 해결하기 위해 조 단위 파라미터 학습을 지원하는 오픈소스 인프라 올모코어 3(Olmo-core 3)를 공식 출시했습니다. 차세대 Olmo 모델의 핵심 기반이 되는 본 프레임워크는 대규모 MoE 학습 효율을 유지하면서도 연구 커뮤니티 누구나 수정하고 활용할 수 있는 완전 개방형 스택으로 설계되었습니다.
"올모코어 3는 계산 효율을 유지하면서 MoE 학습을 조 단위 파라미터 영역까지 확장하도록 설계되었습니다. 이는 차세대 Olmo의 핵심 시스템 중 하나이며, 새로운 모델 이면의 도구와 학습 인프라를 개방하려는 지속적인 노력의 일환입니다."
---
FSDP에서 DDP 기반으로: 아키텍처 재설계의 핵심
기존 올모코어의 MoE 구현은 완전 분할 데이터 병렬(FSDP) 방식을 사용하여 매 미니배치마다 모델 가중치를 모으고(Gather) 다시 샤딩(Reshard)하는 과정을 거쳤습니다. 그러나 올모코어 3는 분산 데이터 병렬(DDP) 기반 아키텍처로 전면 개편되었습니다. 전문가 모듈을 GPU 메모리에 상주시킨 채 입력 데이터만을 해당 전문가가 위치한 GPU로 라우팅함으로써 불필요한 가중치 통신을 원천적으로 제거했습니다.
- 입력 토큰 수신
- GPU 상주 라우터에서 전문가 매핑
- 데이터를 해당 GPU 전문가 버퍼로 전송
- Grouped GEMM 고속 연산 수행
이러한 구조적 전환은 학습 처리량의 비약적인 상승으로 이어졌습니다. 엔비디아 B300 GPU 8대를 사용한 470억(47B) 파라미터 MoE 모델 예비 벤치마크에서, 기존 FSDP 기반 스택은 GPU당 초당 19,400개의 토큰을 처리하는 데 그쳤으나 올모코어 3는 52,000 토큰/초/GPU를 기록하며 약 2.7배(270%) 향상된 처리 성능을 입증했습니다.
| 인프라 버전 | 병렬화 기반 | 47B MoE 처리량 (tokens/sec/GPU) | 성능 향상 배수 |
|---|---|---|---|
| 기존 Olmo-core | FSDP (가중치 재수집 방식) | 19,400 | 기준선 (1.0×) |
| Olmo-core 3 | DDP (전문가 상주 + 데이터 라우팅) | 52,000 | 약 2.7× |
---
대규모 클러스터 스케일링 및 통신 최적화 기법
올모코어 3는 대규모 GPU 클러스터 환경에서 MoE 모델을 효율적으로 분할하고 실행하기 위해 복합적인 최적화 기술을 도입했습니다.
1. 분산 상태 관리 및 메모리 분할
- 전문가 병렬화(Expert Parallelism): GPU 풀 전체에 걸쳐 전문가 모듈을 분산 배치하여, 단일 GPU가 전체 전문가 풀 중 일부만 유지하도록 구성합니다.
- 파이프라인 병렬화(Pipeline Parallelism): 모델의 여러 레이어를 GPU 그룹 단위로 분할하여 개별 GPU의 메모리 점유를 줄입니다.
- 분산 옵티마이저(Distributed Optimizer): 모든 GPU가 옵티마이저 상태 사본 전체를 갖는 대신, 옵티마이저 상태 데이터를 GPU 전반에 분산 저장합니다.
2. 라우팅 오버헤드 최소화
- 행 단위 전문가 병렬화(Rowwise Expert Parallelism): 라우팅된 데이터를 전문가의 입력 버퍼에 직접 배치하여 데이터 재배치 오버헤드를 줄입니다.
- GPU 상주 라우팅(GPU-resident Routing): 라우팅 메타데이터를 GPU 내부 메모리에서 직접 처리하여 CPU가 작업 큐를 전달할 때 정보 복사를 기다리지 않도록 최적화했습니다.
- 그룹화 일반 행렬 곱(Grouped GEMM): 다수의 소규모 전문가 연산을 단일 커널로 병합하여 GPU 연산 유닛 가동률을 극대화합니다.
3. MXFP8 저정밀도 포맷 도입
올모코어 3는 MXFP8 저정밀도 데이터 포맷을 공식 지원합니다. 엔비디아 B300 GPU 4대 기반 벤치마크에서 기존 고정밀 기준선인 BF16 대비 학습 처리량이 약 21% 향상되었으며, 피크 활성 메모리는 103 GiB에서 95 GiB로 감소했습니다. 연구팀은 성능 향상의 주요 원인이 단순 어텐션 연산보다는 피드포워드(FFN) 연산 가속 및 전문가 간 통신량 감소에 기인했다고 밝혔습니다.
---
1조~2.38조 파라미터 벤치마크와 실증적 발견
| Olmo-core (구 FSDP) | 19400 |
| Olmo-core 3 (DDP) | 52000 |
올모코어 3는 토큰당 활성 파라미터 수를 약 32억 개(3.2B)로 고정한 상태에서 전문가 수를 8개에서 128개로 확대하는 실험을 거쳤습니다. 그 결과 총 파라미터 용량이 4.6B에서 47B로 10배 이상 증가했음에도 학습 처리량 저하는 5% 미만에 머물렀습니다.
나아가 초거대 규모 확장을 입증하기 위해 엔비디아 B300 GPU 512대 클러스터에서 토큰당 58.36B 활성 파라미터를 가진 1.2조(1.2T) 파라미터 모델을 벤치마크했습니다. 이 구성에서 올모코어 3는 GPU당 최대 858 TFLOP/s의 연산 처리 효율을 달성했습니다. 또한 전문가 간 통신 라이브러리인 DeepEP v2를 적용한 용량 확장 테스트에서는 최대 2.38조(2.38T) 파라미터 설정까지 인프라 확장이 가능함을 검증했습니다.
인프라 구축 과정에서 규명된 주요 엔지니어링 맹점
Ai2는 기술 보고서를 통해 단순한 이론적 설계가 실제 학습 환경에서 실패할 수 있음을 보여주는 4가지 실증 결과를 공개했습니다.
- 토큰 게리맨더링(Token Gerrymandering): 부하 균등화를 유도하는 라우팅 스코어가 상승하더라도, 실제 작업 부하는 특정 전문가로 쏠리는 역설적 편향 현상이 관측되었습니다.
- 전문가 학습률 감축 무효성: 적은 토큰을 처리하는 전문가의 학습률(Learning rate)을 인위적으로 낮추는 방식은 테스트된 모델군에서 성능 개선에 기여하지 못했습니다.
- 입력값에 따른 연산 시간 변동: 행렬 크기가 동일하더라도 처리되는 실제 수치 데이터의 분포에 따라 GPU 연산 소요 시간이 달라지므로 정확한 성능 비교를 위해서는 동일한 입력 텐서 값 설정이 필수적입니다.
- 통신-연산 중첩(Overlap)의 역효과: 서로 다른 GPU 스트림에서 통신과 연산을 무조건 겹쳐 실행하는 것이 항상 유리하지 않았으며, 일부 설정에서는 전체 종단간 처리 속도를 오히려 저하시켰습니다.
---
오픈 AI 생태계의 전환점과 차세대 Olmo의 방향
과거 오픈 모델 진영의 주요 MoE 모델인 OlmoE(64 전문가)와 밀집(Dense) 모델이었던 Olmo 3를 거쳐, 차세대 Olmo는 올모코어 3 인프라를 바탕으로 한 대규모 MoE 구조로 훈련될 예정입니다. Ai2는 역대 가장 큰 데이터셋과 가장 긴 컨텍스트 윈도우를 갖춘 차세대 모델을 준비 중이며, 모델 가중치뿐 아니라 학습 프레임워크, 분산 훈련 레시피, 엔지니어링 시행착오 데이터를 깃허브(GitHub)와 기술 보고서를 통해 완전 개방했습니다.
이는 폐쇄형 프론티어 랩들이 독점해 온 초대형 MoE 인프라 기술을 학계와 독립 연구진이 직접 재현하고 수정할 수 있는 환경을 마련했다는 점에서 오픈소스 AI 생태계에 중요한 전환점을 제시합니다.