Ai2, 조 단위 MoE 학습용 Olmo-core 3 공개

Ai2가 차세대 Olmo 모델의 기반이 될 오픈소스 학습 프레임워크 Olmo-core 3를 공개했습니다. 이 프레임워크는 Mixture-of-Experts(MoE) 모델을 조 단위 파라미터 규모로 확장하면서도 연산 효율성을 유지하도록 설계되었습니다.

주요 성능 지표

FSDP에서 DDP로 전환하고 전문가(expert) 가중치를 GPU에 상주시켜 반복적인 수집을 없앤 새로운 학습 스택을 도입했습니다. NVIDIA B300 GPU 8개 기준 예비 테스트 결과:

• 470억 파라미터 MoE 모델이 GPU당 초당 52,000개 토큰을 처리했으며, 이는 기존 구현의 19,400개 대비 2.7배 향상된 처리량입니다.
• 토큰당 4개 전문가만 선택하면서 전문가 풀을 8개에서 128개로 확장해도 활성 파라미터는 약 32억 개로 유지되고 총 용량은 470억 개로 증가했으며, 학습 처리량 감소는 5% 미만입니다.

기술 최적화

Olmo-core 3는 대규모 MoE를 GPU 클러스터에 분산하기 위해 여러 기술을 통합했습니다:

• 전문가 및 파이프라인 병렬화: 전문가와 모델 레이어를 GPU에 분산해 디바이스별 메모리 요구 사항을 줄입니다.
• 분산 옵티마이저: 옵티마이저 상태를 GPU 간에 분산하여 각 GPU에 전체 사본을 저장하지 않습니다.
• MXFP8 지원: NVIDIA B300 GPU 4개에서 이 저정밀도 형식을 사용하면 BF16 대비 학습 처리량이 약 21% 증가하고, 피크 활성 메모리는 103 GiB에서 95 GiB로 감소합니다.
• 라우팅 효율성: 행 단위 전문가 병렬화 및 GPU 상주 라우팅 기능으로 데이터 재배치와 CPU 대기 시간을 최소화합니다.

확장 한계 및 발견 사항

이 시스템은 1조 개 이상의 총 파라미터에서 벤치마킹되었으며, 512개 GPU를 사용한 특정 테스트에서는 1.2조 개 파라미터(토큰당 활성 583.6억 개)로 858 TFLOP/s/GPU를 달성했습니다. DeepEP v2를 사용한 짧은 용량 테스트에서는 2.38조 개 파라미터까지 확장성이 입증되었으나, 이는 전체 학습 실행은 아닙니다.

동반된 기술 보고서에서는 라우팅 균형 점수는 개선되지만 실제 워크로드 균형은 악화되는 '토큰 게리멘더링(token gerrymandering)' 현상과, 통신과 계산의 오버래핑이 항상 엔드투엔드 실행 속도를 높이는 것은 아니라는 관찰 결과 등 중요한 발견 사항도 강조했습니다.