Ai2, 트릴리언 파라미터 MoE를 위한 오픈 트레이닝 스택 Olmo-Core 3 출시

Allen Institute for AI는 2026년 10월 1일 Olmo-core 3을 출시했으며, 이는 재설계된 오픈 혼합 전문가(MoE) 트레이닝 시스템을 중심으로 구축된 대형 언어 모델 개발 프레임워크의 업그레이드입니다.

Ai2는 Olmo-core 3가 계산 효율성을 유지하면서 MoE 트레이닝을 트릴리언 파라미터 규모로 확장하도록 설계되었으며, 차세대 Olmo의 핵심 시스템 중 하나라고 설명했습니다. 이번 출시에는 기술 보고서, 인터랙티브 데모, 그리고 오픈 코드가 함께 제공됩니다.

MoE 모델은 모든 입력이 모든 파라미터를 사용할 필요 없이 훨씬 많은 파라미터를 포함할 수 있지만, 전체 모델은 여전히 GPU 메모리에 저장되고 학습 중에 업데이트되어야 하며, 클러스터 전역에서 올바른 전문가에게 입력을 라우팅하는 데 자체적인 통신 및 조정 비용이 발생합니다. Ai2는 이러한 비용이 MoE가 커짐에 따라 계산상의 이점을 크게 감소시킬 수 있으며, Olmo-core 3가 그 격차를 메우도록 설계되었다고 밝혔습니다. 한 Ai2 벤치마크에서는 전문가 풀을 8명에서 128명으로 늘렸지만 토큰당 여전히 4명의 전문가를 선택했으며, 활성 파라미터는 약 32억 개 수준으로 유지하면서 전체 파라미터 용량은 46억 개에서 470억 개로 증가했으며, 학습 처리량은 5% 미만 감소했습니다.

FSDP에서 DDP 기반 트레이닝 스택으로

Ai2의 이전 MoE 구현인 Olmo-core는 완전 샤드 데이터 병렬(Fully Sharded Data Parallelism)을 사용했으며, 각 작은 학습 배치에 대해 모델 가중치를 수집하고 다시 샤드하도록 구성되었습니다. Olmo-core 3는 전문가를 GPU에 상주시키고 관련 데이터를 라우팅하는 분산 데이터 병렬(Distributed Data Parallelism) 기반 시스템으로 전환하여 반복적인 가중치 수집을 방지합니다. 8대의 NVIDIA B300 GPU에서 수행한 예비 테스트에서 Ai2는 470억 파라미터 MoE가 새로운 스택으로 GPU당 초당 52,000 토큰을 처리했으며, 이전 구현에서는 19,400 토큰으로, 이는 약 2.7배의 처리량에 해당한다고 보고했습니다.

Ai2의 희소 모델 작업은 64명의 라우팅된 전문가를 갖춘 MoE 아키텍처를 사용한 OlmoE까지 거슬러 올라갑니다. 반면 Olmo 3는 거의 모든 모델이 각 토큰마다 활성화되는 밀집 아키텍처를 사용했습니다. Olmo-core 3는 훨씬 큰 MoE 모델을 위한 트레이닝 시스템을 프레임워크에 추가합니다. Ai2는 NVIDIA의 Megatron-Core가 대규모 MoE 훈련을 위한 확립된 옵션이라고 언급했으며, Olmo-core 3가 Olmo 뒤의 프레임워크에 통합된 MoE 트레이닝 스택을 제공한다고 설명했습니다.

병렬성, 정밀도 및 메모리 기법

모델과 학습 상태를 하드웨어에 어떻게 분할할지 결정하는 세 가지 기법이 있습니다: expert parallelism은 전문가를 GPU에 분산시키고, pipeline parallelism은 모델의 레이어를 GPU 그룹에 나누며, distributed optimizer는 모든 GPU에 전체 복사본을 저장하는 대신 옵티마이저 상태를 GPU에 분산시킵니다. Olmo-core 3는 올바른 전문가에게 데이터를 라우팅하는 비용도 줄입니다: rowwise expert parallelism은 라우팅된 데이터를 직접 전문가 입력 버퍼에 넣고, GPU-resident routing은 라우팅 메타데이터를 GPU에 유지해 CPU가 데이터를 다시 복사할 때까지 기다리지 않고 작업을 대기시킬 수 있게 하며, grouped GEMM은 다수의 작은 전문가 연산을 결합해 GPU가 이를 보다 효율적으로 실행하도록 합니다. 기술 보고서는 NVSHMEM 기반의 rowwise expert parallelism 설계를 설명하며, 각 토큰을 해당 전문가의 버퍼에 직접 기록하고, 디바이스 스케줄링된 grouped matrix multiplication을 통해 expert parallelism을 완전한 동기화 없이 수행하도록 합니다.

Olmo-core 3는 낮은 정밀도 숫자 형식인 MXFP8을 지원합니다. 4대 NVIDIA B300 GPU에서 전문가 간에 작업을 균등하게 분배한 제어 벤치마크에서 Ai2는 BF16 기준 대비 학습 처리량이 약 21% 증가했으며, 피크 활성 메모리는 103 GiB에서 95 GiB로 감소했으며, 대부분의 향상은 피드포워드 연산 및 전문가 간 데이터 이동에서 비롯되었습니다. 보고서는 또한 토폴로지에 구애받지 않는 체크포인트가 병렬 레이아웃과 무관하게 전역 FP32 텐서를 기록하므로 다른 토폴로지에서 실행을 재개할 수 있다고 덧붙였으며, 제어된 비교에서 활성화 재계산이 활성화 메모리의 거의 3분의 2를 제거하고 피크 메모리를 약 1/4 감소시켰지만 처리량은 약 1/5 감소했다고 밝혔습니다.

트릴리언 파라미터 벤치마크 및 명시된 제한

Ai2는 NVIDIA B300 GPU에서 다양한 구성으로 Olmo-core 3를 벤치마크했으며, 그 중 512 GPU에 걸쳐 토큰당 583.6억 파라미터가 활성화되는 1.2 트릴리언 파라미터 모델을 포함했고, 최고 관측 처리량은 GPU당 858 TFLOP/s였습니다. Ai2는 이 테스트들이 학습된 모델의 품질이 아닌 시스템 성능을 측정하기 위해 무작위 라우팅을 사용했다고 밝혔습니다. 기술 보고서는 16 GPU에서 129억 파라미터 모델부터 512 GPU에서 1.2 트릴리언 파라미터 모델까지 측정된 운영 포인트를 나열하고, 헤드라인 속도는 무작위 라우팅 하에서 측정된 시스템 기준치이며, 제어된 스케일링 곡선이나 시간 대비 품질 주장과는 다르다고 명시합니다.

Ai2는 또한 DeepEP v2를 실험했으며, 이는 GPU 간 전문가들 간 통신을 위한 대체 백엔드로, 총 2.38조 파라미터의 구성을 달성했습니다. Ai2는 이 결과를 지속적인 학습 성능이라기보다 Olmo-core 3이 도달할 수 있는 규모를 보여주는 단기 용량 테스트라고 설명합니다. “Supercharging Olmo-core for Efficient and Scalable MoE Training”이라는 제목의 기술 보고서는 2026년 10월에 발행되었으며, 저자는 Allen Institute for AI와 University of Washington 소속이며, Tao Tianhua가 주요 저자이자 주요 개발자로 기재되어 있습니다.

문서화된 발견 및 차세대 Olmo 계획

보고서는 Ai2가 토큰 게리맨더링이라고 부르는 실패 패턴을 문서화했는데, 이는 균형 잡힌 라우팅을 장려하기 위한 점수가 실제 작업 부하가 덜 균형 잡히게 되면서도 향상될 수 있음을 의미합니다. 기타 문서화된 발견 사항으로는, 전문가들이 처리하는 토큰 수가 적어 학습률을 낮추는 것이 테스트된 모델군에서는 결과를 개선하지 못했으며, 동일한 행렬 차원에서도 처리되는 값이 변하면 GPU 계산 시간이 달라졌고, 별도의 GPU 스트림에서 통신과 계산을 겹치게 하는 것이 항상 학습을 빠르게 하지는 못했으며 일부 테스트에서는 전체 실행 시간을 오히려 늦추기도 했습니다. 보고서는 또한 채택되지 않은 접근 방식도 기술했는데, 여기에는 호스트 링크가 전달할 수 없었던 활성화의 CPU 오프로드와 전문가 연산과 GPU 자원을 경쟁하는 두 가지 겹침 방식이 포함됩니다.

Ai2는 차세대 Olmo가 MoE 아키텍처를 사용할 것이며, 가장 큰 데이터셋과 가장 긴 컨텍스트 윈도우로 학습된 가장 강력한 Olmo가 되기를 목표로 한다고 밝혔습니다. 이 조직은 Olmo-core 3이 완전히 오픈되어 연구자와 개발자가 자체 MoE를 학습하고, 다양한 하드웨어에 맞게 조정하며, 라우팅, 병렬성 및 시스템의 다른 부분을 실험할 수 있다고 말했습니다. Olmo-core GitHub 저장소는 이 프로젝트를 OLMo 생태계를 위한 PyTorch 빌딩 블록으로 설명하고 있으며, Apache-2.0 라이선스를 가지고 있고, 소스에서 직접 또는 PyPI에서 ai2-olmo-core로 설치할 수 있습니다.