[논문 리뷰] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

Mohammad Shoeybi, Mostofa Patwary외 4명|arXiv (Cornell University)|2019. 09. 17.

Topic Modeling참고 문헌 46인용 수 821

한 줄 요약

Megatron-LM은 PyTorch에서 층 내 모델 병렬성을 통해 수십억 매개변수의 트랜스포머 모델을 훈련시키고, 512 GPU에서 8.3B 매개변수를 달성하며 강한 확장성과 여러 NLP 벤치마크에서 최첨단 결과를 달성합니다.

ABSTRACT

Recent work in language modeling demonstrates that training large transformer models advances the state of the art in Natural Language Processing applications. However, very large models can be quite difficult to train due to memory constraints. In this work, we present our techniques for training very large transformer models and implement a simple, efficient intra-layer model parallel approach that enables training transformer models with billions of parameters. Our approach does not require a new compiler or library changes, is orthogonal and complimentary to pipeline model parallelism, and can be fully implemented with the insertion of a few communication operations in native PyTorch. We illustrate this approach by converging transformer based models up to 8.3 billion parameters using 512 GPUs. We sustain 15.1 PetaFLOPs across the entire application with 76% scaling efficiency when compared to a strong single GPU baseline that sustains 39 TeraFLOPs, which is 30% of peak FLOPs. To demonstrate that large language models can further advance the state of the art (SOTA), we train an 8.3 billion parameter transformer language model similar to GPT-2 and a 3.9 billion parameter model similar to BERT. We show that careful attention to the placement of layer normalization in BERT-like models is critical to achieving increased performance as the model size grows. Using the GPT-2 model we achieve SOTA results on the WikiText103 (10.8 compared to SOTA perplexity of 15.8) and LAMBADA (66.5% compared to SOTA accuracy of 63.2%) datasets. Our BERT model achieves SOTA results on the RACE dataset (90.9% compared to SOTA accuracy of 89.4%).

연구 동기 및 목표

• 단일-GPU 메모리 한계를 넘어 수십억 매개변수의 언어 모델 훈련을 동기부여한다.
• 최소한의 변경으로 PyTorch에 맞게 간단하고 효율적인 층 내 모델 병렬 접근 방식을 개발한다.
• GPT-2 및 BERT 유사 아키텍처에서 스케일링 효율성과 성능을 평가한다.
• 언어 모델링 및 다운스트림 작업에서 최첨단 결과를 시연하고 오픈 소스 코드를 공개한다.

제안 방법

• 최소한의 동기화로 MLP와 셀프 어텐션의 GEMM을 GPU 간에 분할하여 층 내 모델 병렬성을 구현한다.
• 효율적인 확장을 가능하게 하기 위해 트랜스포머 계층당 두 개의 순방향 및 두 개의 역전파 All-Reduce 통신을 사용한다.
• 어휘 차원에 따라 입력 및 출력 임베딩 행렬을 병렬화하여 GPU 간 통신을 줄인다.
• 추가 통신을 피하기 위해 각 GPU에서 중복된 계층 정규화와 잔여 연산을 유지한다.
• 메모리 효율성을 위해 동적 손실 스케일링과 활성화 체크포인팅이 있는 혼합 정밀도 트랜스포머 모델을 훈련한다.

실험 결과

연구 질문

• RQ1 PyTorch에서 층 내 모델 병렬성이 커스텀 컴파일러 없이 트랜스포머 모델을 수십억 매개변수로 확장할 수 있는가?
• RQ2 매개변수를 수십억으로 확장했을 때 표준 NLP 벤치마크에서 GPT-2 및 BERT 유사 모델의 성능에 모델 크기가 어떤 영향을 미치는가?
• RQ3 모델이 커질수록 성능을 유지하거나 향상시키기 위해 필요한 아키텍처 조정(예: 계층 정규화 배치)은 무엇인가?
• RQ4 수십억 매개변수까지 수백 대의 GPU에서 훈련할 때 실제적인 확장성 한계(FLOPs, 처리량, 효율성)는 무엇인가?

주요 결과

• 512 GPU에서 8.3B 매개변수까지 교육하며 8방향 모델 병렬화를 사용하고 최대 15.1 PetaFLOPs를 지속하고 강력한 단일-GPU 기준선(39 TeraFLOPs)에 비해 76%의 스케일링 효율을 달성한다.
• 모델 병렬성은 모델 전용 및 모델+데이터 구성 모두에서 강한 약한 확장을 가져오며 가장 큰 구성에서 74%–77%의 스케일링 효율을 달성한다.
• BERT-유사 모델에서 계층 정규화를 신중하게 배치하면 모델 크기가 커짐에 따라 성능이 단조롭게 향상된다.
• GPT-2 스타일 모델은 8.3B 매개변수로 WikiText103에서 최첨단 퍼플렉시티(10.81)와 강력한 LAMBADA 정확도(66.51%)를 달성하며; RACE 정확도도 더 큰 BERT/Megatron 설정에서 90.9%로 향상된다.
• 제안된 아키텍처와 훈련 방식으로 최대 3.9B 매개변수의 BERT 스타일 모델은 여러 GLUE 유사 작업과 RACE에서 최첨단 개발 결과를 달성한다.