MoE 모델은 각 토큰을 소수의 전문가(expert)에게 희소하게 라우팅합니다. 이러한 모델이 커지면서...

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

Zyphra 연구에서는 토큰이 전문가에게 라우팅되는 패턴을 활용하여 모델 변경 없이도 @AMD MI300X GPU에서 해당 통신 속도를 최대 2.63배까지 빠르게 만들었습니다.

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

저희의 테스트에서, 데이터 교환은 단일 노드에서 스텝 시간의 1324%를 차지했고, 네 개의 노드를 거치면서는 4560%를 차지했습니다.

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

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

따라서 저희는 함께 선택되는 전문가들을 동일한 GPU에 배치하고, 해당 토큰의 전문가가 몇 개가 있든 관계없이 각 토큰을 단 한 번만 GPU로 전송합니다. 이를 상관관계 전문가 배치(correlated expert placement)라고 부릅니다. 8개의 GPU에서 이는 최대 58%의 토큰 복사본 전송량을 줄입니다.

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

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

단지 4,096개의 토큰만을 사용했음에도 불구하고, 저희의 배치 및 라우팅 예측은 524,000개의 토큰을 사용한 결과와 약 1% 포인트 이내로 근접합니다.