멀티헤드 라티엔트MoE와 헤드 패럴렐: 통신 효율과 결정성 확보를 위한 새로운 MoE 병렬화

초록

본 논문은 기존 전문가 병렬(Expert Parallel) 방식의 통신 비용·로드 불균형·비결정성 문제를 해결하기 위해, 토큰을 여러 서브‑토큰으로 분할하고 각각 독립적인 MoE 모듈에 할당하는 멀티헤드 라티엔트MoE와 이를 위한 헤드 패럴렐(HP) 병렬 전략을 제안한다. IO‑aware 라우팅과 블록‑스파스 어텐션 기반 전문가 연산을 도입해 HBM 접근을 최소화하고, 통신량을 O(1)로 고정, 트래픽을 완전 균형화하며, 메타데이터 교환 없이 결정적인 통신을 구현한다. 실험 결과, 동일 성능을 유지하면서 EP 대비 1.61배 빠른 학습 속도를 달성하고, granularity를 두 배로 늘려도 1.11배 가속을 기록한다.

상세 분석

본 연구는 대규모 언어 모델 학습에서 조건부 계산을 이용한 Sparse Mixture of Experts(MoE)의 병렬화 한계를 체계적으로 분석한다. 기존 Expert Parallel(EP) 방식은 토큰당 활성화된 전문가 수 k에 비례해 all‑to‑all 통신량이 증가하고, 전문가 큐 길이의 불균형으로 인해 라티시가 증가하며, 라우팅 결과에 따라 통신 패턴이 달라지기 때문에 메타데이터 교환이 필요해 비결정성을 초래한다. 이러한 문제를 근본적으로 해결하기 위해 저자는 두 가지 핵심 아이디어를 제시한다. 첫째, 토큰을 N_h개의 서브‑토큰으로 분할하고, 각 서브‑토큰을 독립적인 MoE 헤드에 할당하는 멀티헤드 라티엔트MoE 구조를 설계한다. 이때 각 헤드는 자체 라우터와 전문가 집합을 갖으며, 파라미터는 전혀 공유되지 않는다. 둘째, 헤드 패럴렐(HP)이라는 새로운 분산 전략을 도입해, 라우팅 이전에 서브‑토큰을 GPU 간에 한 번만 전송한다. 따라서 통신량은 k와 무관하게 O(1)로 고정되고, 모든 GPU가 동일한 양의 데이터를 교환하므로 트래픽이 완전 균형화된다. 또한 라우팅 결과가 통신에 영향을 주지 않으므로 메타데이터 교환이 필요 없으며, 통신이 완전히 결정적이다.

효율성을 확보하기 위해 두 가지 IO‑aware 연산을 구현한다. 첫 번째는 라우팅 단계에서 전체 라우팅 스코어와 활성화 정보를 HBM에 저장하지 않고, 블록 단위로 SRAM에 로드하여 온‑칩에서 top‑k를 실시간으로 계산한다. 이를 위해 스코어와 인덱스를 64비트 정수로 패킹하고, 블록별 로컬 top‑k를 누적하여 최종 top‑k를 얻는 방식으로 HBM 접근을 O(N_e)→O(k)로 감소시킨다. 두 번째는 전문가 연산을 블록‑스파스 어텐션 형태로 재구성하고, FlexAttention 커널을 재활용해 전문가 가중치와 입력을 Q·K^T·V 형태로 처리한다. 이 과정에서도 HBM 접근을 토큰·전문가 차원에서 O(T·d_e)→O(T+d_e)로 축소한다. 두 연산 모두 정확성을 유지하면서 메모리 대역폭을 크게 절감한다.

실험에서는 10B 토큰 규모의 FineWebEdu 데이터셋을 사용해 기존 MoE+EP와 비교하였다. 동일한 활성 파라미터 수와 모델 정확도를 유지하면서, 멀티헤드 라티엔트MoE+HP는 학습 속도가 최대 1.61배 향상되었으며, k=4일 때 통신량이 25% 수준으로 감소했다. granularity를 두 배로 늘린 경우에도 정확도는 6.9%p 상승하면서 1.11배 가속을 달성했다. 이러한 결과는 통신 효율성, 로드 밸런싱, 결정성 확보가 실제 학습 효율에 크게 기여함을 입증한다.

전반적으로 본 논문은 MoE 모델의 병렬화 설계에서 “라우팅 전 통신”이라는 새로운 패러다임을 제시하고, 하드웨어‑친화적인 알고리즘 설계와 기존 GPU 메모리 계층 구조를 활용한 최적화를 통해 대규모 모델 학습의 비용 장벽을 낮추는 실질적인 방안을 제공한다.