효율적인 딥러닝 학습을 위한 병렬화 기본 개념 정리 | Data, Tensor, Pipeline Parallelism

주요 개념

병렬화(Parallelism)는 모델이나 데이터를 여러 개의 GPU로 분산시켜 큰 모델도 학습이 가능하도록 하는 것입니다. 이런 방법론의 필요성은 딥러닝 모델이 커짐에 따라 하나의 모델을 GPU에 넣을 수 없는 상황이 오고, 이를 학습시키기 위한 데이터들의 양이 방대해졌기 때문입니다.

딥러닝에서 병렬화의 방식은 크게 두 가지가 있습니다. 하나는 데이터를 병렬화하는 것이고, 또 다른 하나는 모델의 병렬화입니다. 두 방식은 목적이 서로 다른데, 전자는 데이터를 나눔으로써 학습 속도를 높이기 위함이고 후자는 현존하는 GPU에 올릴 수 없는 초거대 모델의 학습을 위함입니다.

먼저, 데이터의 병렬화는 큰 데이터를 여러 GPU에 분할하여 동시에 처리합니다. 동일한 모델로 처리해야 하기 때문에 모든 GPU로 모델을 복제해야 하고, GPU에서 학습 과정을 통해 나온 결과를 복제된 모델에 동일하게 업데이트 해주는 과정이 필요합니다. (만약 그렇지 않다면... 나눈 데이터로 각자 학습한 서로 다른 모델이 나오는 불상사가 나오겠지요) 그렇기 때문에 메모리 사용량이 증가하고, 동일하게 업데이트 해주는 방법을 고민할 필요가 있습니다.

다음으로 모델의 병렬화는 큰 모델을 여러 GPU로 분할하여 대형 모델 처리합니다. 여기서 문제는 '모델을 어떻게 나눌 것인가'인데, 생각해볼 수 있는 것은 가장 직관적인 모델을 구조(층) 단위로 분할하여 GPU에 할당하는 Pipeline Parallelism과 연산 단위로 나누는 Tensor Parallelism이 있습니다. 모델 병렬화는 하나의 모델로 동작을 해야하는 만큼 다른 GPU로부터 결과값을 이어받아야 하기 때문에 GPU 통신이 필요하고 이 과정에서 오버헤드로 인한 계산 부하 가능성을 고려해야 합니다.

데이터 병렬화

기본 방법

데이터 병렬화를 위해서는 크게 '초기화' - 'Forward' - 'Backward'로 딥러닝의 일반적인 학습 방법에 맞춰 과정이 진행되어야 합니다. 여기서 전체 모델을 뿌리고 결과를 취합하는 컨트롤 타워인 Master GPU가 존재합니다. 우선, 초기화 단계에서는 다음과 같은 일들이 수행되어야 합니다.

 

• 데이터셋을 미니배치 단위로 나눔 (GPU 개수에 맞춰서)
• Master GPU가 각 GPU에 모델 전달 (동일한 가중치)

그 다음 순전파 과정 입니다. Forward 과정에서는 각 GPU 할당 받은 데이터에 대한 각 연산을 병렬적으로 수행합니다.

 

• 각 GPU는 전달받은 데이터의 logit 계산
• Master GPU가 모든 logit 취합
• Master GPU가 전체 logit에 대한 loss 계산

마지막으로 역전파 입니다. Backward에서는 각 GPU에서 계산된 gradient를 하나의 최종 gradient로 합치고, 마스터 노드 병합 및 모델 파라미터 업데이트 수행합니다.

 

• Master GPU는 계산한 loss를 각 GPU 전송
• 각 GPU에서 그래디언트 계산
• Master GPU가 모든 그래디언트 취합
• Master GPU에서 모든 가중치 업데이트 ➡️ 초기화 과정

개선 방법 : DDP

하지만, 위의 방법은 Master GPU가 할 일이 많아서 병목이 필연적으로 발생합니다. 따라서 이를 해결하기 위해 등장한 것이 분산된 데이터 병렬화 Distributed Data Parallelism (DDP) 입니다. DDP도 기본적인 데이터 병렬화와 동일하게 초기화 - Forward - Backward 과정을 거칩니다. 초기화의 경우 위와 동일하게 모든 모델을 GPU에 복제하고, 데이터셋을 미니배치 단위로 나누어 할당합니다.

다음으로 Forward입니다. 기존 과정과 다른 점은 loss를 계산하는 것이 하나의 Master GPU가 아니라 각 GPU마다 계산된다는 점입니다.

 

• 데이터가 모든 GPU로 분산
• 모든 GPU에는 동일한 모델 복제
• GPU마다 데이터와 복제된 모델을 통해 독립적으로 logit 계산
• GPU마다 계산된 logit으로 loss 계산

마지막으로 Backward 입니다. 여러 GPU에 흩어져있는 데이터를 주고받기 위한 방법으로 AllReduce 연산을 활용합니다.

 

• GPU마다 그래디언트 계산
• GPU마다 그래디언트를 AllReduce 연산으로 평균 그래디언트를 구함 (모든 GPU는 같은 평균 그래디언트 동기화)
• GPU마다 평균 그래디언트로 업데이트 진행

텐서 병렬화

텐서 병렬화(Tensor Parellelism)은 행렬 또는 텐서 연산이 행/열을 따라 나누더라도 동일하게 연산이 유지된다는 점에서 고안된 방법입니다. 텐서 연산을 여러 차원의 슬라이스로 나눠 GPU에 할당해서 처리합니다.

 

Tensor Parallelsim [2]

 

텐서를 쪼개는 방향에 따라 세로(Column-wise)와 가로(Row-wise)로 구분되며, 각각 처리하는 방식은 조금 다릅니다. 먼저 Column-wise는 Input 행렬이 주어지면 각 GPU는 행렬의 열 부분을 나눠서 GPU의 output 값들의 열 부분만 계산합니다. 순전파 과정에서 GPU의 output을 이어붙여 하나의 output 행렬 완성하고, 역전파 과정에서 각 GPU의 열단위 그래디언트를 모든 GPU에 동일하게 공유하는 방식으로 연산이 구성됩니다.

반면, Row-wise는 Input 행렬이 주어지면 각 GPU는 행렬의 행 부분을 나눠서 계산합니다. 순전파 과정에서 GPU의 output을 합산해여 하나의 output 행렬 완성하고, 역전파 과정에서 각 GPU의 행단위 그래디언트를 이어붙여 전체 그래디언트를 형성합니다. 같은 듯 조금 다른데, 이는 행렬의 연산을 생각해보면 떠올릴 수 있습니다.

파이프라인 병렬화

파이프라인 병렬화는 모델을 여러 층을 스테이지 단위로 나눠 GPU에 분배하고 데이터들을 나눠 순차적으로 넣어주는 방법입니다. 여기서 스테이지 마다 연산량의 차이가 있는 경우 대기시간(pipeline bubble)이 발생하게 되는데, 이를 최소화할 수 있도록 메커니즘이 구성됩니다. 일반적인 파이프라인 수행 방법은 아래와 같습니다.

 

1) 배치를 더 작은 단위인 마이크로 배치로 분할합니다.
2) Forward는 첫 스테이지 연산이 완료되면 그 다음으로 전달, 두번째 미니 배치는 다시 첫번째로 주입하면서 순차적으로 진행합니다.
3) Backward는 스테이지에 들어간 역순으로 그래디언트가 모델 업데이트합니다.

 

여기서 마이크로 배치로 쪼개는 규모를 줄일수록 대기시간 감소할 수 있습니다만, 너무 많이 줄일 경우 데이터를 전달하는 시간이 증가할 수 있으므로 최적의 나누는 분할 수를 찾을 필요가 있습니다.

추가적으로 Forward와 Backward를 실행하는 시기에 따라 방법론이 분화될 수 있습니다. Synchronous pipeline은 Forward가 끝나기 전에는 Backward 연산이 시작되지 않도록 하고 최종 역전파가 마무리된 후 모델을 한번에 업데이트 합니다. 반면, Asynchronous pipeline : Forward와 Backward를 번갈아가면서 실행하는 방법으로 약간의 bubble이 생길 수 있지만, 전체적인 forward와 backward의 효율을 높일 수 있습니다.

참고자료

[1] 조현수. "모델 최적화 및 경량화". boostcamp AI Tech.
[2] Huggingface