분산학습 대표 유형: DP vs MP

안녕하세요~!

오늘은 분산학습 시리즈 2번째로 분산학습의 대표 유형들에 대해 살펴보는 시간을 가지겠습니다.

분산학습을 잘 모르시는 분들을 위해 개괄적으로 설명할 예정으로 세부적인 개별 구현에 대해서 다루기보다는 각 유형의 컨셉에 대한 이해 그리고, 해당 기법과 관련된 용어들을 소개하는 것을 목적으로 하겠습니다.

 

지난 시간에 학습 데이터와 모델 사이즈가 점점 커지는 학습 트렌드를 설명하면서 여러 서버의 GPU들을 사용해 학습하는 분산학습(multi-node training)이 도래하게 되었다고 말씀드렸습니다. 하나의 GPU만 사용한다면, 모델이 logit을 계산하는 것부터 loss를 구하고 optimizer로 모델을 업데이트하기까지의 모든 metric 연산들은 GPU에 올려서 수행하기만 하면 됐습니다. 그러나 분산학습을 위해서 이제 이 모든 연산들을 여러 GPU에 나누어 학습해야 합니다. 어떻게 학습하는 것이 가장 효율적일까요?

 

정답이 있진 않지만 다양한 방식들이 존재하고 있습니다.

Data Parallelism

Data Parallelism은 지난 시간에도 한번 설명드렸듯이

(1) 모든 GPU가 동일하게 모델 정보를 가지고

(2) 데이터를 쪼개서 GPU가 개별적으로 학습을 진행하는 방식입니다.

실제로 하나의 batch set을 어떻게 학습하는 걸까요?

모델을 학습하기 전에 각 GPU로 동일한 모델과 나눠진 batch set 일부를 전송하게 됩니다. (scatter 연산)

GPU는 자신에게 주어진 데이터로 Forward -> Backward로 gradient를 계산하고,

모든 GPU의 gradient를 모아 합산한 최종 gradient를 모두에게 돌려주면 끝이 납니다. (all-reduce 연산)

 

자, 이게 끝입니다! 컨셉은 되게 명료하죠?

재밌는 것은 pytorch 로 data parallelism을 구현하기 위해 검색하다보면 다음과 같은 두 메소드`DataParallel`, `DistributedDataParallel`를 확인할 수 있습니다. 둘 다 Data Parallelism을 구현하는 메소드이지만, 약간의 차이가 있습니다.

 

`DataParallel` 은 연산과정은 다음과 같습니다.

forward -> (gather) -> loss -> (scatter) -> backward -> (reduce) -> gradient

loss 에서 gradient를 계산하기까지 결과물을 한 곳으로 모으는 gather, reduce 통신이 동반되기 때문에, 이 작업을 주관하는 master node에 많은 부담이 실리게 됩니다. (보통은 GPU 0번이 master node가 됩니다)

https://levelup.gitconnected.com/quick-primer-on-distributed-training-with-pytorch-ad362d8aa032

`DataParallel`은 single process, multi thread 로 동작하기 때문에 실제로 multi-node training에는 사용하기 어렵고, 하나의 호스트에 여러 GPU를 효율적으로 사용하고 싶을 때 즉, single node multi GPU 환경에서 유용합니다.

 

반대로 DistributedDataParallel 의 연산 프로세스는 다음과 같습니다.

forward -> backward -> (all-reduce) -> gradient

https://levelup.gitconnected.com/quick-primer-on-distributed-training-with-pytorch-ad362d8aa032

`DistributedDataParallel`은 multi process 로 동작하며, 특정 노드에 메모리가 과중되지 않습니다. 파이썬의 경우 GIL 문제가 있어  multi-thread 보다 multi-process 환경에서 연산 속도가 더 빠른 경우가 많으므로 gpu가 아닌 cpu-intensive 한 job 들이 많을 경우에도 `DistributedDataParallel`이 더 유용하게 동작할 수 있습니다.

 

Model Parallelism

이름만 봐도 짐작이 가시죠? Model Parallelism은 Data Parallelism과 다르게 각 GPU마다 Model을 나누어 학습하는 방법을 일컫습니다.taParalle

Data Parallelism은 어떤 모델이던 자유롭게 적용할 수 있는 방법이면서 많은 데이터를 빠르게 학습하는 것이 가능합니다. 그러나 모든 GPU가 똑같이 모델을 들고 있어야 한다는 점에서 여전히 사이즈가 큰 모델을 학습하기엔 부족합니다. 

반면 Model Parallelism 은 모델을 쪼개서 학습하기 때문에 큰 모델도 학습이 가능합니다. 그러나 모델을 나눈다는 것은 데이터를 나누는 것보다 큰 비용이 들고, 모델을 잘 이해하고 있는 사람이여야 가능한 작업입니다. 

 

대충 머리로만 생각해도 모델을 이루는 수만개의 weight들을 어떻게 나눠야 할지 조금 막막해지는 면이 있습니다. 모델은 각 유형에 따라 다르게 나눠질 수 있지만 Neural Network를 이루는 layer를 기준으로 다음과 같은 두가지 방식이 존재합니다.

Inter Layer Model Parallelism은 layer를 기준으로 모델을 나누어 학습하는 방법입니다. 위 그림으로 본다면, layer 1과 layer 2는 '0번 GPU'에 layer 3은 '1번 GPU'에 할당된다고 볼 수 있습니다. 반대로 Inter Layer Model Parallelism은 layer와 상관없이 tensor 자체를 쪼개서 GPU에 할당하는 방식입니다.

 

Pipeline Parallelism

Pipeline Parallelism은 Google의 Gpipe에서 소개된 방법으로 Inter Layer Model Parallelism이 가지는 근본적인 문제점을 개선한 병렬화 기법입니다. Neural Network 에서 모델은 항상 이전 layer의 결과를 받아 다음 layer를 연산하기 때문에 유기적입니다. layer를 기준으로 모델을 나누게 되더라도 layer 2의 연산이 끝나지 않으면 layer 3는 연산을 시작할 수 없습니다. 따라서 layer 3의 연산을 책임지는 '1번 GPU'는 layer 2의 `0번 GPU`가 연산 결과값을 전달하기까지 그대로 쉬게 됩니다.

 

이런 Idle Time(GPU가 쉬는 시간. Bubble Time이라고도 합니다.)을 최소화하기 위해서 Gpipe는 

1. Mini-batch를 더 잘게 쪼갠 Micro-batch로 학습을 진행합니다.
2. Mirco-batch 로 연산을 pipelining하여 여러개의 layer를 동시에 처리할 수 있도록 학습합니다.

학습 데이터를 mirco-batch로 한번 더 자르게 되면서 GPU가 결과를 기다리며 쉬는 시간이 적어지게 됩니다. 

Tensor Parallelism

Tensor Parallism은 Intra Layer Model Parallelism 입니다. tensor 를 나눈다는 건, weight를 이루는 metric을 분할한다는 의미와 같습니다. 2D tensor로 column 또는 row를 기준으로 모델을 분할할 수 있습니다. Hugging Face에서 이를 설명하는 좋은 그림이 있어 가져왔습니다.

https://huggingface.co/docs/transformers/parallelism

column을 기준으로 모델(A)을 수직으로 분할하면,

위와 같이 입력 데이터(X)가 동일하게 주입되고(broadcast) 연산 결과를 다시 모으는(all-gather) 형태로 이뤄집니다.

반면 row를 기준으로 모델(A)을 수평으로 분리하면,

데이터부터 분할되면서(scatter) 연산 결과를 합산하는(all-reduce) 형태로 이뤄집니다.

 

조금 더 나아가서 Megatron-LM에서 Transformer Block을 연산할 때 Tensor Parallelism을 활용하고 있는 것까지 이해해 볼까요? 

MLP layer는 'Linear1 → GeLU → Linear2 → Dropout' 순으로 연산이 이뤄집니다. `Linear1 → GeLU` 연산에서는 tensor를 column 기준으로 분할하여 연산하였고, `GeLU → Linear2` 연산에서는 tensor를 row 기준으로 분할하여 연산하고 있습니다. 

 

column 으로 분할 한 후 그 결과값을 row로 이어서 학습하게 되면 row로 나눌 때 데이터를 한번 더 나눠줄 필요가 없습니다. column 기준 연산의 결과값은 Y1, Y2 로 떨어지기 때문에 concat 연산이 필요한데, row 기준 연산의 결과값은 입력 데이터를 나누어 전달받기 때문에 column 기준 연산과 row 기준 연산이 연달아 이뤄질 수 있게 됩니다.

 

마무리..

오늘은 DP, MP 를 기준으로 Parallelism 기법에 대해 알아보았습니다. 사실 Zero Infinity에서 소개된 Zero Redundancy Optimization 까지 한번에 소개하며 글을 마무리하고 싶었으나 이번주는 정말 시간이 없어서 오히려 글을 빠르게 마무리하게 되었습니다😭 조만간 한번 이 글을 수정해보려고 합니다. Zero 기법은 다음편에서 이어가도록 해보겠습니다. 부족한 글을 함께해주셔서 감사합니다.