Mixture of Expert (MoE) 모델에서 토큰을 전문가(expert)에게 이동시키는 비용이 전체 실행 시간을 지배할 수 있습니다.

Mixture of Expert (MoE) 모델에서는 토큰을 각 전문가(experts)에게 이동시키는 비용이 전체 실행 시간(total runtime)을 지배할 수 있습니다.

Zyphra research에서는 패턴을 활용하여 토큰이 전문가들에게 라우팅되는 방식을 개선함으로써, 모델 자체를 변경하지 않고도 @AMD MI300X GPU에서 통신 속도를 최대 2.63배까지 빠르게 만들었습니다.

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

저희가 테스트한 결과, 이 데이터 교환 과정은 단일 노드에서 스텝 시간(step time)의 13-24%를 차지했고, 네 개의 노드를 거치면서는 45-60%를 차지했습니다.

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

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

그래서 함께 선택되는 전문가들을 같은 GPU에 배치하고, 그 전문가들이 아무리 많은 GPU에 걸쳐 있더라도 각 토큰을 단 한 번의 GPU로 전송합니다. 저희는 이를 상관관계 있는 전문가 배치(correlated expert placement)라고 부릅니다. 8개의 GPU에서 이 방법은 최대 58%의 토큰 복사본 전송을 줄여줍니다.

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

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

단 4,096개의 토큰만 사용해도, 우리의 배치 및 라우팅 예측은 524,000개의 토큰을 사용한 결과와 약 1 퍼센트 포인트 이내의 오차를 보였습니다.

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

- 토큰 교환: 1.16배 ~ 2.63배 빠름
- 전체 훈련 스텝: 최대 1.41배 빠름

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

전문가와 토큰의 위치를 구성함으로써 우리는 동일한 작업을 수행하는 데 필요한 통신량을 줄입니다.

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

@ZyphraAI는 캘리포니아주 산호세에 기반을 둔 개방형 초지능 연구 및 제품 회사로, 개인과 조직이 잠재력을 최대한 발휘하도록 돕는 인간 정렬 AI를 구축하는 것을 목표로 합니다.

합류하려면 지원하세요!