LLM 파인튜닝을 위한 GPU 분산 학습 정복하기 PART 1

 

---

 

 

최근 대규모 언어 모델(LLM)의 등장으로 모델의 크기와 데이터의 양이 급증하면서, 제한된 하드웨어에 맞추거나 학습 시간을 단축하기 위해 다양한 분산 학습 기법들이 급속히 발전하고 있습니다. 특히 GPU 분산 학습은 대용량 데이터와 큰 모델에 대한 효율적인 학습을 가능케 하며, 여러 GPU를 효율적으로 활용하는 방법을 제공합니다. 최근에는 다양한 병렬 처리 방식을 조합하여 더 빠른 학습과 더 큰 모델을 지원하는 방법들도 등장하고 있습니다. 이 글은 다양한 분산 학습 방법들을 살펴보고, 상황에 따라 적합한 방법을 선택할 수 있도록 가이드라인을 제공하는 것을 목적으로 합니다.

 

데이터 병렬 처리(Data Parallelism)와 모델 병렬 처리(Model Parallelism)는 딥러닝 모델을 여러 GPU 또는 디바이스에 분산하여 학습하는 방법들을 말합니다. 이 두 가지 방법은 모델의 크기와 GPU 메모리 제한 등에 따라 선택됩니다. 대규모 언어 모델(LLM)과 같이 모델 사이즈가 큰 경우에는 모델 병렬 처리를 활용하여 모델을 분할하고 여러 GPU에 나누어 처리하는 방법을 선택할 수 있습니다. 반면, 훈련 데이터가 큰 경우에는 데이터 병렬 처리를 통해 데이터를 분산 처리하여 학습 속도를 높이는 방법이 효과적일 수 있습니다.

 

본 포스팅은 두 개의 챕터로 이루어져 있으며, GPU 분산 학습 정복하기 시리즈의 첫 번째 포스팅은 "Data Parallelism"에 대해 다루고 있습니다.

 

 

---

 

 

1. Data Parallellism

Data Parallelism은 모델을 여러 개의 GPU에 복제하고, 각 GPU에 학습 데이터를 분할하여 학습하는 방법입니다. 학습 데이터의 크기가 크고, 모델의 크기가 단일 GPU에 들어갈 때 사용할 수 있습니다. 아래 그림의 오른쪽은 Data Parallelism을 시각화한 것입니다. Data Parallelism에서는 데이터셋이 'N'개의 부분으로 나뉘며, 'N'은 GPU의 개수를 나타냅니다. N개의 데이터는 N개의 GPU에 Replicate 된 모델에 할당됩니다.

 

출처: https://xiandong79.github.io/Intro-Distributed-Deep-Learning

 

Pytorch의 built-in-feature인 DataParallel, DP을 사용해서 코드 한 줄만으로도 쉽게 Data Parallelism을 구현할 수 있습니다.

 

import torch.nn as nn

parallel_model = nn.DataParallel(model)

 

torch.nn.DataParallel을 좀 더 자세히 살펴보겠습니다. N개의 데이터는 N개의 GPU에 Replicate된 모델에 할당됩니다. 각 GPU는 gradient를 계산하고, 한 개의 GPU에 모두 모아서 평균화하고, 이를 이용하여 모델의 가중치를 업데이트합니다.

 

출처: https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255

 

하지만, DataParallel을 사용할 때 발생할 수 있는 문제 중 하나는 GPU 사용의 불균형입니다. 딥러닝 학습은 학습 데이터를 바탕으로 딥러닝 모델 안의 Weight parameter를 업데이트하는 과정입니다. 딥러닝 학습이 시작되면, Forward pass와 Backward pass를 지나가며 Weight parameter를 업데이트합니다. DataParallel을 사용하게 되면, Forward pass의 4단계(상단 그림 참조)에서 병렬 계산의 결과들이 모두 하나의 GPU에 수집됩니다. 이 GPU가 터지지 않게, 메모리가 집중되는 GPU의 메모리 사용량에 Batchsize를 맞추면 다른 GPU를 제대로 활용할 수 없습니다. 하나의 GPU의 사용량이 다른 GPU들보다 과도하게 사용되는 것이 아닌, GPU들 간의 메모리 부담을 더 고르게 분배해 GPU 사용의 균형을 맞춰야 합니다.

 

 

 

---

 

 

2. Distributed Data Parallelism

출처: https://www.coursera.org/learn/generative-ai-with-llms/lecture/e8hbI/optional-video-efficient-multi-gpu-compute-strategies

DataParallel의 한계를 개선하기 위해, Pytorch DistributedDataParallelism, DDP를 사용할 수 있습니다.  DistributedDataParallel은 여러 대의 컴퓨터 또는 여러 개의 GPU를 사용하여 모델을 분산 학습하는 것을 지원합니다. 멀티노드(multi-node) 및 멀티-GPU(multi-GPU) 환경에서 동작합니다.

DataParallel이 한 줄의 코드만으로 구현이 가능하다는 편리함에도 불구하고 DistributedDataParallel을 사용해야 하는 이유를 pytorch 문서에서 설명하고 있습니다. 

 

"DataParallel is single-process, multi-thread, and only works on a single machine, while DistributedDataParallel is multi-process and works for both single- and multi- machine training. DataParallel is usually slower than DistributedDataParallel even on a single machine due to GIL contention across threads, per-iteration replicated model, and additional overhead introduced by scattering inputs and gathering outputs."

 

"Recall from the prior tutorial that if your model is too large to fit on a single GPU, you must use model parallel to split it across multiple GPUs. DistributedDataParallel works with model parallel; DataParallel does not at this time. When DDP is combined with model parallel, each DDP process would use model parallel, and all processes collectively would use data parallel."

 

 

DataParallel은 단일 프로세스(single-process) 내에서 멀티 스레드(multi-thread)로 동작하며, 하나의 컴퓨터에서만 작동합니다. 반면, DistributedDataParallel은 멀티 프로세스(multi-process)로 동작하며, 단일 또는 여러 대의 컴퓨터에서 학습할 수 있습니다. DataParallel은 한 컴퓨터 안에서 멀티 스레드로 동작하는데, 이로 인해 스레드 간 경합이 발생하고 모델을 반복마다 복제해야 하며, 입력 데이터를 분산시키고 출력 데이터를 수집하는 과정에서 오버헤드가 발생합니다. 이러한 이유로 DistributedDataParallel에 비해 심지어 단일 컴퓨터에서도 더 느릴 수 있습니다.

 

또한 모델이 하나의 GPU에 들어가기에 너무 크다면, 여러 GPU를 사용한 모델 병렬 처리(Model Parallel) 방법을 사용해야 합니다. 이 경우 DistributedDataParallel은 Model Parallel과 함께 사용할 수 있지만, DataParallel은 사용할 수 없습니다.

 

DistributedDataParallel은 각 GPU에 복제된 모델들이 동시에 Forward pass와 Backward pass를 수행하고, Gradient를 평균화하여 모델을 업데이트합니다. DataParallel에서는 각 GPU에서 계산된 Gradient를 하나의 GPU에 모아서(Gather) 업데이트하는 방식이기 때문에 업데이트된 모델을 매번 다른 GPU들로 복제(Broadcast) 해야 합니다. 이러한 복잡한 과정을 거치지 않고 각 GPU에서 계산된 Gradients를 모두 더해서 모든 GPU에 균일하게 뿌려준다면 각 GPU에서 자체적으로 step()을 실행할 수 있습니다. torch.nn.parallel.DistributedDataParallel은 Ring All-reduce라는 연산을 통해 모든 GPU의 파라미터를 동시에 업데이트합니다.  해당 연산을 수행 한뒤, 결과를 모든 디바이스로 broadcast하는 연산인 All-reduce를 활용하게 되면서 특정 GPU로 부하가 쏠리지 않게 되었습니다. 

 

출처: https://nbviewer.org/github/tunib-ai/large-scale-lm-tutorials/blob/main/notebooks/05_data_parallelism.ipynb

 

DistributedDataParallel을 사용하려면 먼저 분산 환경을 설정해야 합니다. 간단한 예제를 통해 DistributedDataParallel 환경을 구현해보겠습니다. Pytorch에 포함된 분산 패키지(torch.distributed) 설정에 대한 자세한 내용은 "https://pytorch.org/tutorials/intermediate/ddp_tutorial.html"에서 확인하실 수 있습니다.

 

Process Group

먼저 "setup"과 "cleanup" 함수를 정의해야 합니다. 이는 모든 계산 프로세스가 통신할 수 있는 프로세스 그룹(Process group)을 열게 됩니다. 이렇게 파이썬 기반 스크립트의 함수로 정의하는 방식을 편리하게 해주는 HuggingFace Accelerate 라이브러리도 존재하는데, 이 부분은 추후 포스팅하도록 하겠습니다.

import os
import torch.distributed as dist

def setup(rank, world_size):
    "Sets up the process group and configuration for PyTorch Distributed Data Parallelism"
    os.environ["MASTER_ADDR"] = 'localhost'
    os.environ["MASTER_PORT"] = "12355"

    # Initialize the process group
    dist.init_process_group(backend="nccl", rank=rank, world_size=world_size)

def cleanup():
    "Cleans up the distributed environment"
    dist.destroy_process_group()

setup 함수에는 MASTER_PORT, MASTER_ADDR 등 필요한 변수가 정의되어야 합니다. MASTER ADDR는 통신할 주소, MASTER_PORT는 통신할 포트입니다. 이외에도 이후에 설명할 RANK, WORLD_SIZE, LOCAL_RANK 등을 정의하기도 합니다.

 

torch.distributed 패키지를 사용하여 init_process_group를 호출하면, 전체 프로세스가 속한 default 프로세스 그룹이 만들어집니다. init_process_group에는 backend, rank, world_size가 필요합니다.

 

backend는 MPI(Message Passing Interface)라고 하는 Process 간 Message Passing에 사용되는 여러 연산(e.g. broadcast, reduce, scatter, gather 등)이 정의되어 있는 표준 인터페이스를 정의합니다. nccl, gloo 등이 있으며, nccl은 NVIDIA에서 개발한 GPU 특화 Message Passing 라이브러리로, NVIDIA GPU에서 사용 시, 다른 도구에 비해 월등히 높은 성능을 보여줍니다. gloo는 Facebook에서 개발된 Message Passing 라이브러리로, 이밖에도 다양한 backend 라이브러리가 존재합니다. backend별 수행 가능한 연산을 확인해 보고 사용하면 됩니다. 

 

rank는 현재 GPU의 랭킹으로, 다른 모든 사용 가능한 GPU들과 비교하여 현재 GPU의 랭킹을 나타냅니다. world_size는 training을 위한 프로세스 수이며, GPU의 개수를 의미합니다.

 

 

DitributedDataParallel 모듈은 모델을 각 GPU로 복사하고, loss.backward()가 호출될 때 역전파가 수행되며, 이러한 모델 복사본들 간의 resulting gradients가 평균화/축소화됩니다. 이렇게 함으로써 각 장치가 optimizer step 이후 동일한 가중치를 가지게 됩니다.

from torch.nn.parallel import DistributedDataParallel as DDP

def train(model, rank, world_size):
    setup(rank, world_size)
    model = model.to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    optimizer = optim.AdamW(ddp_model.parameters(), lr=1e-3)
    # Train for one epoch
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    cleanup()

 

 

스크립트를 실행하기 위해 Pytorch에는 편리한 torchrun 커맨드 라인 모듈이 있습니다. 사용하고자 하는 노드의 개수와 실행할 스크립트를 전달합니다. 아래 커맨드 라인은 하나의 머신에서 두 개의 GPU에서 훈련 스크립트를 실행하며, 이것이 PyTorch를 사용한 분산 학습의 기본 형태입니다.

torchrun --nproc_per_node=2 --nnodes=1 example_script.py

 

 

---

 

참고:

https://xiandong79.github.io/Intro-Distributed-Deep-Learning

https://www.coursera.org/learn/generative-ai-with-llms/lecture/e8hbI/optional-video-efficient-multi-gpu-compute-strategies

https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255