[논문 리뷰] Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model

Shaden Smith, Mostofa Patwary외 18명|arXiv (Cornell University)|2022. 01. 28.

Topic Modeling인용 수 299

한 줄 요약

본 논문은 MT-NLG 530B를 제시하며, 3D(데이터, 텐서, 파이프라인) 병렬성을 사용해 DeepSpeed와 Megatron으로 학습된 가장 큰 단일 구조 트랜스포머 언어 모델로, 인프라스트럭처, 데이터 큐레이션, 훈련 및 0-/1-/부분샷 성능 및 편향에 대한 평가를 상세히 다룬다.

ABSTRACT

Pretrained general-purpose language models can achieve state-of-the-art accuracies in various natural language processing domains by adapting to downstream tasks via zero-shot, few-shot and fine-tuning techniques. Because of their success, the size of these models has increased rapidly, requiring high-performance hardware, software, and algorithmic techniques to enable training such large models. As the result of a joint effort between Microsoft and NVIDIA, we present details on the training of the largest monolithic transformer based language model, Megatron-Turing NLG 530B (MT-NLG), with 530 billion parameters. In this paper, we first focus on the infrastructure as well as the 3D parallelism methodology used to train this model using DeepSpeed and Megatron. Next, we detail the training process, the design of our training corpus, and our data curation techniques, which we believe is a key ingredient to the success of the model. Finally, we discuss various evaluation results, as well as other interesting observations and new properties exhibited by MT-NLG. We demonstrate that MT-NLG achieves superior zero-, one-, and few-shot learning accuracies on several NLP benchmarks and establishes new state-of-the-art results. We believe that our contributions will help further the development of large-scale training infrastructures, large-scale language models, and natural language generations.

연구 동기 및 목표

• 언어 모델의 규모 확장의 필요성을 동기화하고 530B 매개변수의 단일 모듈 트랜스포머를 학습하는 것을 시연한다.
• 3D 병렬성 방법론(데이터, 텐서, 파이프라인)과 효율적인 학습을 위한 토폴로지 인식 매핑을 설명한다.
• 고품질 프리트레이닝 데이터를 만들기 위한 데이터 큐레이션, 전처리 및 혼합을 상세히 기술한다.
• 극대 규모에서의 학습 역학, 하이퍼파라미터 및 안정성 고려사항을 제시한다.
• 제로-/원샷-/소수 샷 설정에서의 평가 결과를 보고하고 편향 및 생성 역량에 대한 관찰을 논의한다.

제안 방법

• DeepSpeed와 Megatron을 이용해 데이터, 텐서, 파이프라인 병렬성을 결합한 3D 병렬화를 채택한다.
• 토폴로지 인식 매핑을 활용하여 노드 간/노드 내 통신을 최적화한다.
• 수천 개의 GPU에 걸쳐 2048 시퀀스 길이와 전역 배치 크기 1920으로 530B 디코더 전용 트랜스포머를 프리트레이닝한다.
• The Pile과 Common Crawl을 포함한 소스에서 중복 제거 및 태스크-데이터 제거를 적용하여 대규모 다양성 데이터셋을 큐레이션 및 전처리한다(대략 339B 토큰 사용; MT-NLG는 270B 토큰으로 학습).
• 혼합 정밀도(16-bit bfloat16)와 특정 하이퍼파라미터를 가진 Adam 옵티마이저를 사용하고; 그래디언트 클리핑과 가중치 감소를 적용하며; 학습률 워밍업과 코사인 감소를 구현한다.
• lm-evaluation-harness 스위트를 사용하여 여러 NLP 태스크에서 제로-/원샷-/소수 샷 프롬프트로 평가한다.

실험 결과

연구 질문

• RQ1 이 모델과 학습 인프라를 확장하여 530B 매개변수 자동회귀 트랜스포머를 효율적으로 학습하려면 어떻게 해야 하는가?
• RQ2 이 규모에서 고품질 프리트레이닝에 필요한 데이터 큐레이션 및 전처리 전략은 무엇인가?
• RQ3 MT-NLG의 제로-/원샷-/소수 샷 능력은 표준 NLP 벤치마크에서 어떠하며, 기존의 거대 언어 모델과 어떻게 비교되는가?
• RQ4 이 규모에서 MT-NLG의 관찰된 특성(예: 편향, 맥락 내 학습)은 무엇인가?

주요 결과

• MT-NLG는 여러 NLP 벤치마크에서 제로-/원샷-/소수 샷 정확도에서 최첨단 성능을 달성하며, LAMBADA에서 모든 설정에서 새로운 SOTA를 확립한다.
• 모델은 여러 태스크에서 맥락 내 학습 및 생성 역량이 강함을 보인다.
• 3D 병렬성(데이터, 텐서, 파이프라인)과 토폴로지 인식 매핑은 수천 대의 GPU에서 530B 매개변수 모델의 효율적인 학습을 가능하게 한다.
• 데이터 큐레이션, 필터링, 중복 제거 및 태스크-데이터 제거가 모델 성능과 안정성의 핵심 요소로 확인된다.
• 프리트레이닝 동안 검증 손실 곡선이 점진적으로 개선되어 270B 토큰 이후 낮은 교차 엔트로피에 도달한다.