[논문 리뷰] Tutel: Adaptive Mixture-of-Experts at Scale

Chang Ho Hwang, Wei Cui외 13명|arXiv (Cornell University)|2022. 06. 07.

Advanced Neural Network Applications인용 수 31

한 줄 요약

Tutel은 희소 MoE 모델용 무비용 전환 적응 병렬성 및 파이프라이닝 프레임워크를 도입하여 기존 MoE 구성과 알고리즘적 동등성을 유지한 채로 대규모에서 상당한 속도 향상을 달성합니다.

ABSTRACT

Sparsely-gated mixture-of-experts (MoE) has been widely adopted to scale deep learning models to trillion-plus parameters with fixed computational cost. The algorithmic performance of MoE relies on its token routing mechanism that forwards each input token to the right sub-models or experts. While token routing dynamically determines the amount of expert workload at runtime, existing systems suffer inefficient computation due to their static execution, namely static parallelism and pipelining, which does not adapt to the dynamic workload. We present Flex, a highly scalable stack design and implementation for MoE with dynamically adaptive parallelism and pipelining. Flex designs an identical layout for distributing MoE model parameters and input data, which can be leveraged by all possible parallelism or pipelining methods without any mathematical inequivalence or tensor migration overhead. This enables adaptive parallelism/pipelining optimization at zero cost during runtime. Based on this key design, Flex also implements various MoE acceleration techniques. Aggregating all techniques, Flex finally delivers huge speedup at any scale -- 4.96x and 5.75x speedup of a single MoE layer over 16 and 2,048 A100 GPUs, respectively, over the previous state-of-the-art. Our evaluation shows that Flex efficiently and effectively runs a real-world MoE-based model named SwinV2-MoE, built upon Swin Transformer V2, a state-of-the-art computer vision architecture. On efficiency, Flex accelerates SwinV2-MoE, achieving up to 1.55x and 2.11x speedup in training and inference over Fairseq, respectively. On effectiveness, the SwinV2-MoE model achieves superior accuracy in both pre-training and down-stream computer vision tasks such as COCO object detection than the counterpart dense model, indicating the readiness of Flex for end-to-end real-world model training and inference.

연구 동기 및 목표

• 토큰 라우팅과 고르게 분포되지 않는 전문가 부하로 인해 MoE 워크로드가 엑사스케일에서 왜 동적이고 도전적인가를 설명한다.
• DP, EP+DP+MP 및 관련 구성 간의 통합 및 전환을 위한 제로 오버헤드 적응 병렬성 전환 메커니즘을 제안한다.
• 계산과 통신을 공동으로 최적화하는 적응형 파이프라이닝과 All-to-All 전략(2DH 포함)을 개발한다.
• 동적 용량 계수와 동적 top-k 라우팅 지원을 갖춘 실용적이고 프레임워크 독립적인 MoE 시스템(Tutel)을 제공한다.
• 대형 GPU 클러스터에서 SwinV2-MoE를 사용하여 확장성 및 실제 효과를 입증한다.

제안 방법

• MoE 계층의 동적 워크로드 특성과 병렬성 선택에 미치는 영향을 분석한다.
• 구성 전반에서 최적의 성능을 달성하는 데 필요한 최소하지만 완전한 병렬성 전략 집합(DP 및 EP+DP+MP)을 도출한다.
• 전환 시 O(1) 오버헤드를 수반하도록 데이터 배치를 재구성하는 제로-코스트 전환 가능한 병렬성을 도입한다.
• 워크로드에 따라 파이프라이닝 차수와 All-to-All 알고리즘(Linear vs 2DH)을 공동으로 선택하는 적응형 파이프라이닝을 개발한다.
• 용량 범위에 걸친 최적의 병렬성/파이프라이닝 설정의 해시 맵(딕셔너리)을 구축하여 런타임 의사결정을 빠르게 할 수 있게 한다.
• 희소 계산과 Flexible All-to-All를 활용한 빠른 인코드/디코드를 구현하여 메모리 및 통신 오버헤드를 줄인다.

실험 결과

연구 질문

• RQ1 MoE 워크로드가 학습 동역학과 용량 요소에 따라 어떻게 달라지며, 정적 병렬성과 동적 병렬성에 대한 시사점은 무엇인가가?
• RQ2 동일한 데이터 레이아웃과 알고리즘적 정확성을 유지하면서 병렬성 전략 간 제로 비용 전환을 달성할 수 있는가?
• RQ3 동적 MoE 워크로드에서 최상의 처리량을 제공하는 데이터-병렬, 모델-병렬, 전문가-병렬 전략의 조합은 무엇이며 이를 런타임에 어떻게 적응시킬 수 있는가?
• RQ4 적응형 파이프라이닝과 All-to-All 알고리즘(Linear vs 2DH)을 함께 최적화하여 대규모 MoE 모델의 처리량을 극대화할 수 있는 방법은?
• RQ5 실제 MoE 모델(예: SwinV2-MoE)에서 Tutel의 훈련 및 추론 시 실용적인 성능 향상은 무엇인가?

주요 결과

• 단일 MoE 계층에서 16GPU 대비 최대 4.96배의 속도향상, 2,048GPU 대비 5.75배 향상.
• End-to-end SwinV2-MoE 훈련은 Fairseq 대비 훈련에서 최대 1.55배, 추론에서 2.11배의 속도향상을 보인다.
• 용량 계수 f에 따라 달라지는 적응형 병렬성 스위칭은 다양한 설정에서 정적 전략보다 우수한 처리량을 제공한다.
• 적응형 파이프라이닝은 243개 MoE 모델 설정에서 평균적으로 9%에서 100% 이상까지 개선을 제공하며, 최악의 경우에는 최대 599%의 향상을 달성한다.
• 유연한 All-to-All 및 빠른 인코드/디코드 기술은 메모리 사용을 줄이고 처리량을 향상시켜 대형 GPU 클러스터에서 확장 가능한 희소 MoE 구현을 가능하게 한다.
• Tutel과 함께한 SwinV2-MoE는 COCO와 같은 다운스트림 CV 작업에서 조밀 기반선보다 우수한 정확도를 달성하여 이 접근법의 실용적 효과를 시사한다.