딥러닝 모델의 분산학습이란? (Data parallelism과 Model parallelism)

지난 글(자연어처리 모델 학습을 위한 하드웨어 구성은? - NVIDIA Grace)에서 NLP 모델의 학습시 큰 모델 크기와 대규모 학습 데이터 때문에 여러 GPU에 나누어 연산하는 분산학습이 이뤄진다고 간단히 설명했는데 이번 글에서는 본격적으로 딥러닝 분산학습(Distributed learning)에 대해 알아보겠습니다. 딥러닝 분산학습 방법에는 학습 데이터를 여러 GPU에 나눠 학습하는 Data parallelism과 모델을 여러 GPU에 나누는 Model parallelism이 있습니다.

 

딥러닝 학습이란?

분산학습에 대해 알아보기 전에 딥러닝 학습이란 무엇인지부터 정리하죠. 딥러닝 연산은 크게 학습과 추론으로 나뉘는데 이에 대한 설명은 이전글(딥러닝 학습/추론과 이에 따른 하드웨어 차이)를 참고해주세요. 딥러닝 학습은 학습 데이터를 바탕으로 딥러닝 모델 안의 Weight parameter를 업데이트하는 과정입니다. Weight parameter 업데이트는 학습데이터에 표시된 바람직한 결과와 딥러닝 모델의 Output의 차이를 줄여가는 방향으로 진행됩니다. 이미지 인식을 예로 들어보죠. 영상 인식은 사진 속 대상을 분석하여 해당 대상이 여러 개의 인물/사물에 해당할 확률이 얼마인지를 결과로 나타냅니다. 고양이의 사진을 인식한다면 가장 이상적인 결과는 고양이 범주의 확률이 1이고 나머지 범주는 모두 0이 되어야 하지요. 딥러닝 학습을 수행하면 Forward pass를 수행해서 각 범주의 확률을 구한후, 고양이 범주의 확률은 높이고 다른 범주의 확률은 낮추는 방향으로 Weight parameter를 업데이트 합니다. Weight parameter를 업데이트 하는 과정을 Backward pass라고 부르죠.

 

딥러닝 모델의 학습 과정 (source: NVIDIA)

여기에서 딥러닝 학습의 중요한 두 가지 요소인 학습 데이터와 모델이 등장합니다. 학습 데이터는 위 예시에서 수많은 사진에 해당합니다. 고양이, 개 등 다양한 인물/사물이 들어있는 사진이 어떤 인물/사물인지에 대한 설명인 Tag와 함께 있는 학습 데이터가 있어야 Weight parameter를 업데이트할 수 있죠. 학습의 대상이 되는 Model은 추상적으로 말하자면 신경망의 연결 형태입니다. 하지만 그 중 대부분의 정보가 Weight matrix 안의 Parameter로 되어 있기 때문에 모델의 크기를 Parameter의 수로 많이 부릅니다. 

 

마지막 중요한 개념으로는 Batch가 있는데 한 번에 학습 하는 단위(데이터의 양)를 Batch라고 부릅니다. 여기서 한 가지 딥러닝 학습 관련 이론적 개념이 등장하는데 SGD(stochastic gradient descent)입니다. 용어 자체는 매우 어려워보이지만 개념이 그렇게 어렵지는 않습니다. 쉽게 설명하자면 학습 데이터를 전체를 넣고 Forward pass와 Backward pass를 지나가며 Weight parameter를 업데이트하는 것이 아니라, 전체 학습 데이터를 Batch 단위로 잘라서 Forward pass와 Backward pass를 지나면서 여러 번에 걸쳐 업데이트 하자는 개념이죠. 이렇게 하면 보다 다양한 경로를 탐색할 수 있어서 한 번에 모든 데이터를 학습하는 것 대비 모델의 정확도가 높은 장점이 있습니다. 따라서 현재 딥러닝 모델 학습은 전체 학습 데이터를 Batch size로 나눠 여러 번에 걸쳐 Weight parameter를 업데이트 하는 방식을 채택하고 있죠. 딥러닝 모델 학습에 대한 설명이 길어졌는데 다시 분산 학습으로 돌아가죠.

Batch size (source: https://blog.bsion.me)

 

Data parallelism이란?

Data parallelism은 학습해야 할 데이터가 많은 상황에서 학습 속도를 높이기 위해 나온 분산 학습 방법입니다. 간단히 예를 들어 하나의 GPU가 1개의 Data를 학습해야 하는데 1분이 걸린다고 가정하면 1000개의 학습 데이터 Batch를 학습하려면 1000분이 걸립니다. 하지만 100개의 데이터씩 10개의 GPU에 나눠서 학습한다면 100분이 걸리겠죠. Data parallelism은 연산을 하는 일꾼인 GPU의 수를 늘림을 통해 동일 학습 데이터를 보다 빨리 처리하여 Throughput을 높이고 학습 시간을 단축하기 위한 전략입니다.

Model parallelism과 Data parallelism 비교 (source: https://xiandong79.github.io)

 

Synchronization의 필요성

Data parallelism은 학습 시간을 단축한다는 장점이 있지만 매 Weight parameter를 업데이트할 때마다 여러 GPU가 학습한 결과를 종합한 후 다시 나누는 Synchronization이 필요한 단점이 있습니다. 위 설명 예에서도 봤듯이 원래는 1개 GPU가 할 일을 10개가 나눠서 했기 때문에 10개 GPU는 1개의 GPU 처럼 행동해야 합니다. 하지만 학습 데이터를 나눠서 학습했기 때문에 자기가 받은 학습 데이터에 따라 Weight(Gradient) parameter가 업데이트된 결과가 다르겠죠. 따라서 10개 GPU의 업데이트 결과를 모두 모아서 평균을 낸 다음 다시 나눠주는 과정인 Synchronization이 필요합니다. Synchronization이 끝나면 동일한 모델(parameter)를 바탕으로 다시 다음 학습 데이터를 분배 받아 학습을 시작하죠.

 

GPU간 Interconnect의 중요성

위에서 Synchronization을 단점이라고 언급했는데 이는 Synchonization이 시간이 드는 부가적인 작업이기 때문입니다. 하나의 GPU에서 학습했다면 필요하지 않기 때문에 사용자 입장에서 Synchronization은 최소화해야 하는 Overhead입니다. Synchronization은 자세히 보면 Weight parameter(혹은 Gradient)라는 data를 주고받는 과정입니다. 좀 더 전문적인 용어로는 All-Reduce라고 부르는데, 모든 GPU의 Parameter(Gradient)를 더해서 평균을 낸 다음 동일 Parameter(Gradient)를 모든 GPU에 동일하게 나눠주는 과정이죠. 따라서 GPU간의 데이터 전송 속도가 빠르면 그에 반비례해 전송 시간은 줄어듭니다. 여기서 NVIIDA의 NVLink의 중요성이 등장하죠. NVLink는 Chip간 통신에 보통 사용되는 PCIe에 비해 수십배 속도가 빠른데 이렇게 빠른 속도의 GPU간 Interconnect를 통해 Synchronization 시간과 전체 학습 시간을 줄일 수 있죠.

NVIDIA NVLink (source: NVIDIA)

Model parallelism이란?

Model parallelism은 Model size가 너무 커서 하나의 GPU 메모리에 다 들어가지 않는 경우 여러 GPU에 모델 parameter를 나누어 연산하는 것을 말합니다. 이전글(자연어처리 모델 학습을 위한 하드웨어 구성은? - NVIDIA Grace)에서도 말했듯이 특히 자연어처리(NLP) 모델의 경우 GPU 메모리인 HBM에 비해 사이즈가 너무 커서 여러 GPU에 나누어 저장한 후 연산할 수 밖에 없습니다. Model parallelism을 구현하는 방식에는 Tensor parallelism과 Pipeline parallelism이 있으며 두 방식을 모두 적용한 2D(dimension) model parallelism도 가능합니다.

Tensor parallelism(초록색)과 Pipeline parallelism(파란색) 설명 (source: [1])

Tensor parallelism

Tensor parallelism은 커다란 Weight matrix를 여러 GPU로 나누어 연산을 한 후 그 결과값을 합치는(Concatenate) 방식입니다. Weight matrix를 중심으로 보면 각 결과값의 Row는 Weight matrix의 Row와 Input Column의 곱이죠. 따라서 Weight matrix를 Row를 기준으로 나누어 여러 GPU에서 동일한 Input에 대해 곱하여 결과를 낸 후, 결과값들을 붙여서 합하면 원래의 결과값이 나오게 됩니다. 원래 해야 할 Matix mutiplication을 여러 GPU에서 동시에 하기 때문에 Model parallelism과 마찬가지로 Throughput이 증가해서 학습 시간이 단축되는 효과가 있죠.

 

Pipeline parallelism

Pipiline parallelism은 순차적으로 이어진 Layer 혹은 Layer 안의 Stage를 각각 다른 GPU 혹은 GPU 그룹이 나눠서 하는 방식입니다. 예를 들어서 Layer 1을 GPU 1 이 연산하고 그 결과값을 GPU 2에 주면 GPU 2는 그 값을 Input으로 해서 다시 Layer 2를 연산하죠. 이 방식은 순차적으로 해야할 일을 나눠하는 개념이기 때문에 앞에서 설명한 Tensor parallelism과 달리 Throughput 증가 효과가 없습니다. 하지만 자연어처리(NLP)와 같이 큰 모델을 학습하기 위해 모델을 각 GPU에 나눠 담아야 할 필요가 있을 경우, 여러 Layer로 이뤄진 딥러닝 모델의 특성을 반영해 Pipilie parallelism 방식으로 분산할 수 있죠.

 

Model parallelism의 단점

Model parallelism은 Data parallelism에 비해 Synchornization 빈도가 훨씬 높다는 단점이 있습니다. Data parallelism은 Forward pass와 Backward pass를 거친 후 최종 결과인 Weight(Gradient)만을 공유하고 다시 재분배하면 됩니다. 하지만 Model parallelism은 모델을 나눴기 때문에 각 GPU가 생성한 중간 결과물을 연산이 끝날 때마다 매번 Synchoronization 해야 하는 부담이 있습니다. 예를 들어 50개의 Layer로 이러진 딥러닝 모델 연산을 각 GPU가 1 Layer씩 담당해 연산한다고 하면 Forward, Backward 2회 이동하므로 총 100회의 Synchronization이 필요하죠. Data parallelism은 총 1회였지만요. Nvidia는 알고리즘 최적화를 통해 이러한 Model parallelism의 Synchronization 부담을 최적화하고 이론적 최대 연산량의 50%까지 효율을 높였다고 논문[1]에서 밝히고 있는데 실제 성능을 살펴봐야겠습니다.

 

위와 같은 Model parallelism의 단점과 GPU 수의 부담 때문에 지난 글()에서 말한 바와 같이 CPU 메모리를 적극 활용하는 Grace를 NVIDIA는 발표하기도 했죠. 앞으로 증가하는 딥러닝 모델의 크기와 학습 데이터 양에 대응하기 위해 어떠하 새로운 분산학습 방법이 나올지 기대됩니다.

 

[1] Narayanan, D. et al., 2021, Efficient large-scale language model training on GPU clusters