MoE 모델의 통신 병목 현상 해결을 위한 상관관계 기반 전문가 배치 및 라우팅 최적화

Mixture of Expert (MoE) 모델에서 토큰이 각 전문가(experts)에게 이동하는 비용이 전체 실행 시간의 지배적인 요인이 될 수 있습니다.

Zyphra 연구팀은 토큰이 전문가로 라우팅되는 패턴을 활용하여, 모델 자체를 변경하지 않으면서 @AMD MI300X GPU에서 해당 통신 속도를 최대 2.63배까지 빠르게 만드는 방법을 사용했습니다.

MoE 모델은 각 토큰을 소수의 전문가에게 희소하게(sparsely) 라우팅합니다. 이러한 모델이 성장함에 따라, 전문가들은 여러 GPU와 노드에 분산되며, 토큰들은 자신의 전문가들에게 전송되고 다시 돌아와야 합니다.

저희가 측정한 결과, 데이터 교환 과정만으로도 단일 노드의 스텝 시간의 1324%를 차지했으며, 네 개의 노드를 거치면서는 4560%에 달했습니다.

이러한 라우팅 결정에는 패턴이 존재합니다. 예를 들어, 한 레이어에서 가능한 8,128개의 전문가 쌍 중 단지 64쌍만이 토큰의 42%를 차지하는 경우도 있었습니다. 만약 전문가가 독립적으로 선택되었다면, 이 쌍들은 1.6%만을 담당했을 것입니다.

저희는 이러한 구조를 활용하여 전문가가 배치되는 위치를 개선하는 방법을 보여줍니다.

따라서 함께 사용되는 전문가들을 같은 GPU에 배치하고, 토큰이 몇 개의 전문가를 가지고 있든 상관없이 각 토큰을 단 한 번의 GPU로 전송합니다. 저희는 이를 '상관관계 기반 전문가 배치(correlated expert placement)'라고 부릅니다. 이 방법을 통해 8개의 GPU에서 최대 58%의 토큰 복사본 전송량을 줄일 수 있었습니다.

라우팅은 레이어 간에도 예측 가능합니다. 한 레이어에서 사용된 전문가가 다음으로 어떤 전문가를 필요로 할지 알려줍니다. '토큰 셔플링(Token shuffling)'은 각 토큰을 해당 전문가들을 보유한 GPU로 이동시키며, 이 과정은 이미 어텐션(attention) 이후에 실행되는 전송 과정 내부에 포함되므로 추가적인 네트워크 트래픽이 발생하지 않습니다.

결정적으로, 저희는 이러한 패턴들이 훈련 중에 유용하게 활용될 만큼 충분히 일찍 나타나며 측정 비용이 저렴하다는 것을 발견했습니다.

단지 4,096개의 토큰만을 사용했을 때도, 우리의 배치 및 라우팅 예측 결과는 524,000개의 토큰을 사용한 결과를 약 1% 포인트 이내로 구현할 수 있었습니다.

이러한 성능 향상은 훈련 구성에 따라 달라지며, 각 토큰이 더 많은 전문가를 사용하고 전문가들이 여러 노드에 걸쳐 분산될 때 가장 크게 나타납니다. Megatron-LM을 8개에서 64개의 GPU에서 실행했을 때의 결과는 다음과 같습니다:

• 토큰 교환(Token exchange): 1.16배 ~ 2.63배 빠름
• 전체 훈련 스텝(Full training step): 최대 1.41배 빠름

저희의 접근 방식은 손실이 없습니다. 동일한 토큰들이 동일한 전문가를 방문합니다. 아키텍처(architecture), 라우팅 결정 및 훈련 목표는 변경되지 않습니다. 모델은 동일한 함수를 계산합니다.

전문가와 토큰이 존재하는 위치를 체계화함으로써, 동일한 작업을 수행하는 데 필요한 통신량을 줄일 수 있습니다.

전체 블로그:
https://zyphra.com/our-work/expert-coupling-in-moe-pretraining
…
전체 논문:
https://arxiv.org/abs/2610.09372

@ZyphraAI는 캘리포니아주 샌프란시스코에 기반을 둔 개방형 초지능(superintelligence) 연구 및 제품 회사로, 개인과 조직이 잠재력을 최대한 발휘하도록 돕는 인간 정렬 AI(human-aligned AI)를 구축하는 것을 목표로 합니다.

저희와 함께할 기회에 지원하세요!