MoonshotAI/MoonEP

해결하는 문제

MoonEP는 전문가 병렬화(EP) 중 혼합 전문가(MoE) 모델의 부하 불균형 문제를 해결합니다. 기존 MoE에서는 일부 전문가가 다른 전문가보다 훨씬 많은 토큰을 받게 되어 "핫 랭크"가 발생하며, 이는 전체 학습 또는 추론 프로세스를 지연시키고 GPU 메모리 분할 또는 메모리 초과(OOM) 오류를 유발할 수 있습니다.

작동 방식

MoonEP는 루팅 분포에 관계없이 모든 GPU 랭크가 동일한 수의 토큰을 받도록 보장합니다. 이를 위해 다음과 같은 기술을 사용합니다:

• 동적 중복 전문가: 온라인에서 소수의 중복 전문가를 계획하고, 필요로 하는 랭크에 미리 프리페치하여 모든 랭크 간에 작업 부하가 완벽하게 균형되도록 합니다.
• 온라인 계획: 최소한의 오버헤드로 이러한 중복 전문가의 계획을 고성능 GPU 커널이 처리합니다.
• 제로 코피 통신: 토큰을 원격 랭크의 최종 위치로 직접 전송하기 위해 융합된 permute/unpermute 연산을 사용하며, 통신 버퍼의 뷰를 계산에 직접 제공함으로써 비용이 큰 데이터 복사를 제거합니다.
• 정적 메모리 형상: 토큰용 버퍼 크기를 고정함으로써 메모리 분할을 방지하고 예측 가능한 메모리 사용을 보장합니다.

대상 사용자

NVLink를 사용하여 다수의 GPU에서 안정적인 학습 성능을 유지하면서 대규모 MoE 모델을 구축하는 개발자 및 연구자에게 적합합니다.

주요 특징

• 불균형에 강함: 루팅 불균형이 증가해도 반복 시간이 일정하게 유지됩니다.
• 제로 코피 효율성: 통신 버퍼와 사용자 버퍼 간의 복사를 제거하여 더 빠른 원시 통신이 가능합니다.
• 정적 메모리: 고도의 불균형 상황에서도 GPU 메모리 분할 및 OOM을 방지합니다.
• 통합 가중치 프리페치: 중복 전문가에 대한 프리페치 및 기울기 감소를 위한 특화된 커널을 포함합니다.