MoonshotAI/MoonEP

해결하는 문제

MoonEP는 Mixture-of-Experts (MoE) 모델에서 Expert Parallelism (EP)의 효율성 저하 문제, 특히 부하 불균형 문제를 해결합니다. 일반적인 EP에서는 일부 GPU 랭크가 다른 랭크보다 훨씬 많은 토큰을 수신하게 되는("가장 뜨거운 랭크" 문제) 상황이 발생하여, 지연 시간 증가, 메모리 분할, 트레이닝 중 Out-of-Memory (OOM) 오류가 발생할 수 있습니다.

작동 방식

MoonEP는 라우터의 분배 방식에 관계없이 모든 GPU 랭크가 동일한 수의 토큰을 처리하도록 보장합니다. 이를 위해 다음과 같은 기술을 사용합니다:

• 동적 중복 전문가: 라이브러리는 온라인으로 소수의 중복 전문가를 계획하고 미리 로드하여 랭크 간의 부하를 완벽하게 균형화합니다.
• 제로 코피 통신: 토큰을 원격 랭크의 최종 전문가 그룹 위치로 직접 전송하기 위해 융합된 permute/unpermute 연산을 사용하며, 계산에 버퍼 뷰를 반환하여 불필요한 데이터 복사를 제거합니다.
• 온라인 계획: 고성능 GPU 커널이 거의 무시할 수 있는 오버헤드로 전문가와 토큰의 최적 분배를 결정합니다.
• 정적 메모리 형상: 토큰용 버퍼 크기를 고정함으로써 레이어별 호스트 동기화의 필요성을 제거하고 메모리 분할을 방지합니다.

대상 사용자

이 라이브러리는 대규모 MoE 모델을 구축하고자 하는 개발자 및 연구자들을 위한 것으로, 여러 NVIDIA GPU(향후 Zhenwu PPUs)에서 통신 오버헤드를 최적화하고 안정적이고 균형 잡힌 트레이닝 및 추론을 보장해야 하는 경우에 적합합니다.

주요 특징

• 완벽한 부하 균형: 모든 랭크가 정확히 동일한 수의 토큰을 수신하여 반복 시간이 라우팅 불균형에 영향을 받지 않습니다.
• 제로 코피 디스패치: 통신 버퍼에서 사용자 버퍼로의 복사를 제거하여 원시 통신 속도를 향상시킵니다.
• OOM 방지: 정적 메모리 형상으로 인해 불균형한 MoE 트레이닝에서 흔히 발생하는 활성화 형상 변화로 인한 메모리 분할을 방지합니다.
• 통합된 가중치 프리페치: 중복 전문가 가중치를 로컬 슬롯에 효율적으로 프리페치하는 전용 메커니즘을 포함합니다.