Megatron-LM
최근 언어모델에서는 큰 트랜스포머 모델을 학습하는게 중요하다. 엄청나게 빠르고 메모리가 큰 GPU가 있다면 걱정이 없겠지만 Google, Facebook, Nvidia, OpenAI등 과 같이 글로벌 빅테크 기업들이 아니라면 대부분의 사람들에게 컴퓨팅과 메모리 자원이 제한적이다.
Megatron-LM 논문에서는 매우 큰 트랜스포머 모델을 학습하기 위한 테크닉들을 보인다. intra-layer model 병렬화를 이용해 쉽고, 효율적으로 적용할수 있으며, 새로운 컴파일러나 라이브러리의 변경이 필요하지 않다. 제안된 방법들을 적용하기 위해서 최소한의 native pytorch 변경을 통해 적용할수 있다.
제안된 방법을 통해서 저자들은 512GPU를 이용해 8.3 Billions 파라미터를 가진 모델을 학습했다. 8.3B을 가진 GPT-2와 3.9B을 가진 BERT를 학습했고, WikiText103, LAMBDA, RACE 에서 SOTA를 달성했다 .
Introduction
최근에 큰 모델들의 중요성이 커지고 있으나 트랜스포머 모델이 많은 메모리를 사용하며, GPU가 가진 메모리또한 제한 적이기 때문에 추가적인 메모리 관리 테크닉들이 요구된다.
딥러닝 학습에 많이 사용되는 ADAM optimizer의 경우 모델을 학습할때 파라미터마다 모멘텀과 다른 optimizer state들을 저장하기 위해 추가적인 메모리를 사용한다. 그리고 다른 모델 병렬화 방법들은 프로세서 위에 올라가 있지 않아도 되는 optimizer states들과 weight들을 나누는 방법을 통해 큰 모델들을 학습한다.
예시로는 GPipe와 Mesh-Tensorflow가 있으며, 다른 종류의 모델 병렬화를 프레임워크로 제안한다.
하지만 위와 같은 방법들은 모델들을 다시 작성해야하는 단점과 컴파일러와 프레임웍 부분에서 아직 개발 중인 부분도 있다.
Megatron-LM의 저자들은 intra-layer model-parallelism을 이용해서 단순하면서도 효율적인 구현을 제안하고 있으며 트랜스포머 모델을 이용해 c++ 코드 변경 없이, pytorch에서 학습할수 있는 효율적인 방법을 제안한다. 이러한 접근 방법은 GPipe와 같은 모델 병렬화 방법과는 다르다.
제안된 방법의 확장성을 증명하기 위해 v100 32GB GPU에서 1.2B의 모델의 베이스라인으로 학습시켰으며, 베이스라인으로 부터 8.3B까지 모델을 512GPU에서 8가지 모델 병렬화 방법을 적용했다.
Megatron-LM Contribution
- 단순하고 효율적으로 모델 병렬화를 구현: 기존에 pytorch 트랜스포머 구현체에서 몇가지 부분들을 수정해서 쉽게 적용가능.
- 모델과 데이터 병렬 테크닉에 대해 깊이 있게 분석: 512GPU를 통해 76% 확장 효율성을 검증.
- BERT와 같은 모델에서 layer normalization의 배치가 모델의 증가에 따라 정확도에 미치는 영향을 보임
- 모델사이즈 증가에 따른 모델의 성능향상 확인
- 몇몇 태스크들에서 SOTA를 달성
- 학습과 평가 파이프라인 코드를 오픈소스로 공개: https://github.com/NVIDIA/Megatron-LM
Background and Challenges
Neural Language Model Pretraining
최근의 언어모델들은 단지 워드 임베딩에 머무르는것이 아니라 수십억 파라미터를 가진 언어모델들로 진보하고 있다. 이러한 변화에 따라 하드웨어, 시스템 테크닉, 확장에 효율적으로 운영될수 있는 프레임워크들이 요구된다. 그리고 컴퓨팅 파워의 증가에 대한 대응도 필요하다.
Transformer
파란부분은 단일 트랜스포머 레이어를 나타내며, 보라색 부분은 fully connected layer에 해당된다.
오리지날 트랜스포머에서는 ReLU function과 outputs에 layer normalization을 적용했다면, Megatron-LM에서는Multi-Head Attention, feed-forward layers의 inputs에 layer normalization을 사용했다.
Data and Model Parallelism in Deep Learning
딥러닝에서 학습에서 많은 하드웨어 가속기에서 학습하기 위해서 크게 2가지 패러다임이 존재한다:
1️⃣ Data Parellelism: 미니 배치를 여러 워커에 분산시키는것
2️⃣ Model Parallelism: 모델의 메모리 사용과 계상 부분을 여러 워커에 분산 시키는것.
미니배치를 적절하게 증가시켜서 워커들에게 적절하게 나누어 주는 방법은 트레이닝 데이터의 증가에 따라 거의 선형적으로 증가하는것을 볼수 있다. 하지만 큰배치 학습이 최적화 과정에서 문제를 보일수 있다: 모델의 정확도를 낮추거나 모델이 수렴하는데 더 많은 시간이 필요로 하게된다. 이러한 점들 때문에 학습시 처리량 증가의 이점이 상쇄 된다.
또 다른 연구에서는 Data parallelism 과 activation checkpointing을 결합한 방법을 사용하였다. 메모리 사용량을 줄이기 위해 activation을 저장하지 않고 backward 과정에서 다시 계산하는 방법을 사용했다. 하지만 이러한 방법들도 문제가 있었으며 다른 솔루션에서는 메모리가 차지하는 공간을 줄이기 위해 파라미터를 공유하는(Parameter Sharing)을 사용한다. 그러나 이것 또한 모델의 전체 크기에 제한되었다는 단점이 있다.
Megatron-LM에서는 여러 가속기들에 모델을 여러개로 나누는 model parallelism을 사용한다. 이 방법을 통해 메모리를 보다 효율적으로 사용하고 마이크로 배치의 크기와 상관없이 병렬화를 적용할수 있다.
model parallelism의 2가지 패러다임
model parallelism에는 layer-wise pipeline parallelism과 텐서 계산을 분산해서 하는 방법이 있다.
🚰 pipeline model parallelism
파이프라인 모델 병렬화는 연산들의 그룹이 다른 디바이스로 전달되기 전에 수행된다. 하지만 이러한 방법들도 일관적이지 않은 문제가 있었지만 GPipe에서는 synchronous gradient deccent을 사용해서 일관적이지 않은 이슈를 극복했다.
통신과 계산연산들의 효율적인 파이프라인을 다루기 위해서는 추가적인 로직이 필요하며 파이프라인 버블로부터 효율성이 줄어들고, 옵티마이저의 변경으로 인해 정확도에 영향을 줄수도 있다.
🧽 Distributed tensor computation
Distributed tensor computation 의 경우 여러 디바이스를 통해서 계산을 가속시키거나 모델을 크기를 증가하기 위해 텐셔 연산을 분할하는 방법은 수직적이며 보다 일반적인 접근 방법이다.
이전에는 FlexFlow, Mesh-Tensorflow와 같은 방법들이 있었으며, Megatron-LM에서는 Mesh-Tensorflow의 아이디어를 이용해서 기존의 Pytorch Transformer 구현체에 몇가지 수정을 통해 모델 병렬화를 구현할 수 있었다.
