[논문 리뷰] M6-T: Exploring Sparse Expert Models and Beyond

Yang An, Junyang Lin외 13명|arXiv (Cornell University)|2021. 05. 31.

Recommender Systems and Techniques참고 문헌 34인용 수 14

한 줄 요약

이 논문은 전문화된 프로토타입 기법을 제안하며, 이는 전문가를 k개의 프로토타입으로 분할하고 k개의 top-1 라우팅을 사용하여 계산 비용을 유지하면서도 희소 전문가 모델의 성능을 향상시키는 방법이다. 이 방법을 통해 1.1조 파라미터 모델이 단지 480개의 V100 GPU에서 기존의 동일 규모 기준 모델 대비 5배 빠르게 학습되며, 2048개의 TPU 코어로 학습된 이전 최고 성능(SOTA) 모델들을 능가한다.

ABSTRACT

Mixture-of-Experts (MoE) models can achieve promising results with outrageous large amount of parameters but constant computation cost, and thus it has become a trend in model scaling. Still it is a mystery how MoE layers bring quality gains by leveraging the parameters with sparse activation. In this work, we investigate several key factors in sparse expert models. We observe that load imbalance may not be a significant problem affecting model quality, contrary to the perspectives of recent studies, while the number of sparsely activated experts $k$ and expert capacity $C$ in top-$k$ routing can significantly make a difference in this context. Furthermore, we take a step forward to propose a simple method called expert prototyping that splits experts into different prototypes and applies $k$ top-$1$ routing. This strategy improves the model quality but maintains constant computational costs, and our further exploration on extremely large-scale models reflects that it is more effective in training larger models. We push the model scale to over $1$ trillion parameters and implement it on solely $480$ NVIDIA V100-32GB GPUs, in comparison with the recent SOTAs on $2048$ TPU cores. The proposed giant model achieves substantial speedup in convergence over the same-size baseline.

연구 동기 및 목표

• 혼합 전문가(MoE) 모델에서 모델 품질과 학습 효율성에 영향을 주는 핵심 요인을 조사하기 위해.
• 1000억 파라미터를 초과하는 희소 전문가 모델을 효율적으로 확장하는 데 도전하기 위해.
• 초거대 규모의 MoE 아키텍처에서 계산 비용을 줄이면서도 모델 성능을 향상시키기 위해.
• 소비자용 GPU 클러스터를 사용하여 트리리언 파라미터 MoE 모델을 TPU 기반 시스템 대신 효율적으로 학습할 수 있는 가능성을 입증하기 위해.

제안 방법

• 전문가 프로토타입 기법은 각 전문가를 k개의 서로 다른 프로토타입으로 분할하여 표준 top-k 라우팅 대신 k개의 top-1 라우팅을 가능하게 한다.
• 각 입력 토큰은 게이팅 함수를 사용하여 k개의 별도 프로토타입으로 라우팅되며, 각 프로토타입 별로 독립적으로 라우팅 결정이 내려진다.
• 최종 출력은 softmax-게이팅된 어텐션을 사용해 k개의 프로토타입 출력을 가중합으로 계산한다.
• 전체 top-k 계산을 피하고 k개의 독립된 top-1 라우팅 연산을 사용함으로써 FLOPs를 일정하게 유지한다.
• 안정성을 확보하기 위해 학습률을 감소시킨(0.005) Adafactor 최적화 기법과 비례 확대된 가중치 초기화(σ=0.002)를 사용하여 모델을 학습한다.
• RDMA 인터커넥트를 사용한 Whale 분산 프레임워크를 활용해 480개의 NVIDIA V100-32GB GPU 클러스터에서 학습을 수행한다.

실험 결과

연구 질문

• RQ1 MoE 모델에서 로드 불균형이 모델 품질에 심각하게 악영향을 미치는가?
• RQ2 활성화된 전문가 수(k)와 전문가 용량(C)은 MoE 성능에 어떤 영향을 미치는가?
• RQ3 계산 비용을 증가시키지 않고도 전문가 프로토타입 기법이 모델 품질을 향상시킬 수 있는가?
• RQ4 초거대 규모 모델(예: 1000억 파라미터 이상)에서 전문가 프로토타입 기법은 얼마나 효과적인가?
• RQ5 트리리언 파라미터 MoE 모델을 TPU 풀 대신 GPU 클러스터에서 효율적으로 학습시킬 수 있는가?

주요 결과

• 기존의 가정과는 달리, 로드 불균형은 MoE 모델의 모델 품질에 심각한 영향을 주지 않는 것으로 나타났다.
• 활성화된 전문가 수(k)와 전문가 용량(C)을 증가시킬수록 모델 성능이 크게 향상된다.
• 전문가 프로토타입 기법을 통해 1.1조 파라미터 MoE 모델의 수렴 속도가 동일 규모의 기준 모델 대비 5배 빨라졌다.
• 1.1조 파라미터 전문가 프로토타입 모델은 480개의 V100-32GB GPU에서 학습되었으며, 2048개의 TPU 코어로 학습된 이전 SOTA 모델들을 능가하는 성능을 달성했다.
• 이 방법은 계산 비용을 일정하게 유지하면서도 특히 대규모 설정에서 모델 품질을 향상시켰다.
• 동일한 FLOPs를 유지하면서도 수렴 속도가 크게 향상되어 학습 동역학의 효율성이 향상됨을 입증했다.