서론
2023년도 초반인 현재 ChatGPT 는 장안의 화제이다. 언어의 대한 이해 뿐 아니라 개발자들의 코드를 수정해주는 등 엄청난 품질로 사람들을 놀래키고 있다. ChatGPT 는 GPT3 라는 딥러닝 모델 구조 기반이다. GPT3 는 175billion (약 1750억개) 의 parameter 를 가지고 있는 모델이다. 매우 거대한 모델이라고 할 수 있다. GPT 뿐 아니라 BERT, T5 등 다양한 거대 언어 모델 (Large Language Model) 들이 있다. 아래 그림을 보면 2021 년 이후로는 1 Trillion (1조개) 를 넘어가는 초초거대 모델이 등장할 정도로 모델의 사이즈는 점점 커지고 있는 추세다.
최근 language model 이 계속 커지는 이유는 [2]를 통해 추론해볼 수 있는데, 해당 논문에서는 transformer 기반 모델은 훈련 dataset 을 늘리는 것 보다 모델의 크기를 키우는게 performance 향상에 도움이 된다고 말하고 있다. 하지만, 거대 모델은 그만큼 모델 훈련에 있어 많은 자원을 요구한다. 일단 모델 사이즈가 커서 그를 다 담을 수 있는 memory 용량이 필요하고 또한, 매우 많은 파라미터들을 훈련하는데 많은 시간이 쓰이기 때문에 빠른 서비스를 위해서는 이 훈련 시간을 줄일 수 있는 방법이 필요하다. 그렇다면 현재는 이런 거대 모델 훈련에 어떤 방법을 제시하고 있을까? 본 글은 [3]에 기반하여 거대 모델 훈련을 위한 GPU 기반의 대표적인 병렬 훈련 기법들을 설명하고자 한다.
본문
거대 모델 훈련에서의 Challenges 는 크게 두가지가 있다.
- Memory Capacity (한정된 메모리 용량)
80G A100 이라도, 사이즈가 너무 큰 모델을 GPU main memory 에 한번에 담기는 불가능하다.
또한, training 경우에는 아래 그림처럼 model 말고도 optimization, activation memory 가 추가적으로 필요하다. - Computation Time (긴 계산 시간)
host & device memory 를 swapping 해서라도 memory 용량문제를 해결할 순 있지만 그렇게 하면 말도 안되게 긴 훈련 시간이 필요하다. 예로, GPT3 같은 경우에 하나의 V100 GPU 로 훈련할 시 모델을 훈련시키는데 288년의 시간이 걸린다고 한다.
이런 두 가지 한계점을 극복하기 위해 'GPU 를 1개만 쓰지 않고 여러개를 병렬적으로 돌려서 훈련을 시키자' 가 가장 기본적인 방법이 될 수 있다. 그렇게 나온 병렬 기법은 크게 3가지로 나눌 수 있다.
- Data Parallelism : 훈련 데이터 Batch 를 분할해서 각 device(ex.GPU) 에 할당해서 훈련을 돌리고 결과값을 합치는 것
- Tensor Parallelism : 훈련할 모델을 tensor-wise 로 분할해서 각 device(ex.GPU) 에 할당해 각 device 가 모델 일부분의 훈련값을 담당하는 것, 보통 같은 layer 내부를 분할해서 할당한다. (intra-layer)
- Pipeline Parallelism : 모델을 layer-wise 로 분할해서 각 Layer 를 GPU에 할당한 뒤, batch 를 micro-batch 로 나눠 각 gpu 가 병렬적으로 batch 계산을 수행하는 것
Tensor & Pipeline Parallelism 은 결국 전체 모델을 분할해서 할당하는 것이기 때문에 둘을 합쳐 Model parallelism 으로 분류하기도 한다. 본문에서 이 3가지 기법에 대해 설명하고자 한다.
Data Parallelism
Data Parallelism (편의상 DP로 통칭) 은 말그대로 훈련 data를 병렬 처리하는 기법이다. 각 GPU 에 모델 전체를 할당하고 전체 batch 를 일정양으로 나눠서 각 GPU에 할당하는 것이다. 예로 만약 global batch 가 128 이고 GPU 가 4개인 경우 각 GPU 가 32 개(=128/4)의 데이터를 할당받아서 모델에 넣고 loss 혹은 gradient 를 구하는 것이다. 후에 이 값들을 모두 모아 평균내어 최종 gradient를 구하고 backpropagation 해준 뒤 모든 GPU가 가지고 있는 모델을 다시 동기화 시켜주면 된다. (값을 모으는 것과 최종 동기화를 어디서 어떻게 하느냐는 구현 전략에 따라 다르다) 이를 통해, global batch를 한개의 gpu 에서 계산하느라 걸리는 시간을 1/총 GPU 갯수 로 줄일 수 있어 총 훈련 시간을 줄일 수 있다. 하지만, 단점은 결국 모든 GPU가 모델 전체 데이터를 가지고 있어야 한다는 것이다. 즉, 각 GPU 의 메모리 총량이 전체 모델 데이터 및 훈련 데이터를 감당할 수 있어야 한다. 하지만 초거대 모델로 갈수록 이는 어렵다.
Tensor Parallelism
Tensor Parallelism (TP) 는 [6] 에서 제안된 기법으로 전체 모델을 일부 나눠서 각 GPU에 할당하는 것이다.
예로, GPU 2개에 TP 를 한다고 했을 때, 모델 전체 레이어의 윗 colums (or rows) 부분은 GPU 0번에 할당하고, 나머지 밑 부분은 GPU 1에 할당하는 것이다. 그리고, 중간중간 전체 output (최종 혹은 중간 layer 들의 output) 이 합쳐질 필요가 있을 경우 all-reduce 혹은 all-gather 계산을 통해 sync 를 맞춰준다. 이는 layer 별로 나누는 것이 아니고 같은 layer 내에서 GPU 할당이 달라지는 경우이기 때문에 intra-layer (레이어 내부) parallel 으로 불린다.
모델을 나눠서 gpu 에 할당하니까 memory 용량 문제를 해결할 수 있고, 각 gpu 마다 계산량이 줄어들 수 있다. 하지만 중간중간 all-reduce 와 같은 과정이 필요하게 되는데 이는 각 GPU 간의 (혹은 각 server 간의) 통신을 요구하기 때문에 commnication cost(통신비용)가 증가한다. 통신 비용을 이야기 할 때, 우리는 intra-server 통신 (서버 내부 device 들의 통신) 과 inter-server (외부 서버의 gpu 간의 통신) 을 나눠 이야기할 수 있다. 간단히 intra-server 는 같은 서버 내부에 있는 GPU 들끼리 데이터를 주고받는 것인데, 이것은 Nvidia NVLink 와 같은 멀티 GPU 통신을 지원하는 기술을 통해 상대적으로 속도가 빠르다. 그렇지만 inter-server 통신은 다른 물리적으로 떨어져 있는 server 의 GPU 간의 통신을 뜻하는데 이는 infiniband 와 같은 초고속 연결망을 지원해도 intra-server 에 비하면 매우 느리다. 근데 TP 는 all-reduce 와 같은 비용이 큰 통신을 요구하므로 inter-server 보다 intra-server 만 지원하도록 하는 것이 좋다. 실제로 [3] 에서는TP 분할 갯수를 서버 내부에 있는 GPU 갯수로 제한하는 것을 추천하고 있다.
Pipeline Parallelism
Pipeline Parallelism (PP) 은 모델을 layer 별로 나눠서 할당하는 것이다. 예로 아래 그림처럼 GPU가 6개 있을 때, 한 DP group 에서 PP 를 수행하면, DP group1 에서 GPU0 번째는 모델의 레이어 1,2번째를 할당하고 는 GPU1 에는 3,4 레이어 할당 GPU2 에는 layer 5,6을 할당하는 것이다. 이렇게 할당하고 전체 batch 를 micro-batch로 잘게 나눠서 전체 batch 를 pipeline 화하여 계산하는 것이다.
그렇게 되면, GPU0이 micro-batch1을 layer1~2 를 계산하고 계산된 값을 GPU1 으로 전달해서 GPU1이 micro-batch1 을 layer3~4에 대해 계산할 땐, GPU0 은 동시에 놀지 않고 micro-batch2 에 대해서 계산을 처리하는 것이다. 이를 통해 각 GPU 들의 노는 시간(idle time)을 줄이는 것이다.
아래는 pipelining 한 forward-backward 가장 간단한 스케듈링 그림이다. 즉 PP는 각 GPU에 필요한 memory 용량을 줄이고, 계산 시간도 각 micro-batch를 병렬적으로 수행함으로서 줄일 수 있다는 장점이 있다. 하지만, 그림을 보면 알 수 있듯이, 결국엔 synchronizaion 을 위해 앞 Layer 들 그리고 뒷 Layer 들의 계산을 기다려야하는 구간인 pipeline-bubble 이 존재할 수 밖에 없다. 이 pipeline-bubble 을 최대한 줄이면서 GPU utilization 은 높이고 총 계산 시간을 줄이는 것이 PP의 최적화의 목표이다. 그에 대한 연구도 활발하게 진행되는 중이다. 또한 synchronization 을 위해 pipeline flush 라는 과정이 매 batch 끝에 필요하다.
참고
[1] https://amatriain.net/blog/transformer-models-an-introduction-and-catalog-2d1e9039f376/
Transformer models: an introduction and catalog — 2023 Edition
amatriain.net
[2] Scaling Laws for Neural Language Models (arXiv, Jan 2020)
[3] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (SC, Apr 2021)
[4] Low-Memory Neural Network Training : A Technical Report (arXiv, Apr 2022)
Parallelisms — NVIDIA NeMo
docs.nvidia.com
[6] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (GTC, Mar 2020)
