[리뷰] MS Tutel: A efficient mixture-of-experts implementation for large DNN model training
MS는 최근 ‘21/11/22에 대규모 DNN 모델 개발을 지원하는 고성능 MoE 라이브러리인 Tutel을 발표하였다. MoE는 계산비용이 parameters 개수에 sub-linear하여 딥러닝 모델을 Trillion 파라미터 이상으로 확장시킬 수 있는 유일한 접근 방식이다. Tutel은 MS의 Azure NDm A100 v4 시리즈에 대해 고도로 최적화되어 있다. 다음은 Azure NDm A100 v4의 사양은 다음과 같다.
Tutel MoE 최적화
Tutel은 Fairseq 및 FastMoE와 같은 다른 고수준의 MoE 솔루션을 보완하며 주로 MoE 특정 계산 및 all-to-all 통신의 최적화, 기타 다양하고 유연한 알고리즘 MoE 지원에 초점을 맞추고 있다. 또한 Tutel은 다른 MoE 솔루션에 통합하기 쉽도록 간결한 인터페이스를 가지고 있다.
연산을 위한 MoE-전용 최적화
Pytorch와 Tensorflow와 같은 딥러닝 프레임워크에 효율적인 MoE 구현이 없기 때문에 MoE기반 딥러닝 모델은 중복 계산으로 인한 성능 오버헤드가 발생한다. Tutel은 MoE-전용 연산을 위한 operators를 제공하기 위해 몇개의 고도로 최적화된 GPU 커널을 디자인하여 구현하였다.
예를들어 Tutel은 “게이팅 출력” 디스패치 시간 복잡도를 O(N³) → ON(N²)로 줄였다. 또한 fairseq 구현과 비교하여 속도가 24배 빠른 “ cumsum-minus-one operator”를 구현하였다. Tutel은 NVRTC(CUDA C++용 런타임 컴파일 라이브러리)를 활용하여 맞춤형 MoE 커널 Just-In Time을 좀더 최적화합니다.
이와 같은 최적화로 인해 Azure NDm A100 v4 플랫폼에서 fairseq와 Tutel의 결과를 비교하였을 때, single MoE layer를 가진 Tutel이 fairseq을 뛰어넘는 성능을 보인다.
- 8x A100 GPUs : 8.49x speedup
- 512x A100 GPUs: 2.75x speedup
Azure NDm A100 v4 clusters 환경에서 all-to-all comm. 최적화
Tutel은 CPU-GPU 바인딩과 적응형 라우팅(AR. Adaptive Routing) 튜닝을 포함하여, Azure NDm A100 v4 cluster 상에서 large-scale MoE 학습을 위한 all-to-all collective comm.을 최적화한다.
MoE에서 all-to-all 통신의 경우, 각 GPU별 통신의 총 데이터 트래픽 사이즈는 변하지 않지만 각 GPU 쌍간 데이터 사이즈는 GPU의 개수 증가와 함께 더 작아진다. 이는 MoE 학습 성능을 더 저화시킨다. Tutel은 AR 기술을 활용하여 작은 메시지 그룹에 대한 통신 효율성을 개선하고 NDm A100 v4 시스템에서 고성능 all-to-all 통신을 제공한다. CPU-GPU 바인딩과 AR 튜닝을 통해 수백 MiB 크기의 메시지 크기에 대해 512 A100 GPU로 2.56x ~ 5.93x의 전체 속도 향상을 달성하였다.
다양하고 유연한 MoE 알고리즘 지원
- Top-K 게이팅 알고리즘를 위한 임의의 K 세팅 (대부분의 구현은 Top-1 및 Top-2만 지원)
- batch-prioritized routing, input dropout, input jitter
- FP16, FP32, Mixed precision (다음 릴리즈에 BF16 지원)
- NVIDIA CUDA 와 AMD ROCm 디바이스 지원
Meta(구 facebook) 언어모델를 Tutel로 통합
MS는 Meta와 함께 Tutel을 fairseq 툴킷에 통합하였다. Meta는 Tutel을 사용하여 Azure NDm A100 v4에서 GPT-3와 유사한 어텐션 기반 뉴럴 아키텍처를 가진 large-scale 언어모델을 학습하고 있다.
MS는 Tutel의 end-to-end 성능을 평가하기 위해 Meta의 1.1 Trillion-parameters MoE 언어모델을 사용하였다. Meta의 MoE 언어 모델의 구성은 다음과 같다. 두개의 레이어마다 1개의 MoE 레이어를 포함하고 각 GPU는 1개의 expert를 가지고 있다.
Tutel의 성능은 다음과 같다. 8 A100 GPUs으로 131 %의 speedup, 512 A100 GPUs에서40%의 speedup 향상을 얻었다. 512 A100 GPUs에서 speedup 향상이 낮은 이유는 all-to-all 통신이 병목이기 때문이며 다음 버전에서 좀더 최적화될 것이다.
Tutel은 DeepSpeed 프레임워크에도 통합되었으며, Tutel 및 관련 통합은 특히 대규모 모델을 효율적으로 확장하려는 사람들과 Azure 서비스(Tutel 이 Azure NDm A100 v4에 최적화되어 있으므로)에 도움이 될 것이다.
레퍼런스
[1] Tutel: An efficient mixture-of-experts implementation for large DNN model training
[2] Microsoft expands its AI-supercomputer lineup with general availability of the latest 80GB NVIDIA A100 GPUs in Azure, claims 4 spots on TOP500 supercomputers list
