[리뷰] Scalable and Efficient MoE Training for Multitask Multilingual Models
GPT-3 (175B) 파라미터 모델은 256 NVIDIA A100 GPUs로 학습할 경우, 168 일이 걸리며 1T 파라미터 모델을 학습하는데 512 A100 GPUs로 1년이상이 걸린다. MoE 모델은 GPT-3와 같은 dense 모델을 학습하는데 필요한 엄청난 비용에 대한 잠재적인 솔루션을 제공한다. 예를 들어 Switch Transformer(1.6T)는 10B 파라미터를 학습하는데 필요한 compute budget보다 더 적은 budget을 요구한다.
System Challenges
MoE 모델에서 parameters의 개수는 두 가지 factor에 의해 결정된다.
- 1) Expert 레이어없는 기본 모델(Base Model)의 사이즈
- 2) 모델에서 experts의 개수
1)을 증가시키는 것은 모델 사이즈와 계산 비용을 둘다 증가시키고, 2)는 모델 사이즈를 증가시키지만 계산 비용을 증가시키지 않는다. 하지만 experts의 수만을 늘리는 것만으로는 품질 향상에는 한계가 있다. MoE 모델의 정확도를 달성하면서 계산 비용을 최소화하기 위해서 Base Model과 experts의 수를 모두 스케일링해야 한다.
Dense 모델을 스케일링하기 위해 디자인된 tensor-slicing과 pipeline parallelism와 같은 기술은 MoE 모델의 sparse 특성때문에 experts의 개수를 스케일링하는데 효과적이지 않다. 반면 experts의 개수를 증가시켜 모델 사이즈를 스케일링하기 위해 디자인 된 expert parallelism은 non-expert parameters에는 적용되지 않는다. 확장가능한 MoE 시스템을 만들기 위해서 모델의 experts와 non-experts 둘다를 대상으로 parallelism 기술을 결합해야 하는 것이 필요하다.
GShard와 Switch Transformer와 같은 기존 MoE 시스템은 expert, data, model parallelism을 모두 조합하거나 3가지 parallelism 중 몇 개를 조합하여 사용한다. 그러나 이러한 방식은 두가지 주요 제약이 있다.
- 1) Data parallel GPUs에 base model(expert parameters 없는 모델의 일부)을 복제하므로 메모리가 낭비된다.
- 2) Base model을 1.4B parameters 이상으로 스케일링하기 위해 model parallelism을 요구하며 모델 코드 리팩토리링이 요구된다.
이 제한사항에 추가하여 모든 기존 MoE 시스템은 잠재적으로 디바이스 메모리 wall에 의해 제한된다. 이것은 dense 모델과 극명하게 대조되는데 dense 모델의 경우, 학습 가능한 가장 큰 모델 크기는 오늘날의 컴퓨팅 클러스터에서 모델을 학습하는데 걸리는 시간으로 제한된다. (1T parameters 모델은 512 A100 GPUs 조차도 학습하는데 2년이 걸릴 수 있다.)
MoE는 수 T parameters 모델조차 같은 클러스터에서 몇 주이내에 학습할 수 있다. 하지만 그러한 모델은 512 A100 GPUs 메모리에 맞지 않으므로 parallelism의 선택과 관계없이 MoE 모델 scale은 가용 디바이스 메모리에 의해 근본적으로 제한된다.
학습 & 추론 Challenges
MoE 모델은 학습과 추론 방법 측면에서 다음과 같은 challenges를 가지고 있다.
- 1) MoE 모델은 제한된용량과 불균형한 expert 활용이라는 본질적인 문제가 있다. 더 높은 capacity factor를 사용하는 것에 의해 완화시킬 수 있지만 GPU 간 중복 계산과 통신을 추가하여 비효율적인 학습을 초래한다.
- 2) 모델 parameters의 개수를 증가하는 것에 의해 모델이 overfitting 되지 않고 정해진 데이터 양으로 더 많은 일반적인 표현을 학습하는 것이 중요하다. 많은 정규화(Regularization) 기술이 수년이상 개발되었음에도 불구하고 MoE 아키텍처는 아키텍처의 고유한 특성을 활용하여 모델을 정규화할 수 있는 좋은 기회를 제공한다. 더 모델 parameters의 수가 증가함에 따라,sample 효율적이고 안정적인 학습을 달성하기 위해 모델 parameters의 초기화가 중요하다. 하지만 지금까지 제대로 연구되지 않고 있다.
MoE 모델의 expert 레이어들은 한번에 오직 1 or 2개의 experts만을 활성화하여 일정 수준의 자율성과 전문성을 자연스럽게 학습한다. 이는 모델 초기화와 런타임 효율성을 위해 더 나은 모델 방법을 찾기 위해 흥미로운 기회를 제공한다.
MS가 제안하는 MoE를 위한 솔루션
- DeepSpeed MoE: Expert 와 Non-Expert parameters을 대상으로 5개의 다른 형태의 parallelism간 빈틈없는 오케스트레이션을 통해 MoE 모델을 수 Trillion parameters로 스케일링한다. 동시에 GPU 메모리 wall을 뛰어넘기 위해 CPU 메모리를 사용한다.
- Overflow된 토큰을 더 효율적으로 처리하고 샘플 효율성을 개선하는 토큰 선택을 위한 효율적인 학습 방법 & 학습 효율성을 위해 experts와 추론 시간을 개선하기 위해 expert pruning 전략을 합친 구성 가능한 요소로써 experts를 활용하는 방법
- 멀티태스크와 다국어 언어모델을 MoE 모델 아키텍처로 scale-up하기 위한 효과적인 training 레시피
MS는 Z-code M3(Multilingual Multitask MoE)라고 불리우는 높은 효율의 large-scale 언어모델의 패밀리를 제안된 DeepSpeed MoE 시스템과 효과적인 학습 방법 및 레시피를 활용하여 학습한다. Z-code M3는 50개 언어에 대한 10B parameters와 multiple pre-training task를 가진 인코더-디코더 모델로 학습되었다. MS는 기계 번역 및 교차 언어 요약(cross-lingual summarization)과 같은 다양한 다운스크림 task에서 모델 성능을 탐색한다.
DeepSpeed MoE: A scalable system for training MoE Models
DeepSpeed MoE에 적용가능한 Parallelism 방식
DeepSpeed MoE는 Expert, Data, Model(tensor-slicing), ZeRO, ZeRO-Offload를 조합하여 5개 형태의 parallelism을 지원한다. 이중에서 Expert과 ZeRO, ZeRO-offload 방식은 다음과 같이 동작한다.
Expert Parallelism: expert parallelism은 다른 experts를 다른 디바이스에 위치하여 병렬로 실행한다. Experts가 다른 GPU 디바이스상에 위치하기 때문에 All-to-All primitive를 사용하여 명시적인 comm.이 필수적이다. Expert parallelism은 experts의 개수를 증가시켜 모델 사이즈를 확장시키며, non-expert 파라미터를 스케일링하기 위해 model parallelism을 사용한다.
Data Parallelism(DP)은 다수의 GPU를 통해 스케일링한다. 이 접근 방식은 data parallelism 디바이스 전반에 걸쳐 model states를 복제하여 i) DP 전반에 걸친 메모리 중복성(Memory Redundancy) 및 ii) 주요 코드 리팩토링이 필요한 non-expert 파라미터를 증가시키기 위해 모델 병렬화를 요구합니다.
ZeRO(Zero Redundancy Optimizer): 위의 문제를 해결하기 위해 expert, model, data parallelism을 추가하기 위해 DeepSpeed MoE는 ZeRO parallelism을 지원한다. DeepSpeed MoE는 optimizer states와 gradients를 분할하는 ZeRO-2를 사용한다.
DeepSpeed MoE는 data parallelism를 ZeRO로 대체하여 data parallelism 과정에 걸친 optimizer state & gradient redundancy를 제거한다. 대규모 transformer 모델의 optimizer states & gradients는 학습시간 동안 전체 메모리의 85% 이상을 소비한다. 그들을 복제하기 전에 GPUs 간에 이들을 분할하여 non-expert 파라미터에 크게 늘릴 수 있도록 메모리를 확보한다.
DeepSpeed MoE는 Switch Transformer의 non-expert 파라미터들보다 더 큰 15B non-expert 파라미터까지 지원할 수 있다. 이는 model parallelism의 복잡성을 추가하지 않고 1.4B의 non-expert 파라미터만 지원할 수 있는 기존 MoE 시스템에 비해 10배 이상 크다. Model parallelism과 결합하면 기본 모델(MoE 모델의 non-expert 부분)만으로도 100B 이상의 매개변수를 가질 수 있으며 이는 기존 시스템에서는 불가능합니다.
ZeRO-Offload로 GPU 메모리벽을 초월
DeepSpeed MoE는 ZeRO를 사용하지만 여전히 aggregated 디바이스 메모리에 제한을 받는다. 따라서 Deepspeed MoE는 GPU 메모리를 뛰어넘기 위해 ZeRO의 확장인 ZeRO-Offload를 지원한다.
앞에서 이야기한 것과 같이 optimizer states와 gradients는 전체 GPU 메모리에 85%를 차지하므로 이들을 CPU 메모리로 offload할경우, 총 GPU 메모리 소비가 크게 줄어들어 DeepSpeed MoE는 대규모 모델을 지원할 수 있다. 예를들어 DeepSpeed MoE는 512 NVIDIA A100 GPUs 상에 3.5T 파라미터 이상을 가진 MoE 모델을 지원하여 현존하는 MoE 시스템(0.4T)과 비교하여 8배 더 큰 모델 사이즈를 지원한다. 다른 이야기로는 DeepSpeed MoE는 8배 더 적은 리소스로 동일한 모델 스케일을 달성한다. (512 GPUs → 64 GPUs)
Efficiency
DeepSpeed MoE는 탁월한 계산 효율성을 달성하며 다음 조건에서 120 TFLOPS를 달성한다.
- transformer MoE with hidden size of 4096
- batch size per expert of 8192 tokens
Expert 당 batch size는 global batch size와 experts의 개수, top-1 vs top-2 게이팅의 함수이며, MoE 아키텍처를 디자인하고 학습을 위한 batch size를 고르기 위한 파라마터를 고려해야 한다.
Multi-GPU Scalability
DeepSpeed MoE는 고도로 최적화된 expert parallelism과 ZeRO 구현으로 인해 거의 완벽한 확장성을 제공하며 효율적인 comm. primitives를 신중하게 사용하여 통신 overhead를 최소화시킨다. 아래 그림에서 보는 것과 같이 GPU 수의 증가에 따라 throughput이 거의 선형으로 증가한다.
레퍼런스
[1] Scalable and Efficient MoE Training for Multitask Multilingual Models
