Olmo-core 3 오픈 믹스처오브익스퍼트 학습 인프라스트럭처

TL;DR

Olmo-core 3은 새로운 오픈소스 학습 스택으로, 트리리언 파라미터 규모의 효율적인 믹스처오브익스퍼트(MoE) 모델을 가능하게 하며, 이전 구현 대비 최대 2.7배 높은 토큰 처리량을 제공하고 MXFP8 정밀도와 같은 고급 최적화를 지원합니다.

Olmo-core 3의 정의와 중요성

Olmo-core 3은 앨런 연구소의 MoE 학습 프레임워크를 재설계한 것으로, 전문가 풀이 수십에서 수백으로 증가할 때에도 토큰당 활성 파라미터 수는 대체로 일정하게 유지하면서 계산 효율성을 유지하도록 설계되었습니다. 전통적으로 MoE의 장점을 약화시키는 통신 및 라우팅 오버헤드를 줄임으로써, 트리리언 파라미터 규모의 희소 모델을 학계 연구소와 소규모 연구 팀이 경제적으로 활용할 수 있게 합니다.

핵심 아키텍처 변경사항

FSDP에서 DDP로 전환

  • 이전 Olmo-core는 각 미니배치마다 모델 가중치를 반복적으로 수집하고 재공유하는 완전 분산 데이터 병렬 처리(FSDP)를 사용했습니다.
  • Olmo-core 3은 분산 데이터 병렬 처리(DDP)를 채택하여 전문가 가중치를 GPU에 유지하고 필요한 토큰 데이터만 라우팅합니다. 이로써 비용이 큰 가중치 수집 단계를 제거합니다.

통합 MoE 스택 vs. Megatron-Core

  • NVIDIA의 Megatron-Core는 대규모 MoE의 기준 구현체로 남아 있습니다.
  • Olmo-core 3은 이전 FSDP 기반 버전보다 성능이 뛰어난 엔드투엔드로 밀접하게 통합된 MoE 스택을 제공합니다. 8개의 NVIDIA B300 GPU에서의 벤치마크에서, 47B 파라미터 MoE는 52,000 토큰/초/GPU를 달성했으며, 이전에는 19,400 토큰/초/GPU였습니다 – 2.7배의 처리량 증가입니다.

적용된 확장 기술

병렬화 전략

기술 목적
전문가 병렬화 전문가 풀을 GPU 간에 분산하여 각 GPU가 전문가의 일부만 저장합니다.
파이프라인 병렬화 모델 레이어를 GPU 그룹 간에 분할하여 GPU당 메모리 사용량을 줄입니다.
분산 최적화기 최적화기 상태를 GPU 간에 분산하여 모든 장치에 완전한 복사가 필요 없도록 합니다.

이 세 가지 메커니즘이 함께 작동하여, 모든 GPU가 전체 모델이나 최적화기 상태를 보유할 필요 없이 MoE 모델을 확장할 수 있게 합니다.

라우팅 및 계산 최적화

  • 행 단위 전문가 병렬화 – 라우팅된 토큰을 전문가 입력 버퍼에 직접 씁니다. 데이터 재구성 단계를 줄입니다.
  • GPU 내부 라우팅 – 라우팅 메타데이터를 GPU에 유지하여 CPU가 버퍼링 지점이 되지 않도록 합니다.
  • 그룹화된 GEMM – 여러 전문가에서 나온 많은 작은 행렬 곱셈을 하나의 큰 GEMM 호출로 묶어 GPU 활용도를 향상시킵니다.

MXFP8을 통한 정밀도 최적화

  • MXFP8은 계산과 GPU 간 통신량을 줄이는 저정밀도 수치 형식입니다.
  • 네 개의 NVIDIA B300 GPU에서 실시한 통제된 벤치마크에서 MXFP8을 활성화하면 BF16 기준 대비 ~21% 높은 엔드투엔드 처리량을 달성했으며, 피크 활성 메모리는 103 GiB에서 95 GiB로 감소했습니다.
  • 대부분의 성능 향상은 더 빠른 피드포워드 계산과 전문가 간 데이터 이동 감소에서 비롯되었습니다.

입증된 규모

구성 총 파라미터 수 토큰당 활성 파라미터 수 GPU 수 피크 처리량
1.2T 파라미터 MoE 1.2T 58.36B 512 858 TFLOP/s / GPU
2.38T 파라미터 (DeepEP v2) 2.38T — — —
  • 1.2T 파라미터 실행은 시스템 성능을 시험하기 위해 무작위 라우팅을 사용했으며, 훈련된 모델의 품질을 반영하지 않습니다.

순수 속도를 넘어서는 경험적 발견

  • 토큰 지방화 – 라우팅 균형 점수는 개선될 수 있지만 실제 작업 부하 균형은 악화될 수 있으며, 이는 지표 설계의 함정을 드러냅니다.
  • 전문가 학습률 스케일링 – 희소하게 사용되는 전문가의 학습률을 낮추는 것이 테스트된 모델 계열에서 더 나은 결과를 가져오지 않았습니다.
  • 입력값에 따른 시간 차이 – 입력 값이 다를 경우 동일한 행렬 구조에 대해 GPU 커널의 실행 시간이 변동하는 현상이 나타났으며, 이는 값이 일치하는 벤치마크의 필요성을 강조합니다.
  • 통신-계산 겹침 – 이 두 스트림을 겹치는 것이 종종 전체 실행 속도를 저하시켰으며, 더 많은 겹침이 항상 유리한 것은 아님을 보여줍니다.

이러한 관찰 결과는 함께 제공되는 기술 보고서에 자세히 설명되어 있으며, 향후 MoE 학습 연구를 안내합니다.

오픈소스 공개 및 커뮤니티 영향

앨런 연구소가 오픈 라이선스 하에 스택을 공개함으로써, 연구자들이 자체적으로 트리리언 파라미터 MoE를 훈련하고, 대체 하드웨어에 시스템을 적응시키며, 새로운 라우팅 또는 정밀도 기법을 실험할 수 있도록 합니다.

전망

Olmo-core 3은 다음 세대 Olmo 모델의 기반을 형성하며, 새로운 MoE 아키텍처를 더 큰 데이터셋과 더 긴 컨텍스트 창과 결합할 것입니다. 이 스택의 모듈식 설계는 하드웨어 발전과 함께 진화할 수 있어, 연구 커뮤니티가 확장 가능한 희소 모델 학습의 최전선에 머물 수 있도록 합니다.

주요 요약

  • Olmo-core 3은 토큰당 활성 파라미터 수를 일정하게 유지하면서 MoE 학습을 트리리언 파라미터 규모로 확장합니다.
  • DDP, 전문가/파이프라인 병렬화 및 라우팅 최적화를 통해 이전 FSDP 기반 스택 대비 최대 2.7배 높은 토큰 처리량을 달성합니다.
  • MXFP8 정밀도는 추가로 **~21%**의 속도 향상과 메모리 사용량 감소를 제공합니다.
  • 프레임워크는 완전히 오픈소스이며, 대규모 희소 모델 연구를 위한 프로덕션 수준 도구를 커뮤니티에 제공합니다.

Sources