OmniMoE: 확장 가능한 원자 전문가 조정을 통한 효율적인 Mixture-of-Experts (MoE) 모델

OmniMoE: An Efficient MoE by Orchestrating Atomic Experts at Scale

Mixture-of-Experts (MoE) 아키텍처는 파라미터 효율성을 향상시키기 위해 점점 더 세분화된 구조로 발전하고 있습니다. 그러나 기존의 MoE 설계는 전문가의 전문화 수준과 하드웨어 실행 효율성 간의 근본적인 트레이드오프를 가지고 있습니다. 본 논문에서는 전문가의 세분화 수준을 논리적인 극한까지 끌어올리는 시스템-알고리즘 공동 설계 프레임워크인 OmniMoE를 제안합니다. OmniMoE는 벡터 수준의 원자 전문가를 도입하여 단일 MoE 레이어 내에서 확장 가능한 라우팅 및 실행을 가능하게 하며, 동시에 범용 처리를 위한 공유된 Dense MLP 브랜치를 유지합니다. 이러한 원자적인 설계는 용량을 극대화하지만 라우팅 복잡성과 메모리 접근에 심각한 문제를 야기합니다. 이러한 문제를 해결하기 위해 OmniMoE는 시스템-알고리즘 공동 설계를 채택합니다. (i) 대규모 인덱스 공간을 분해하여 라우팅 복잡성을 O(N)에서 O(sqrt(N))으로 줄이는 카테시안 곱 라우터, 그리고 (ii) 전문가 중심 스케줄링을 통해 분산된 메모리 접근 방식을 효율적인 Dense 행렬 연산으로 변환합니다. 7개의 벤치마크에서 검증된 결과, 17억 개의 활성 파라미터를 가진 OmniMoE는 7개의 벤치마크에서 50.9%의 제로샷 정확도를 달성하여, 거친 입자(e.g., DeepSeekMoE) 및 미세 입자(e.g., PEER) 기반 모델을 능가했습니다. 특히, OmniMoE는 PEER에 비해 추론 지연 시간을 73ms에서 6.7ms로 줄여 10.9배의 속도 향상을 보여주며, 대규모 미세 입자 MoE 모델이 빠르고 정확할 수 있음을 입증했습니다. 본 논문의 코드는 https://github.com/flash-algo/omni-moe 에서 공개적으로 이용할 수 있습니다.

Mixture-of-Experts (MoE) architectures are evolving towards finer granularity to improve parameter efficiency. However, existing MoE designs face an inherent trade-off between the granularity of expert specialization and hardware execution efficiency. We propose OmniMoE, a system-algorithm co-designed framework that pushes expert granularity to its logical extreme. OmniMoE introduces vector-level Atomic Experts, enabling scalable routing and execution within a single MoE layer, while retaining a shared dense MLP branch for general-purpose processing. Although this atomic design maximizes capacity, it poses severe challenges for routing complexity and memory access. To address these, OmniMoE adopts a system-algorithm co-design: (i) a Cartesian Product Router that decomposes the massive index space to reduce routing complexity from O(N) to O(sqrt(N)); and (ii) Expert-Centric Scheduling that inverts the execution order to turn scattered, memory-bound lookups into efficient dense matrix operations. Validated on seven benchmarks, OmniMoE (with 1.7B active parameters) achieves 50.9% zero-shot accuracy across seven benchmarks, outperforming coarse-grained (e.g., DeepSeekMoE) and fine-grained (e.g., PEER) baselines. Crucially, OmniMoE reduces inference latency from 73ms to 6.7ms (a 10.9-fold speedup) compared to PEER, demonstrating that massive-scale fine-grained MoE can be fast and accurate. Our code is open-sourced at https://github.com/flash-algo/omni-moe.