소개
딥 뉴럴 네트워크(DNN)를 훈련할 때 중요한 하이퍼파라미터 중 하나는 배치 크기입니다. 일반적으로 배치 크기가 너무 크면 네트워크가 과적합될 수 있지만, 너무 작으면 수렴 속도가 느려집니다.
고해상도 이미지나 메모리를 많이 차지하는 다른 유형의 데이터를 처리할 때, 대부분의 대규모 DNN 모델 훈련이 GPU에서 이루어진다고 가정하면, 사용 가능한 GPU 메모리에 따라 작은 배치 크기를 적용하는 데 문제가 발생할 수 있습니다. 작은 배치는 수렴 속도를 늦추므로, 우리는 다음과 같은 세 가지 주요 방법을 사용하여 유효 배치 크기를 증가시킬 수 있습니다:
- 여러 개의 소형 GPU에서 작은 배치로 모델을 병렬로 실행 — DP 또는 DDP 알고리즘
- 더 큰 GPU 사용 (비용 문제)
- 여러 단계를 통해 그래디언트 누적
이제 1번과 3번 방법을 더 자세히 알아보겠습니다. 만약 큰 GPU를 가지고 있어 필요한 모든 데이터를 수용할 수 있다면, DDP 부분을 읽고 완전한 코드 부분에서 PyTorch에서의 구현 방식을 확인하여 나머지 부분은 건너뛰어도 됩니다.
데이터 병렬성(DP)
먼저 주 GPU를 정의합니다. 그다음 다음 단계를 수행합니다:
- 10개의 데이터 포인트(소배치)와 모델의 복사본을 주 GPU에서 다른 2개의 GPU로 이동
- 각 GPU에서 순방향 전달을 수행하고 출력을 주 GPU로 전달
- 주 GPU에서 총 손실을 계산하고, 각 GPU로 손실을 다시 보여 매개변수의 그래디언트를 계산
- 그래디언트를 주 GPU로 다시 전달 (이것은 모든 훈련 예제의 그래디언트 평균값입니다), 이를 합산하여 전체 배치 30개의 평균 그래디언트를 얻음
- 주 GPU의 매개변수를 업데이트하고 다음 반복을 위해 다른 GPU로 이 업데이트를 전달
이 과정에는 몇 가지 문제점과 비효율성이 있습니다:
- 데이터 - 주 GPU에서 전달되고 다른 GPU 간에 분배됩니다. 또한 총 손실 계산과 매개변수 업데이트가 주 GPU에서 발생하므로 주 GPU의 활용도가 다른 GPU보다 높습니다
- 각 반복마다 다른 GPU에서 모델을 동기화해야 하므로 훈련 속도가 느려집니다
분산 데이터 병렬성(DDP)
분산 데이터 병렬성은 데이터 병렬 알고리즘의 비효율성을 개선하기 위해 도입되었습니다. 이전과 동일한 설정을 계속 사용합니다 — 3개의 GPU를 사용하여 배치당 30개의 데이터 포인트. 차이점은 다음과 같습니다:
- 주 GPU가 없습니다
- 더 이상 주 GPU가 없으므로, 각 GPU에서 데이터를 비중첩 방식으로 직접 디스크/RAM에서 병렬로 로드합니다 — DistributedSampler가 이 작업을 대신해줍니다. 내부적으로는 로컬 랭크(GPU ID)를 사용하여 GPU 간에 데이터를 분배합니다 — 30개의 데이터 포인트가 주어지면, 첫 번째 GPU는 점 [0, 3, 6, ..., 27]을 사용하고, 두 번째 GPU는 [1, 4, 7, ..., 28], 세 번째 GPU는 [2, 5, 8, ..., 29]를 사용합니다
- 순방향 전달, 손실 계산 및 역방향 전달은 각 GPU에서 독립적으로 실행되며, 그래디언트 평균 계산은 비동기적으로 수행된 후 모든 GPU에서 업데이트가 이루어집니다
DDP가 DP에 비해 가지는 장점 때문에 현재 DDP가 우선적으로 사용되므로, DDP 구현만 보여드리겠습니다.
그래디언트 누적
만약 하나의 GPU만 가지고 있더라도 더 큰 배치 크기를 사용하고 싶다면, 다른 옵션은 특정 단계 수에 대한 그래디언트를 누적하는 것입니다. 이는 여러 소배치의 그래디언트를 누적하여 유효 배치 크기를 증가시키는 효과를 가집니다. 위의 예에서는 3번의 반복을 통해 10개의 데이터 포인트 그래디언트를 누적하여, 유효 배치 크기가 30인 DDP 훈련에서 설명한 결과와 동일한 결과를 얻을 수 있습니다.
DDP 구현 코드
아래에서는 1 GPU 코드와 비교했을 때 DDP 구현 시 차이점만 소개하겠습니다. 전체 코드는 아래 일부 섹션에서 찾을 수 있습니다. 먼저 프로세스 그룹을 초기화하여 다른 프로세스 간의 통신을 허용합니다. int(os.environ["LOCAL_RANK"])를 사용하여 주어진 프로세스에서 사용되는 GPU를 검색합니다.
init_process_group(backend="nccl")
device = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(device)
그다음, 다중 GPU 훈련을 지원하도록 모델을 DistributedDataParallel로 감싸야 합니다.
model = NeuralNetwork(args.data_size)
model = model.to(device)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device])
마지막 부분은 DDP 부분에서 언급한 DistributedSampler를 정의하는 것입니다.
sampler = torch.utils.data.DistributedSampler(dataset)
훈련의 나머지 부분은 변경되지 않습니다 — 전체 코드는 이 글의 마지막 부분에 포함하겠습니다.
그래디언트 누적 구현 코드
역방향 전달이 발생할 때, loss.backward()를 호출하면 그래디언트가 각 텐서에 저장됩니다. 실제 업데이트는 optimizer.step()가 호출될 때 발생하며, 이후 optimizer.zero_grad()를 사용하여 텐서에 저장된 그래디언트를 다음 반복의 역방향 전달과 매개변수 업데이트를 위해 0으로 설정합니다.
따라서 그래디언트를 누적하려면, 그래디언트를 0으로 설정하지 않고 필요한 수의 그래디언트를 누적하기 위해 loss.backward()를 호출합니다. 그런 다음 이를 평균하여 누적 그래디언트 반복에서 평균 그래디언트를 얻습니다 (loss = loss/ACC_STEPS). 그다음 optimizer.step()를 호출하고 그래디언트를 0으로 설정하여 다음 그래디언트 누적을 시작합니다.
ACC_STEPS = dist.get_world_size() # == GPU 수
# 데이터 반복
for i, (idxs, row) in enumerate(loader):
loss = model(row)
# 누적 단계에 따라 손실 조정
loss = loss/ACC_STEPS
loss.backward()
# ACC_STEPS 동안 그래디언트 계속 누적
if ((i + 1) % ACC_STEPS == 0):
optimizer.step()
optimizer.zero_grad()
전체 코드 예제
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import argparse
import torch.optim as optim
import numpy as np
import random
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torch.distributed import init_process_group
import torch.distributed as dist
class CustomDataset(Dataset):
def __init__(self, data):
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, index):
sample = self.data[index]
return index, sample
class SimpleNeuralNetwork(nn.Module):
def __init__(self, input_dim):
super().__init__()
self.linear = nn.Linear(input_dim, 1, bias=False)
self.linear.weight.data.fill_(1.0)
def forward(self, x):
output = self.linear(x)
loss = output.sum()
return loss
class CustomSampler(Sampler):
def __init__(self, dataset, batch_size, num_devices=2):
self.total_samples = len(dataset)
self.batch_size = batch_size
self.num_devices = num_devices
def __iter__(self):
indices = []
for i in range(0, self.total_samples, self.batch_size * self.num_devices):
indices.append(np.arange(self.total_samples)[i: i + self.batch_size*self.num_devices :self.num_devices])
indices.append(np.arange(self.total_samples)[i+1: (i+1) + self.batch_size*self.num_devices :self.num_devices])
return iter(np.concatenate(indices))
def __len__(self):
return self.total_samples
def train_model(config):
input_dim = config.input_size
if config.distributed:
init_process_group(backend="nccl")
device_id = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(device_id)
else:
device_id = "cuda:0"
# 재현성을 위한 시드 고정
seed = config.seed
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
cudnn.benchmark = True
# 데이터 생성
training_data = torch.rand(input_dim, input_dim)
# 모델 초기화
model = SimpleNeuralNetwork(config.input_size)
model = model.to(device_id)
if config.distributed:
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device_id])
# 옵티마이저 설정
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
dataset = CustomDataset(training_data)
if config.distributed:
sampler = torch.utils.data.DistributedSampler(dataset, shuffle=False)
else:
# `DistributedSampler`와 정확히 동일한 재현성을 위해 `CustomSampler` 정의
sampler = CustomSampler(dataset, config.batch_size)
# 데이터 로더 설정
data_loader = DataLoader(
dataset,
batch_size=config.batch_size,
num_workers=0,
pin_memory=True,
sampler=sampler,
shuffle=False,
collate_fn=None,
)
if not config.distributed:
gradients = []
# 누적 단계 수 (GPU 수와 동일)
# 이 수로 손실을 나누어야 여러 GPU에서 평균화된 것과 동일한 그래디언트를 얻을 수 있음
accumulation_steps = config.accumulation_steps
optimizer.zero_grad()
for epoch in range(config.epochs):
if config.distributed:
data_loader.sampler.set_epoch(epoch)
for i, (indices, batch) in enumerate(data_loader):
if config.distributed:
optimizer.zero_grad()
batch = batch.to(device_id, non_blocking=True)
if config.distributed:
rank = dist.get_rank() == 0
else:
rank = True
loss = model(batch)
if config.distributed:
# `DistributedDataParallel`로 감싼 모델 덕분에 그래디언트를 자동으로 평균화합니다
loss.backward()
else:
# 누적 단계에 따라 손실 조정
loss = loss/accumulation_steps
loss.backward()
if i == 0 and rank:
print(f"에포크 {epoch} {100 * '='}")
if not config.distributed:
if (i + 1) % accumulation_steps == 0: # 누적 단계 수만큼 수행했을 때만 단계 실행
# 전체 에포크에 대한 그래디언트 누적
optimizer.step()
optimizer.zero_grad()
else:
optimizer.step()
if not config.distributed and config.verbose:
print(100 * "=")
print("모델 가중치 : ", model.linear.weight)
print(100 * "=")
elif config.distributed and config.verbose and rank:
print(100 * "=")
print("모델 가중치 : ", model.module.linear.weight)
print(100 * "=")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--distributed', action='store_true',)
parser.add_argument('--seed', default=0, type=int)
parser.add_argument('--epochs', default=2, type=int)
parser.add_argument('--batch_size', default=4, type=int)
parser.add_argument('--input_size', default=16, type=int)
parser.add_argument('--accumulation_steps', default=3, type=int)
parser.add_argument('--verbose', action='store_true',)
config = parser.parse_args()
print(config)
train_model(config)
결론
본 글에서는 DP, DDP 알고리즘과 그래디언트 누적을 간략히 소개하고, 다중 GPU가 없는 경우 유효 배치 크기를 증가시키는 방법을 보여주었습니다. 중요한 점은 최종 결과가 동일하더라도 다중 GPU를 사용한 훈련이 그래디언트 누적보다 훨씬 빠르다는 것입니다. 따라서 훈련 속도가 중요하다면, 다중 GPU를 사용하는 것이 훈련을 가속화하는 유일한 방법입니다.