[Coursera] LLM에 필요한 효율적인 Multi-GPU 활용 전략

앤드류 응 교수님의 강의를 듣고 공부한 글임을 알려드립니다.

Generative AI with LLMs In Generative AI with Large Language Models (LLMs), created in partnership with AWS, you’ll learn the fundamentals of how generative AI works, and how to deploy it in real-world applications.

 

오늘은 효율적으로 Multi-GPU를 사용하는 전략에 대해 공부해 보겠습니다.

저번 공부에 따르면, Large Language Model(LLM)을 학습하려면 Multi-GPU 사용이 꼭 필요합니다.

Single GPU로도 학습할 수 있는 작은 모델이더라도 Multi-GPU를 사용한다면 더 빠르고 효율적인 학습이 가능하기 때문에, 여러 개의 GPU를 어떻게 효율적으로 활용할 수 있느냐를 아는 것은 아주 중요한 문제입니다. 😁

Multi-GPU라고 하면 가장 먼저 떠오르는 개념은 바로 Distributed Data Parallel(DDP)입니다. 먼저 이 DDP에 대해서 간략하게 살펴보고 DDP의 메모리 낭비를 보완한 FSDP에 대해 공부해 보겠습니다.

 

Distributed Data Parallel(DDP)

DDP는 Pytorch에서 데이터를 분산시켜 병렬처리할 수 있도록 도와주는 모듈입니다. LLM에 적용한다면, Large 데이터셋을 분산시키는 역할을 맡게 됩니다.

자세한 기능은 Pytorch의 분산 데이터 병렬 처리하기 에서 확인할 수 있습니다.

간략하게 소개하자면,

1. 각각의 GPU에 모델을 복사합니다.
2. Batch Data를 각 GPU에 Parallel 하게 전달합니다.
3. 각 Dataset은 Parallel 하게 처리되고,
4. 각 GPU의 결과를 결합(Combine)하는 Synchronization이 진행된 후 모델을 Update 합니다.

각 GPU의 모델을 복사해서 업데이트하기 때문에, DDP는 모델 학습에 필요한 weights, gradients, optimizer states 등이 있어야 합니다.

다음 그림처럼 Dataloader의 도움을 받아 각 GPU의 데이터를 분산시켜 처리하는 과정으로 동작합니다.

DDP 과정

만약 DDP를 사용했는데도 모델의 크기가 너무 커서 GPU의 한계가 있다면, DDP와 유사한 다른 방법을 사용할 수 있습니다.

그중 하나가 바로 모델을 sharding 하는 것입니다.

 

Fully sharded Data Parallel(FSDP)

shared가 아니라 sharded입니다. 🤣

왜냐면 저는 처음에 shared(공유된)인 줄 알고 모델을 ‘공유’ 하는 느낌인가? 헷갈렸었고 나중에 ‘공유된’이 아니라 ‘샤드 된’이라는 것을 알았습니다.

조각난, 파편 등을 뜻하는 의미로 모델을 쪼개서 분산 처리를 하는 방법이구나를 알았습니다.

모델 sharding의 대표적인 방법에 파이토치의 Fully Sharded Data Parallel(이하 FSDP)가 있습니다.

FSDP는 ZeRO(Zero Redundancy Optimizer) 논문의 기술에서 착안된 방법으로,

• 메모리를 Optimize 하는 방법
• GPU 여러 개에 모델의 state를 distributing 하거나 sharding
• Zero data overlap 사용

을 통해서 모델이 단일 GPU 칩에서 학습할 수 없을 때, GPU 여러 개로 모델학습을 확장할 수 있도록 합니다.

그럼 먼저 ZeRO에 대해서 살펴보도록 합니다.

 

ZeRO란?

ZeRO 논문

마이크로소프트의 DeepSpeed에 따르면 ZeRO는 이렇습니다.

ZeRO leverages the aggregate computation and memory resources of data parallelism to reduce the memory and compute requirements of each device (GPU) used for model training. ZeRO reduces the memory consumption of each GPU by partitioning the various model training states (weights, gradients, and optimizer states) across the available devices (GPUs and CPUs) in the distributed training hardware. Concretely, ZeRO is being implemented as incremental stages of optimizations, where optimizations in earlier stages are available in the later stages.

일반적으로, 모델을 학습할 때 Optimizer가 다른 파라미터(model parameter, gradients 등) 보다 약 2배 많은 메모리를 요구한다고 합니다. ZeRO는 이점에 주목한 것 같습니다. (아마도...?)

앞서 살펴봤던 기존 DDP의 방법은 모델의 Replication Strategy(복제전략)으로, 각 GPU에 모델 사본을 유지하는 방법을 택하고 있습니다. 👉 Full copy of model and training parameters

이 방법의 단점은 바로, 메모리를 중복적으로 소비하게 된다는 것입니다. 👉 Redundant memory consumption

이와 달리, ZeRO는 단순히 복제하는 방법을 택하지 않음으로써 이러한 중복을 제거했다고 합니다.

모델의 파라미터, 그래디언트, 옵티마이저 states를 gpu에 단순 copy가 아니라 sharding 함으로써 말이죠!

 

ZeRO 논문에 따르면 Baseline과 비교했을 때 3가지 단계로 모델을 sharding 하는 과정을 그림으로 보여주고 있습니다.

1. ZeRO 1단계 : Optimizer states만 파티셔닝 → 메모리를 최대 4배까지 줄일 수 있습니다.
2. ZeRO 2단계 : Gradients 파티셔닝 → 1번의 방법에 2번까지 한다면, 메모리를 8배까지 줄일 수 있습니다.
3. ZeRO 3단계 : Model parameter를 포함한 모든 Componen 파티셔닝 → 1+2+3번의 방법부터는 GPU 수에 따라 메모리가 선형적으로 감소합니다.

 

이렇게 쪼개면 모델을 Synchronization 하는 단계에서 오버헤드가 더 크지 않을까? 하는 생각이 들 수 있습니다.

강의에 따르면, 오버헤드는 DDP의 방식과 비슷하다고 합니다. 이 부분은 논문에서 더 찾아봐야겠습니다. 😁

 

FSDP

다시 FSDP 방법으로 돌아와서 더 구체적으로 동작 과정을 살펴보겠습니다.

DDP처럼 데이터를 분산시키긴 하지만, 이제 모델을 sharding 한다는 점에서 위에서 봤던 DDP 그림에 살짝 변화가 생깁니다.

FSDP 과정

이제 하나의 LLM이 각 GPU에 통째로 복사되지 않습니다.

각각의 batch data를 처리하기 위해서 model의 모든 states를 다 가지고 있는 DDP와 달리, FSDP는 forward, backward 패스 전에 모든 GPU에서 데이터를 수집하는 단계가 존재합니다. 👉 그것이 바로 그림에서 Get Weights에 해당합니다.

• Get Weights 단계
  • 각 GPU는 on-demand 방식으로 GPU의 데이터를 요청하여 unsharded data로 구체화합니다.
  • 작업이 끝나면, unsharded non-local 데이터를 다시 원래의 Sharded data로 GPU에 전달합니다.
• Synchronization 단계
  • backward 패스 후 마지막 단계에서 진행합니다.
  • DDP와 동일한 방식으로 GPU 전체에 걸쳐서 gradient를 Synchronization 합니다.

이렇게 함으로써 GPU 메모리의 중복 소비, 쓸데없는 낭비를 줄일 수 있게 됩니다.

모델을 sharding 하는 정도를 scaling factor로 조절할 수 있다고 합니다. 만약 full sharding을 하게 된다면 메모리는 절약할 수 있지만 GPU 간의 communication volume이 증가할 수 있다는 단점이 있습니다.

이 점을 고려해서 FSDP를 활용할 수 있을 것 같습니다.

 

참고

Generative AI with LLMs - DeepLearning.AI

Model Parallelism

Zero Redundancy Optimizer