다중 GPU를 효율적으로 사용하는 방법: DP부터 FSDP까지
안녕하세요. 테서의 연구개발팀에서 의료 용어 해석을 진행하고 있는 노수철입니다.
딥러닝 모델을 학습할 때 준비해야 할 사항은 여러 가지가 있겠지만, 좋은 하드웨어 자원도 필수 사항입니다. 하드웨어 자원은 항상 부족한 경우가 많기 때문에, 이를 효율적으로 사용하는 것 역시 중요한 사항입니다. 이번 포스트에서는 보다 하드웨어 자원을 효율적으로 사용하여, 딥러닝 모델을 학습할 수 있는 방법을 알아보도록 하겠습니다.
배경
딥러닝의 이론적 배경은 꽤 오래전에 등장했지만, 인기가 많아지게 된 원인 중 하나로 하드웨어의 발전을 빼놓을 수 없습니다. 하드웨어의 발전이 이루어지면서 점점 더 깊고 복잡한 모델이 등장하게 되었고, 2020년에는 무려 1,750억 개의 파라미터를 가진 OpenAI의 GPT-3가 등장하기도 하였습니다.
GPT-3와 같은 대형 언어 모델(Large Language Model, LLM)을 학습하기 위해서는, 필요한 GPU 자원은 어마무시합니다. 학습 파라미터만 하더라도 FP32로 가정하면 약 700GB의 메모리가 필요하기 때문에, GPU 메모리가 엄청나게 크지 않는 이상, 단일 GPU를 이용한 모델 학습은 불가능한 상황입니다. 이를 해결하기 위해 보다 효율적으로 여러 개의 GPU를 활용하여 딥러닝 모델을 학습시키는 방법이 등장하기 시작하였습니다. 이번 포스트에서는 PyTorch를 이용하여 여러 개의 GPU를 효율적으로 활용하는 방법에 대해 하나씩 살펴보도록 하겠습니다.
이 포스트는 torch==1.12.1을 기준으로 작성되었습니다.
Data Parallel과 Distributed Data Parallel
딥러닝 모델을 학습할 때 배치 사이즈는 정확도와 관련이 있으며, 학습 속도와도 밀접한 관련이 있습니다. 보통 배치 사이즈가 클 수록 한 epoch당 학습 속도도 빨라지기 때문에, 배치 사이즈를 최대한 키우는 것이 GPU를 효율적으로 활용하는 방법과 직결됩니다. 따라서 정해진 모델에서 최대한 배치 사이즈를 키워보는 것이 이번 포스팅의 목표라 할 수 있습니다.
먼저 기본적으로 비교를 위해 단일 GPU에서의 모델 학습을 살펴보도록 하겠습니다. 번역 모델을 활용하기 위해 다음과 같이 PyTorch에서의 더미 데이터셋을 만들어 보았습니다.
import torch
from torch.utils.data import Dataset
class DummyDataset(Dataset):
def __init__(self, num_samples=64000, num_tokens=256, max_len=256):
super().__init__()
self.num_samples = num_samples
self.input_ids = torch.randint(0, num_tokens, size=(num_samples, max_len))
self.decoder_input_ids = torch.randint(0, num_tokens, size=(num_samples, max_len))
self.labels = torch.randint(0, num_tokens, size=(num_samples, max_len))
def __len__(self):
return self.num_samples
def __getitem__(self, idx):
return {
"input_ids": self.input_ids[idx],
"decoder_input_ids": self.decoder_input_ids[idx],
"labels": self.labels[idx],
}더미 데이터셋을 이용하여 간단한 학습 코드를 만들면 다음과 같습니다. 사용하는 모델은 t5-large로, 파라미터 수는 약 7억 개 입니다.
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import T5ForConditionalGeneration
def train(device, batch_size, epochs=100):
model = T5ForConditionalGeneration.from_pretrained("t5-large")
model.to(device)
dataset = DummyDataset()
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(epochs):
for data in tqdm(dataloader):
data = {k: data[k].to(device) for k in data}
optimizer.zero_grad()
output = model(**data)
torch.mean(output.loss).backward()
optimizer.step()이를 단일 GPU인 V100에서 돌리면 최대 배치 사이즈는 14 정도로 확인됩니다.
그렇다면 여러 대의 GPU를 활용하여 하나의 단일 모델을 어떻게 학습할 수 있을까요? 먼저 간단하게, 각 GPU에 서로 다른 데이터를 전송하는 방법을 떠올릴 수 있습니다. 이러한 방식을 Data Parallel (DP)라 하며, PyTorch에서는 DataParallel 클래스로 사용 가능합니다.
PyTorch를 이용한 학습 코드는 아래와 같이 수정이 가능합니다. 먼저 모델을 DataParallel 클래스로 감싸는 과정이 필요하고, forward의 결과도 기존과는 조금 다른 형태이기 때문에 이를 수정해주어야 합니다. loss 값을 각 GPU에서 가져오기 때문에 결과의 크기가 GPU의 개수만큼 나오게 되므로 이를 압축해주면 됩니다.
import torch.nn as nn
def train(device, batch_size, epochs=100):
model = T5ForConditionalGeneration.from_pretrained("t5-large")
model = nn.DataParallel(model) # DataParallel로 모델을 감싸야 함
model.to(device)
dataset = DummyDataset()
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(epochs):
for data in tqdm(dataloader):
data = {k: data[k].to(device) for k in data}
optimizer.zero_grad()
output = model(**data)
# output이 scalar 형태가 아니므로, 압축이 필요함
torch.mean(output.loss).backward()
optimizer.step()4대의 V100 기준으로 DP를 활용하면 배치 사이즈를 최대 36까지 늘릴 수 있었습니다. 단일 GPU에서 배치 사이즈를 14까지 늘릴 수 있었던 모델인데, 왜 4대의 GPU에서는 14×4=56보다는 훨씬 못미치는 36까지밖에 늘릴 수 없을까요? 이는 DP의 작동 방식을 살펴보면 이해가 가능합니다.
아래에 DP의 작동 방식을 순서대로 나열하여 설명해보겠습니다.
- 입력 데이터와 모델의 파라미터를 각 GPU에 전달해줘야 된다.
- backward 과정에서는 각 GPU에 전달된 데이터와 관련된 gradient를 나누어 주어야 한다.
- 마지막으로 모든 gradient를 모으고 업데이트를 해주어야 한다.
DP의 작동 방식을 보면, 첫 번째 GPU가 하는 일이 많아보입니다. 이와 같은 방식은 GPU의 개수가 많아지면 많아질수록 첫 번째 GPU에 부담되는 일이 커질 수 밖에 없는 구조입니다.
또한 PyTorch에서는 multithreading을 이용하여 DataParallel 클래스를 제공하고 있습니다. Python에서의 GIL 특성 상, multithreading에 대한 단점 역시 만만치 않아 보입니다. 이러한 점들을 해결하기 위해, PyTorch에서는 multiprocessing을 이용한 Distributed Data Parallel (DDP)을 제공하고 있습니다.
multiprocessing을 사용하기 때문에, Python에서의 multithreading의 단점이 사라지게 됩니다. 또한, 하나의 GPU를 위해 하나의 프로세스를 생성하므로 DP에서 첫 번째 GPU에게 모든 일을 맡기게 했던 문제도 사라지게 됩니다. 다만, 각 GPU에서 계산한 결과를 합치는 과정이 필요하기 때문에, GPU끼리 통신하기 위한 백엔드 라이브러리가 필요합니다.
DDP를 활용하기 위해 기존의 코드를 변경하면 다음과 같습니다.
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
def train(rank, world_size, batch_size, epochs=100):
global_rank = rank # multi-nodes인 경우, 수정 필요
dist.init_process_group(
backend="nccl",
init_method="tcp://127.0.0.1:33445",
rank=global_rank,
world_size=world_size,
)
torch.cuda.set_device(rank)
model = T5ForConditionalGeneration.from_pretrained("t5-large").to(rank)
model = DDP(model, device_ids=[rank])
dataset = DummyDataset()
sampler = DistributedSampler(dataset, shuffle=True, drop_last=True)
dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(epochs):
sampler.set_epoch(epoch)
for data in tqdm(dataloader):
data = {k: data[k].to(rank) for k in data}
optimizer.zero_grad()
output = model(**data)
output.loss.backward()
optimizer.step()DP보다 코드가 조금 더 복잡해졌으며, rank와 world size와 같은 용어도 등장하였습니다. 간단하게 설명하면, world size는 전체 GPU의 개수를 나타내며 서버(노드)가 한 대 GPU가 4대라면 world size는 4, 서버가 두 대이고 각 서버의 GPU가 4대라면 world size는 8이 됩니다. rank는 GPU의 번호를 뜻하며, local rank의 경우 해당 서버 내에서의 GPU 번호, global rank의 경우 전체 서버에서의 GPU 번호를 뜻합니다. 추가로 sampler를 정의하고 DataLoader에 전달해야 하며, 학습 시 매 epoch마다 set_epoch 함수를 사용하는 과정이 필요합니다.
위 함수를 실행하기 위해서는, 다음과 같이 PyTorch에서 제공하는 multiprocessing을 활용해야 합니다.
import torch.multiprocessing as mp
num_gpus = 4
world_size = 4
batch_size_per_gpu = 12
mp.spawn(train, nprocs=num_gpus, args=(world_size, batch_size_per_gpu))4대의 V100 기준으로, 각 GPU의 배치 사이즈를 12까지 늘릴 수 있었으며, 이는 전체 배치 사이즈 48을 의미합니다. DP에서 배치 사이즈가 36이었던 것을 감안하면, 조금 더 늘어나게 되었습니다.
Fully Sharded Data Parallel
그렇다면, 여기서 더 GPU를 효율적으로 사용할 수 있을까요? DDP에서는 모델 파라미터, gradient, optimizer에서 사용하는 states 등을 모두 각 GPU에서 보관하고 계산하는 데 사용합니다. 만약 이들을 서로 다른 GPU에서 보관하고, 필요할 때만 다른 GPU에서 해당 states를 불러와서 사용한다면 어떨까요? 통신에 대한 overhead가 늘어나는 대신, GPU가 다룰 수 있는 모델의 크기는 더 커질 수 있지 않을까요? 이와 같은 방식을 Fully Sharded Data Parallel (FSDP)이라 하며, 아래 그림과 같은 동작 방식으로 이루어져 있습니다.
DDP에서는 기본적으로 sampler를 통해 각 GPU에 서로 다른 데이터가 전송되며, 각 데이터를 이용해서 모델 파라미터의 gradients A, B, C, D를 계산합니다. 이후 All Reduce 연산을 통해 gradients A, B, C, D에 대한 평균을 구한 뒤, 모든 GPU에 전달됩니다. 이후 optimizer의 step을 통해 각 GPU에서 모델 파라미터가 업데이트 되고, 똑같은 gradients 값을 사용했기 때문에, 똑같은 모델 정보가 보장됩니다.
반면 FSDP에서는 모델의 모든 정보가 하나의 GPU에 있는 것이 아니라, 여러 GPU에 분산되어(sharded) 있습니다. 따라서 forward 과정에서 모델의 각 layer를 통과할 때마다 다른 GPU에 저장되어 있는 파라미터를 가져와 사용하고 제거합니다 (All Gather 연산). 이후 backward 과정에서 다시 gradients를 계산하기 위해 다른 GPU에 저장되어 있는 파라미터를 가져와서 사용하고 (All Gather 연산), 각 GPU에서 계산된 gradients를 다시 원래 속해 있던 GPU에 전달하기 위해서 Reduce Scatter 연산을 사용합니다. 최종적으로 각 GPU에는 각 GPU가 갖고 있던 모델에 대한 gradients만 남기 때문에, 이후 optimizer의 step 연산을 통해 모델의 파라미터를 업데이트할 수 있습니다. (각 연산에 대한 자세한 설명은 NCCL 문서를 참고해 주세요.)
FSDP는 DeepSpeed, Fairseq, VISSL, PyTorch Lightning와 같이 여러 프래임워크/라이브러리에서 이용할 수 있습니다. 이 포스트에서는 PyTorch에서 어떻게 FSDP를 사용할 수 있는 지 알아보도록 하겠습니다.
FSDP는 모델만 분산하여 각 GPU에 저장할 뿐, 기본적인 구조는 DDP와 매우 유사하기 때문에, 변경점이 그리 많지 않습니다.
import functools
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.t5.modeling_t5 import T5Block
def train(rank, world_size, batch_size, epochs=100):
global_rank = rank
dist.init_process_group(
backend="nccl",
init_method="tcp://127.0.0.1:33445",
rank=global_rank,
world_size=world_size,
)
torch.cuda.set_device(rank)
model = T5ForConditionalGeneration.from_pretrained("t5-large")
wrap_policy = functools.partial(transformer_auto_wrap_policy, transformer_layer_cls={T5Block})
model = FSDP(model, auto_wrap_policy=wrap_policy, device_id=rank)
dataset = DummyDataset()
sampler = DistributedSampler(dataset, shuffle=True, drop_last=True)
dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
for epoch in range(epochs):
sampler.set_epoch(epoch)
for data in tqdm(dataloader):
data = {k: data[k].to(rank) for k in data}
optimizer.zero_grad()
output = model(**data)
output.loss.backward()
optimizer.step()위 코드에서 wrap policy란, FSDP에서 모델을 어떻게 나눌 것인가에 대한 policy로, PyTorch에서는 기본적으로 transformer 모델에 대한 policy를 지원하고 있기 때문에, T5Block을 이용하여 policy 함수를 생성하고 사용하였습니다. 이 코드를 활용하면 최대 64 배치 사이즈로 모델 학습이 가능해집니다.
그렇다면 최대 배치 사이즈 64로 학습이 가능한 FSDP가 DDP보다 epoch당 학습 속도가 더 빠를까요? 실제로는 그렇지 않을 수도 있습니다. DDP와는 달리, FSDP에서는 GPU 사이의 데이터 전송이 더 빈번하게 이루어지기 때문에, 이로 인해 기존의 DDP보다 더 느릴 수도 있습니다. 서버가 여러 대(multi-nodes)라면, 서버 간의 통신 속도까지 고려하여 FSDP가 DDP보다 더욱 느려질 수도 있습니다. (이에 대한 실험이 궁금하다면 다음 글을 참고해 주세요.)
배치 사이즈를 키우는 것 외에 또 어떤 장점이 있을까요? 모델을 분산해서 학습하기 때문에, FSDP에서는 학습할 수 있는 모델 사이즈를 더 키울 수 있습니다. 실제로 지금까지 살펴보았던 샘플 코드를 이용하여 모델 파라미터가 30억 개에 달하는 t5–3b 모델의 경우, 4대의 V100을 이용하면 DP와 DDP에서는 학습이 불가능하지만, FSDP에서는 배치 사이즈 8로 학습이 가능해집니다.
마무리
지금까지 GPU를 효율적으로 사용할 수 있는 방법인 DP, DDP, FSDP에 대해 차례로 살펴보았습니다. DP보다는 DDP가 무조건적으로 좋아보이나, (PyTorch 공식 문서에서도 DP 대신 DDP를 추천하고 있습니다.) DDP와 FSDP 사이에는 장단점이 존재해 보입니다. 학습 속도, 모델 사이즈, 배치 사이즈, GPU간의 통신 속도, 서버 간의 통신 속도 등 다양한 변수를 고려하여 DDP와 FSDP 중 어떤 방법을 사용할 것인지는 각자의 몫에 달려 있습니다.
이외에 GPU를 효율적으로 사용할 수 있는 방법으로는, FP16, BF16, INT8과 같이 FP32에 비해 연산량이 적은 type을 사용하는 방법이 있습니다. 위 코드에서 사용하기 위해서는 어떻게 적용할 수 있을까요? 단순히 type만 바꾼다고 해결될까요? 또, 무조건적으로 bit 수가 낮은 type을 사용한다고 해서 모델 학습 속도에 유리할까요? 이에 대한 답도 한 번 생각해 보시는 게 좋을 것 같습니다.
