앨런 인공지능 연구소(Ai2)는 대규모 혼합 전문가(MoE) 모델 학습 시스템의 통신 병목을 해결하고 조 단위 파라미터 범위까지 확장할 수 있도록 분산 인프라 구조를 재설계한 개방형 프레임워크 Olmo-core 3을 발표했습니다1. MoE 모델은 모든 입력에 대해 전체 파라미터를 동원하는 대신 일부 전문가(Expert) 네트워크만 선택적으로 활성화하여 연산 효율을 높입니다1. 그러나 수천 개의 GPU 클러스터로 확장될 때 입력 토큰을 적절한 전문가로 라우팅하고 데이터를 동기화하는 통신 오버헤드가 커지면서 연산상의 이점이 상쇄되는 한계가 있었습니다1.
Olmo-core 3은 이러한 통신 오버헤드와 메모리 병목을 극복하기 위해 학습 스택 구조를 바꾼 업그레이드 버전입니다1. 기존의 가중치 분할 방식인 FSDP 기반 구현에서 전문가 가중치를 GPU 메모리에 상주시키고 데이터만 전달하는 DDP 기반 시스템으로 전환하여 가중치 수집(Weight Gathering) 작업의 반복을 제거했습니다13. 이와 함께 저정밀 수치 포맷 MXFP8 지원과 메타데이터 기반 데이터 전처리 검증 기능을 도입하여 학습 처리량과 전처리 신뢰성을 한층 끌어올렸습니다13.
FSDP에서 DDP로의 재설계: 통신 병목 해결 원리
기존 Olmo-core의 MoE 구현체는 FSDP(Fully Sharded Data Parallel) 방식을 기반으로 구축되었습니다1. FSDP는 모델 가중치와 그래디언트를 여러 GPU에 분할 저장하여 단일 장치의 메모리 한계를 극복하는 데 유리합니다1. 하지만 토큰별로 활성화되는 전문가가 달라지는 MoE 아키텍처에서는 미니배치를 처리할 때마다 분산된 전문가 가중치를 매번 불러오는 가중치 수집 과정이 반복 발생했고, 이는 GPU 간 통신 지연을 유발하여 전체 학습 처리량을 저하시키는 원인이 되었습니다1.
Olmo-core 3은 이 문제를 해결하기 위해 데이터 병렬 처리를 담당하는 OLMoDDP 학습 스택을 신규 도입하며 DDP(Distributed Data Parallelism) 기반 시스템으로 전면 전환했습니다13. DDP 구조에서는 전문가 네트워크를 각 GPU 메모리에 고정 배치하고 토큰 데이터만 해당 전문가가 있는 GPU로 라우팅하여 반복적인 가중치 수집 작업을 원천 차단합니다1. 여기에 글로벌 로드 밸런싱을 수행하는 EMO(Expert-pool routing) 기법과 전문가 병렬화(Expert parallelism), 파이프라인 병렬화(Pipeline parallelism), 분산 옵티마이저 기술을 유기적으로 결합하여 메모리 부담과 라우팅 비용을 동시에 줄였습니다13.
벤치마크 조건과 실증적 성능 변화
구조적 개선의 효과는 통제된 실험 환경의 구체적인 수치로 증명됩니다. Ai2가 NVIDIA B300 GPU 8개 환경에서 470억(47B) 파라미터 규모의 MoE 모델을 테스트한 결과, Olmo-core 3의 OLMoDDP 스택은 GPU당 초당 52,000 토큰의 처리량을 기록했습니다1. 이는 기존 FSDP 기반 구현의 초당 19,400 토큰 대비 약 2.7배 향상된 결과입니다1. 이때 처리량(Throughput)은 첫 응답 시간을 나타내는 지연 시간(Latency)이 아니라 단위 시간당 소화 가능한 데이터 학습 연산 효율을 의미합니다.
전문가 풀(Expert Pool) 확장 테스트에서도 연산 효율성이 안정적으로 유지되었습니다. 전문가 수를 8개에서 128개로 늘리고 토큰당 4개의 전문가만 선택하도록 설정한 실험에서, 토큰당 활성 파라미터 수는 약 32억 개 수준으로 유지된 채 전체 파라미터 용량이 46억 개에서 470억 개로 약 10배 확대되었습니다1. 이 과정에서 발생한 학습 처리량 감소율은 5% 미만에 불과했습니다1. 또한 NVIDIA B300 GPU 4개를 사용한 저정밀 수치 포맷 MXFP8 제어 벤치마크에서는 연산 개선 효율이 가장 높은 구간에 MXFP8을 적용했을 때 BF16 기준선 대비 학습 처리량이 약 21% 증가했으며 피크 활성 메모리는 103 GiB에서 95 GiB로 감소했습니다1.
조 단위 파라미터 확장 검증과 통신 라이브러리 실험
Olmo-core 3은 수백 대 규모의 대형 클러스터 환경에서도 확장 성능을 입증했습니다. NVIDIA B300 GPU 512개를 투입한 벤치마크에서 토큰당 583억 6천만 개의 활성 파라미터를 갖춘 1조 2천억(1.2T) 파라미터 모델 구성을 평가한 결과, GPU당 최고 858 TFLOP/s의 높은 연산 효율을 보였습니다1. 다만 해당 측정 결과는 모델의 답변 품질이나 학습 수렴도를 나타낸 것이 아니며, 순수한 시스템 연산 성능을 측정하기 위해 랜덤 라우팅(Random routing) 조건을 적용한 테스트입니다1.
고성능 통신 라이브러리인 DeepEP v2를 적용한 연계 실험에서는 전체 파라미터 2조 3천억 개(2.38T)를 초과하는 구성까지 도달했습니다1. 다만 이 테스트는 전체 사전 학습 일정을 수행한 것이 아니라 시스템이 수용할 수 있는 하드웨어 용량의 한계를 점검한 단기 용량 테스트(Short-capacity test)라는 점을 명확히 구분해야 합니다1. 대규모 장기 학습 시의 실제 수렴 특성과 지속 처리량은 추가 검증이 필요합니다.
개발진은 기술 보고서를 통해 대규모 MoE 학습 시 유의할 실증적 관찰 결과도 공유했습니다. 라우팅 불균형을 막기 위한 정량적 점수가 개선되더라도 실제 토큰 분배는 오히려 불균형해지는 현상(Token gerrymandering)이 관찰되었으며, 처리 토큰 수가 적은 전문가의 학습률(Learning rate)을 낮추는 시도가 모델 성능 개선으로 이어지지 않았습니다1. 또한 GPU 연산 시간은 동일한 행렬 크기라도 입력 데이터 값에 따라 달라지므로 정확한 성능 비교를 위해 입력값 조건의 일치가 필수적이라는 점을 밝혔습니다1.
데이터 안정성, 기술적 제약 및 도입 판단 기준
학습 인프라 개편과 더불어 데이터 전처리 정확도와 최신 하드웨어 호환성을 강화하는 기능도 반영되었습니다. Olmo-core v3.0.0부터는 데이터 패킹 시 토큰 배열 내 EOS(End of Sequence) 스캔 방식 외에도 메타데이터 파일 기반의 문서 경계 설정 옵션(use_array_if_local=False)을 지원합니다3. 기존의 토큰 스캔 방식은 문장 잘림으로 EOS 토큰이 누락되었을 때 서로 다른 문서가 원치 않게 병합되는 문제가 있었습니다3. SFT 캐시 데이터 기준 측정 결과, 기존 방식에서는 97.98%의 토큰만 정제되어 도달했던 반면 메타데이터 방식을 활용하면 100.00%의 토큰이 손실 없이 학습 인스턴스에 전달됨을 확인했습니다3.
Olmo-core 3은 데이터 수집부터 사전 학습, 정렬 튜닝까지의 전체 흐름이 공개된 완전 개방형 LLM인 Olmo 3 모델 제품군(7B 및 32B 용량의 Base, Think, Instruct 버전)의 핵심 개발 기반입니다12. 이 프레임워크는 대규모 MoE 모델을 직접 구현하면서 분산 환경의 가중치 수집 병목을 해결하고자 하는 연구소 및 기업 엔지니어링 팀에 실질적인 가치를 제공합니다.
도입 검토 시 프레임워크의 환경 제약 사항을 사전에 확인해야 합니다. 최소 실행 요구 사양으로 PyTorch 버전이 2.10.0 이상으로 올랐으며, 실험적 기능인 KDA(Kernel-level Data Augmentation)를 활성화하려면 CUDA 13 CuTe 컴파일러가 요구됩니다3. 또한 Blackwell 아키텍처 환경에서는 flash-attn 4 버전 호환성 수정 내용이 포함되어 있으므로 기존 가변 길이 인자 전달 방식을 키워드 인자 기반으로 재설정했는지 확인 후 도입 환경을 구성하는 것이 권장됩니다3.
참고
1Introducing Olmo-core 3: Open, scalable training infrastructure for large MoEs
이 노트는 보도·비평·교육·연구 목적으로 작성되었습니다. 저작권 문의가 있으시면 에서 알려주세요.