[논문 리뷰] HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System

Xiaonan Nie, Pinxue Zhao외 3명|arXiv (Cornell University)|2022. 03. 28.

Domain Adaptation and Few-Shot Learning인용 수 11

한 줄 요약

HetuMoE는 고성능이며 분산형인 Mixture-of-Experts (MoE) 훈련 시스템으로, Hetu 기반으로 구축되어 여러 가지 게이팅 전략을 지원하며 일반 GPU 클러스터에서 훈련을 가속화한다. 계층적 AllToAll 통신과 최적화된 커널을 도입하여 배치 크기가 32일 때 스위치 게이팅 전략 하에서 DeepSpeed-MoE 대비 최대 8.1배의 속도 향상을 달성한다.

ABSTRACT

As giant dense models advance quality but require large amounts of GPU budgets for training, the sparsely gated Mixture-of-Experts (MoE), a kind of conditional computation architecture, is proposed to scale models while keeping their computation constant. Specifically, the input tokens are routed by the gate network and only activates part of the expert network. Existing MoE training systems only support part of mainstream MoE models (e.g. Top k) training under expensive high-bandwidth GPU clusters. In this paper, we present HetuMoE, a high-performance large-scale sparse MoE training system built on Hetu. HetuMoE provides multiple gating strategies and efficient GPU kernel implementations. To further improve the training efficiency on commodity GPU clusters (e.g, with only 1 NiC), we introduce the hierarchical AllToAll communication that combines hierarchical networks and aggregating messages. Compared with existing state-of-the-art MoE systems, HetuMoE obtains at least 15% speedup. Specifically, HetuMoE outperforms DeepSpeed-MoE up to 8.1x under the switch gate with a batch size of 32. Our code is available at: https://github.com/PKU-DAIR/Hetu.

연구 동기 및 목표

• 기존 MoE 훈련 시스템이 일반 GPU 클러스터에서 비효율적이고 확장성에 한계가 있는 문제를 해결하기 위해.
• Switch, GShard, Top-k 등 주요 게이팅 전략을 하나의 프레임워크에서 지원하기 위해.
• 저대역폭 네트워크에서 분산 MoE 훈련의 통신 병목 현상을 줄이기 위해 AllToAll 통신을 최적화하기 위해.
• NVLink나 InfiniBand 같은 고속 인터커넥트에 의존하지 않고도 최신 기술 수준의 훈련 성능을 달성하기 위해.
• 표준 GPU 클러스터에서 트리리언 파라미터 MoE 모델을 효율적이고 확장 가능하며 실용적으로 훈련할 수 있도록 하기 위해.

제안 방법

• 노드 내부 및 노드 간 통신을 조합한 계층적 AllToAll 통신 패턴을 도입하여 소규모 메시지를 집계하고 대역폭 활용도를 향상시킨다.
• Top-k 게이팅, 레이아웃 변환, AllToAll와 같은 핵심 MoE 연산을 위한 최적화된 GPU 커널을 사용하며, k=1 및 k=2에 대해 알고리즘적 개선을 적용한다.
• 이중 레이어 데이터 레이아웃 변환: 먼저 노드 내에서 토큰을 전문가별로 그룹화하고, 그 후 노드 간 AllToAll을 수행하여 통신 당 메시지 크기를 줄인다.
• Hetu의 분산 딥러닝 프레임워크를 활용하여 최소한의 코드 변경으로도 원활한 MoE 훈련을 지원하고 고성능을 달성한다.
• Switch, GShard, M6, BASE Layer, Hash Layer, SAM, Dense-to-Sparse 등 다양한 게이팅 전략을 지원하여 넓은 모델 호환성을 확보한다.
• 게이팅 계산, 레이아웃 변환, GPU 간 통신, 전문가 처리, 역방향 레이아웃 변환까지 MoE 훈련 파이프라인 전반을 최적화한다.

실험 결과

연구 질문

• RQ1 NVLink나 InfiniBand 같은 고대역폭 인터커넥트가 필요 없이도 일반 GPU 클러스터에서 MoE 훈련 시스템이 고성능을 달성할 수 있는가?
• RQ2 계층적 AllToAll 통신이 다중 노드 MoE 훈련에서 통신 오버헤드를 얼마나 효과적으로 줄이는가?
• RQ3 Top-k 게이팅 및 레이아웃 변환에 대한 커널 수준 최적화가 MoE 훈련 효율성에 얼마나 기여하는가?
• RQ4 다양한 게이팅 전략과 배치 크기에서 HetuMoE가 DeepSpeed-MoE, FastMoE, Tutel과 같은 최신 기술 수준의 시스템과 비교해 성능가 얼마나 뛰어나게 되는가?
• RQ5 통합된 시스템이 다양한 MoE 게이팅 전략을 효율적으로 지원하면서도 고성능 훈련 스루풋을 유지할 수 있는가?

주요 결과

• HetuMoE는 다양한 설정과 게이팅 전략에서 기존 MoE 시스템 대비 최소 15%의 속도 향상을 달성한다.
• 배치 크기가 32인 스위치 게이팅 전략 하에서 HetuMoE는 DeepSpeed-MoE 대비 최대 8.1배 빠른 훈련 속도를 기록한다.
• 계층적 AllToAll 최적화는 4×8 GPU 환경에서 기존 AllToAll 대비 1.66배, 8×8 GPU 환경에서 2배의 속도 향상을 달성한다.
• HetuMoE의 최적화된 Top-k 커널은 k=1 및 k=2에서 PyTorch 구현 대비 평균 25% 향상된 성능을 보인다.
• HetuMoE의 레이아웃 변환 커널은 최신 기술 수준의 구현 대비 26% 이상의 성능 향상을 기록한다.
• HetuMoE는 Switch, GShard, M6, BASE Layer, Hash Layer, SAM, Dense-to-Sparse 등 다양한 MoE 게이팅 전략을 성공적으로 지원하여 넓은 모델 호환성을 확보한다.