Fast MoE Inference via Predictive Prefetching and Expert Replication

기술 요약: 예측적 프리페칭 및 전문가 복제를 통한 고속 MoE 추론

문제 제기
혼합 전문가 (MoE) 아키텍처는 계산 오버헤드의 비례적 증가 없이 방대한 파라미터 수를 가능하게 함으로써 대규모 언어 모델 (LLM) 을 확장하는 데 핵심적인 기반이 되었습니다. 그러나 MoE 추론은 실제 배포를 방해하는 치명적인 비효율성에 직면해 있습니다:

1. 비최적 GPU 활용: 전문가 활성화의 고유한 희소성으로 인해, 특정 시점에 활성화되는 전문가의 소수 집합만 가동되며 GPU 자원의 대다수가 유휴 상태로 방치됩니다.
2. 부하 불균형 및 지연: 동적 라우팅은 종종 불균등한 부하 분배를 초래하여, 여러 토큰이 동일한 물리적 전문가를 위해 대기열에 서게 되고, 이로 인해 토큰 정지 (stalling) 와 지연 시간 증가가 발생합니다.
3. 라우팅 오버헤드: 전통적인 소프트맥스 기반 라우팅 메커니즘은 계산 복잡성과 통신 오버헤드 (all-to-all 라우팅) 를 도입하며, 이는 전문가 수가 증가함에 따라 확장성이 좋지 않습니다.

증류 (distillation) 나 커널 수준의 개선과 같은 기존 최적화 기법들은 종종 제한적인 개선만 제공하거나 모델 정확도를 희생합니다.

방법론: MoE-MPMC
저자들은 예측적 프리페칭과 동적 전문가 복제를 통해 이러한 비효율성을 해결하는 MoE-MPMC(Multiple Producers Multiple Consumers 를 갖춘 Mixture of Experts) 프레임워크를 제안합니다. 이 시스템은 두 개의 동시 스레드를 사용하여 작동합니다:

1. 추론 스레드:

   • 현재 배치에 대한 전문가 할당 및 복제 수를 포함하는 사전 계산된 해시 테이블을 검색합니다.
   • 예측된 활성화된 전문가만 GPU 메모리에 동적으로 로드하고, 사용하지 않는 전문가들은 CPU 메모리로 오프로딩합니다.
   • 활성화된 전문가들이 GPU 상에서 물리적으로 복제되는 수정된 MoE 레이어를 사용하여 추론을 실행합니다. 이를 통해 동일한 전문가가 필요한 여러 토큰을 별도의 물리적 복제본이 동시에 처리하여 대기열 지연을 제거합니다.

2. 해시 빌딩 스레드:

   • 추론 스레드와 병렬로 실행되어 다음 배치의 토큰을 처리합니다.
   • 희소 어텐션을 갖춘 **단순 순환 유닛 (SRU)**을 사용하여 다음 배치에서 활성화될 전문가를 예측합니다.
   • 이전 작업 (SiDA-MoE) 에서 사용된 LSTM 기반 예측기를 SRU 로 대체하여, 상호 의존성 감소와 더 간단한 게이트 메커니즘으로 인해 더 빠른 훈련 및 추론 속도를 달성합니다.
   • 토큰 수요에 기반하여 각 예측된 전문가의 복제 수를 계산하고 해시 테이블을 업데이트한 후, 이를 추론 스레드를 위해 대기열에 넣습니다.

주요 기여
본 논문은 세 가지 주요 기여를 제시합니다:

1. SRU 기반 예측: 각 토큰에 대한 전문가 선택을 신속하게 예측하기 위해 단순 순환 유닛 (SRU) 을 도입했습니다. 이는 더 느린 LSTM 예측기를 대체하여 예측 단계를 가속화하면서도 정확도를 유지합니다.
2. 심층 물리적 복제: 토큰 수요에 기반하여 GPU 상에서 예측된 전문가를 물리적으로 복제하는 전략입니다. 이를 통해 배치 처리 중 거의 모든 토큰이 대기 지연을 경험하지 않도록 보장하여 병렬성과 GPU 활용을 극대화합니다.
3. 용량 캡핑: 사전 정의된 GPU 메모리 제약에 따라 전문가 복제 수를 제한하는 실용적인 메커니즘입니다. 이는 오버플로우를 유발하지 않으면서 하드웨어 제한 내에서 복제를 최적화하여 거의 최적의 자원 활용을 달성합니다.

실험 결과
저자들은 SST2 및 MRPC 와 같은 데이터셋을 사용하여 대규모 SwitchTransformers 모델 (Switch-base-8, 64, 128, 256) 에서 MoE-MPMC 를 평가했습니다. 주요 발견 사항은 다음과 같습니다:

• GPU 활용: 이 방법은 거의 완전한 GPU 활용 (일부 구성에서 약 100% 에 달함) 을 달성하여, 베이스라인 SwitchTransformers(128 개 전문가의 경우 활용도가 1.3% 로 낮게 나타남) 및 SiDA-MoE 접근법보다 상당한 개선을 보였습니다.
• 추론 속도: 이 접근법은 베이스라인 아키텍처 대비 최대 3 배의 추론 속도 향상을 제공합니다.
• 정확도: 이 방법은 베이스라인 아키텍처 성능의 약 90-95% 를 유지하며, LSTM 대비 우수한 예측 능력을 가진 SRU 덕분에 SiDA-MoE 보다 정확도가 더 잘 유지됩니다.
• 처리량: 전문가 복제를 통한 병렬 토큰 처리 능력으로 인해 전문가 수가 증가함에 따라 처리량이 크게 증가합니다.

의의 및 주장
본 논문은 MoE-MPMC 가 MoE 기반 대규모 언어 모델의 실용적이고 경제적으로 타당한 배포를 향한 중요한 진전이라고 주장합니다. 반응형 부하 균형에서 예측적 프리페칭 및 복제의 사전 대응 전략으로 전환함으로써, 이 방법은 다음과 같은 효과를 거둡니다:

• 메모리 접근 병목 현상 및 토큰 정지를 완화합니다.
• 근본적인 MoE 계산 그래프를 변경하지 않고 GPU 자원 활용을 최적화합니다.
• 토큰 예측 작업과 동시에 미세 조정 및 사전 훈련이 발생할 때 전문가 선택 예측이 실현 가능하고 효과적임을 입증합니다.

저자들은 이 작업을 미래 수요를 예측하고 자원 할당을 동적으로 조정하여 희소 아키텍처에서 흔히 발생하는 유휴 GPU 메모리 및 부하 불균형의 함정을 해결함으로써, SiDA-MoE 와 같은 반응형 방법보다 진보된 것으로 위치시킵니다.