GPT-3보다 3배 더 큰 언어모델: MT(Megatron-Turing)-NLG (530B)
MS는 NVIDIA와 협력을 통해 기존 GPT-3(175B)보다 3배 더 큰 530B 파라미터의 Megatron-Turing Natural Language Generation Model(MT-NLG)를 발표하였다.
Large-Scale 언어모델 학습의 문제점
Large-Scale 언어모델은 빠르게 파라미터 개수를 증가시켜 모델의 성능을 개선하고 있다. 하지만 Large-Scale 언어모델을 학습하기 위해선 2 가지 challenge가 있다.
- 1) Single GPU 메모리 용량의 부족
- 2) 알고리즘/SW/HW 스택을 모두 함께 최적화하지 않으면 학습 시간이 비현실적으로 증가함
MS와 NVIDIA는 위와 같은 Large-scale 언어 모델 학습의 문제점을 해결하고자 NVIDIA의 GPU 가속 학습 인프라스트럭처(Megatron-LM)와 MS의 분산 학습용 SW 스택(DeepSpeed)을 통합하여 전례없는 학습 효율성을 달성하였다.
또한 두 회사는 학습을 위해 필요한 수천억개 Token을 가진 고품질의 자연어 학습용 말뭉치(corpora) 셋을 구축하고 최적화 효율과 안정성을 개선하기 위한 학습 레시피를 공동 개발하였다.
MT-NLG의 모델 구조 및 데이터셋
MT-NLG는 MS의 T-NLG(17B)과 Megatron-LM을 계승하였으며 GPT-3(175B)보다 3배 더 큰 Transformer기반 언어모델이다. MT-NLG는 다음과 같은 자연어 task에서 SOTA를 달성하였다.
- Completion prediction
- Reading comprehension
- Commonsense reasoning
- Natural language inference
- Word sense disambiguation
MT-NLG의 모델 구조
MT-NLG는 GPT-3와 유사하게 transformer decoder를 사용하였다. 다음은 MT-NLG와 GPT-3의 모델 구조를 비교한 결과이다.
MT-NLG의 구체적인 학습 방법은 다음과 같다.
- Tensor parallelism : 8-way
- Pipeline parallelism: 35-way
- Sequence length: 2048
- Global batch size: 1920 (초기 global batch size는 32)
- Learning rate warmup 시 1 Billion Token 사용
MT-NLG의 데이터셋
총 339B Token으로 구성된 15개의 데이터셋을 사용하여 MT-NLG를 학습하였다. 데이터셋의 용량은 총 1.5TB에 해당한다고 한다.
MT-NLG을 위한 Large-scale 학습 인프라스트럭처
HW system
MT-NLG를 학습하기 위해 NVIDIA Selene와 Microsoft Azure NDv4와 같은 SOTA 슈퍼컴퓨팅 클러스터가 필요하다.
1)NVIDIA
- Selene (DGX SuperPOD x4. DGX A100 560대)
- 4480 GPUs (A100 80GB ver. 총 358.4 TB)
- HDR Infiniband in a full fat tree configuration
2)Microsoft Azure NDv4
- NVIDIA Selene와 유사한 구조
SW Design
530B의 MT-NLG를 학습하기 위해 NVIDIA Megatron-LM와 MS DeepSpeed를 조합하여 data/pipeline/tensor-slicing parallelism을 구현하였다.
- Intra Node 스케일: Megatron-LM의 tensor-slicing parallelism 이용
- Inter-Node 스케일: DeepSpeed의 pipeline parallelism 이용
- Training Speed 스케일: DeepSpeed의 data parallelism 이용
예를들어 530B인 MT-NLG 모델을 스케일하기 위해 8-way tensor-slicing, 35-way pipeline parallelism을 사용하였다. 즉, Model parallelism을 위해 A100 GPU 280대를 사용하였다.
System Throughput
530B 모델 학습 시 시스템의 end-to-end throughput은 Data Parallelism 수준이 증가할수록 iteration time은 감소하나 GPU의 throughput(GPU utilization는 40.3% → 38.7% → 36.2%)은 낮아지는 것을 알 수 있다.
Data Parallelism 수준이 높아질수록 GPU의 throughput이 떨어지는 것은 파라미터를 aggregation하기 위한 연산과 노드간 통신 시간이 늘어남에 따라 발생하는 것으로 예측된다.
시사점
AI 모델의 하이퍼스케일링이 더 나은 성능으로 가져다줌으로써 AI 모델의 발전이 무어의 법칙을 훨씬 능가하면서 더 많은 연산능력을 필요로 하고 있다. 현재 Google의 TPU를 제외하고 학습에서 사용할 수 있는 HW 리소스는 GPU가 유일하다. Large-scale 모델 학습을 위해 다음과 같은 업계의 움직임이 전망된다.
1. HW 리소스 효율화
Megatron-LM과 DeepSpeed를 조합하는 것과 같이 Large-scale 모델 학습에 최적화된 SW 프레임을 이용하여 GPU의 throughput을 향상시키고 이기종 컴퓨팅을 도입하여 HW 리소스를 보다 효율적으로 사용할 것이다.
2. Hyperscaler들의 자체 Training 칩 개발
Large-scale 모델 개발에 많은 시간과 비용이 들기 때문이며 많은 Hyperscaler가 자체 AI 가속기를 개발하여 고가의 GPU 도입에 따른 비용을 절감하고 AI 모델 개발 시간을 단축할 것이다.
3. Large-scale 모델에 적합한 메모리 풀 구조
높은 학습비용은 compute와 memory의 coupling으로부터 발생한다. 최근 Cerebras가 compute와 memory를 decoupling 시켜 Model Parallelism이 없이 Data Parallelism으로 학습할 수 있도록 MemoryX와 SwarmX와 같은 새로운 시스템을 소개하였다. Large-scale 모델 학습에 가장 필요한 자원은 메모리로 거대한 메모리 풀을 효율적으로 사용할 수 있는 구조가 연구될 것으로 예상된다. 예를들어 NVIDIA의 Grace CPU는 고속의 NVLINK를 이용하여 CPU의 메모리를 Far 메모리로 사용하여 부족한 GPU 메모리 용량을 확보할 수 있을 것이다.
레퍼런스
[1] Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, the World’s Largest and Most Powerful Generative Language Model
[2] Behold the Megatron: Microsoft and Nvidia build massive language processor
