pytorch Distributed DataParallel 설명 (multi-gpu 하는 법)

from torch.utils.data.distributed import DistributedSampler

train_dataset = datasets.ImageFolder(traindir, ...)
train_sampler = DistributedSampler(train_dataset)

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.batch_size, shuffle=False,
    num_workers=args.workers, pin_memory=True, sampler=train_sampler)

pytorch에서 모델을 학습시킬 때 multi-gpu를 사용하려면 DataParallel이나 DistributedDataParallel을 사용해야 한다.

[참고](https://tutorials.pytorch.kr/beginner/dist_overview.html)에 나와있듯이, 둘의 차이는

 

DataParallel은 single-machine에서 multi-gpu,

DsitributedDataParallel은 single-machine에서 multi-gpu (전자보다 복잡한 대신 속도 더 빠름) or multi-machine에서 multi-gpu를 사용할 때 쓴다고한다.

machine이란 서버한대를 의미한다. 보통 CPU하나에 4 or 8 gpu가 부착된다.

 

- DataParallel의 경우

net = torch.nn.DataParallel(model,device_ids=[0,1,2])

output = net(input_var) # input_var can be on any device, including CPU

net = torch.nn.DataParallel(model,device_ids=[0,1,2])
output = net(input_var) # input_var can be on any device, including CPU

Code. 1

면 된다. 굉장히 간단하다.

 

- 반면 DistributedDataparallel의 경우 복잡하다. [참고](https://medium.com/daangn/pytorch-multi-gpu-%ED%95%99%EC%8A%B5-%EC%A0%9C%EB%8C%80%EB%A1%9C-%ED%95%98%EA%B8%B0-27270617936b)

 

그 전에 multi-gpu에 대해 알아야 할 개념을 윗 글을 참고하여 정리해보았다.

 

여러 개의 GPU에서 딥러닝을 실행하려면, 모델을 복사해서 각  GPU에 할당해야한다. 그 뒤 batchsize를 batch_size/num_gpu만큼 나눈다. 이것을 scatter 한다고 표현한다. (실제로 scatter 함수가 있다.)

각 GPU에서 모델이 입력을 받아 출력하는 것을 forward 한다고 표현하고, 이 출력들을 하나의 GPU로 모은다. 이렇게 여러 tensor들(출력들)을 하나의 device로 모으는 것을 gather라고 한다.

그리고 back propagation (gradient를 구하는 과정)을 GPU 개별로 하기 때문에 각 GPU에 있는 모델들이 각각의 gradient를 가지고 있다. 모델을 업데이트 시키려면 또 이 gradient들을 하나의 GPU로 모아서 업데이트를 해야 한다.

DataParallel에서는 이 replicate -> scatter -> gather을 Code. 1 로 간단하게 실행할 수 있다.

def data_parallel(module, input, device_ids, output_device):
    replicas = nn.parallel.replicate(module, device_ids)
    inputs = nn.parallel.scatter(input, device_ids)
    replicas = replicas[:len(inputs)]
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    return nn.parallel.gather(outputs, output_device)

Detail of Code. 1  

하지만, gather가 하나의 GPU로 모아주기 때문에 이 GPU는 메모리 사용량이 너무 많아지게 된다.

원글의 Matthew 씨 (이하 원작자)는 Bert 모델을 사용하여 다음과 같이 학습 코드를 짰다.

import torch
import torch.nn as nn

model = BERT(args)
model = torch.nn.DataParallel(model)
model.cuda()

...

for i, (inputs, labels) in enumerate(trainloader):
    outputs = model(inputs)          
    loss = criterion(outputs, labels)     
    
    optimizer.zero_grad()
    loss.backward()                        
    optimizer.step()

Code. 2

 

그 결과는

 

GPU를 한 곳에 몰아주는 nn.DataParellel의 설정 때문에 극심한 GPU memory 불균형이 일어난다.

 

이 해결책으로는 Custom으로 DataParallel을 하면 된다 (근데 custom으로 하려면 복잡하다..).

이것이 근데 패키지로 배포되어서 아주 편하게 쓸 수 있다. 👏 (패키지 다운링크: https://github.com/zhanghang1989/PyTorch-Encoding)

DataParellel에서 하나의 GPU로 출력을 모아준 이유는 loss function을 계산했어야했기 때문인데, 기존방식을 보면

기존 방식

 

이 패키지는 loss function도 병렬로 연산하게 해줌으로써, 이 불균형 문제를 해결하였다.

따라서 각 GPU에서 계산한 loss로 바로 backward 연산을 가능하게 하였다. 

custom DataParallel 사용 후 모습

주의할 점은 이 패키지를 쓸 때, nn.DataParallel(이 기본코드는 한 개의 GPU로 모으도록 기본 설정이 되어 있기 때문)이 아닌 parallel에서 import한 함수를 써야하는 것이다.

import torch
import torch.nn as nn
from parallel import DataParallelModel, DataParallelCriterion

model = BERT(args)
model = DataParallelModel(model)
model.cuda()

criterion = nn.NLLLoss()
criterion = DataParallelCriterion(criterion) 

...

for i, (inputs, labels) in enumerate(trainloader):
    outputs = model(inputs)          
    loss = criterion(outputs, labels)     
    
    optimizer.zero_grad()
    loss.backward()                        
    optimizer.step()

Code. 3

그 결과는?

메모리가 각 GPU에 잘 분산되었지만, GPU-Util을 보면 GPU성능을 효율적으로 잘 사용하지 못하는 것을 알 수 있다.

이 성능을 100%로 끌어 올리려면 Distributed 패키지를 사용해야한다.

 

먼저 용어 가이드 (참조 링크: https://better-tomorrow.tistory.com/entry/Pytorch-Multi-GPU-%EC%A0%95%EB%A6%AC-%EC%A4%91)

 

노드 : GPU가 장착되어 있는 machine을 node라고 함. machine이 두 대이면 node 2

World Size : 사용되는 프로세스들의 개수 = 분산 처리에서 사용되는 총 GPU 개수

Rank : process (GPU)의 ID

• global rank : 전체 node에서의 process id
• local rank : 각 note에서의 process id

 

 

Pytorch Distributed Training tutorial :

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

Writing Distributed Applications with PyTorch — PyTorch Tutorials 1.12.1+cu102 documentation

Distributed Training for ImageNet 예제:

https://github.com/pytorch/examples/blob/main/imagenet/main.py

GitHub - pytorch/examples: A set of examples around pytorch in Vision, Text, Reinforcement Learning, etc.

아래는 Bert에 DDP를 적용한 코드에 대한 설명이다.

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel


def main():
    args = parser.parse_args()

    ngpus_per_node = torch.cuda.device_count()
    args.world_size = ngpus_per_node * args.world_size
    mp.spawn(main_worker, nprocs=ngpus_per_node, 
             args=(ngpus_per_node, args))
    
    
def main_worker(gpu, ngpus_per_node, args):
    global best_acc1
    args.gpu = gpu
    torch.cuda.set_device(args.gpu)
    
    print("Use GPU: {} for training".format(args.gpu))
    args.rank = args.rank * ngpus_per_node + gpu
    dist.init_process_group(backend='nccl', 
                            init_method='tcp://127.0.0.1:FREEPORT',
                            world_size=args.world_size, 
                            rank=args.rank)
    
    model = Bert()
    model.cuda(args.gpu)
    model = DistributedDataParallel(model, device_ids=[args.gpu])

    acc = 0
    for i in range(args.num_epochs):
        model = train(model)
        acc = test(model, acc)
    

Code. 3

def main_worker 부분에서 dist.in_process_group으로 각 GPU에서의 분산학습을 위한 초기화를 해준다.

backend : 'nccl' ('MPI', 'groo' 등이 있는데, 보통 이것을 쓴다.)

init_method :  0 순위 프로세스의 ip 주소와 port로 통신. node가 한 개일 땐, localhost ip인 127.0.0.1 이나 0.0.0.0으로 설정하고 port는 열려있는 것 아무거나 사용하면 됨. (참조 링크: https://csm-kr.tistory.com/47)

mp.spawn 부분에서 main_workers 는 보통 gpu개수 X 4로 해줍니다.

다시 def main_worker 부분을 보면, 

각 GPU가 할당될때마다 gpu id ([0,1,2,3])를 받아서 set_device가 되는 것을 볼 수 있다.

우변의 args.rank는 node의 rank를 의미한다. arg.rank는 global rank를 받는다는 것을 알 수 있다.

 

최종적으로 DistributedDataParallel로 모델을 감싸주면, 알아서 분산 forward와 분산 backward를 해준다.

 

또 Data입력시에도 그냥 DataLoader를 쓰면안되고, DistributedSampler를 이용한 Dataloader를 사용해야 한다.

 

from torch.utils.data.distributed import DistributedSampler

train_dataset = datasets.ImageFolder(traindir, ...)
train_sampler = DistributedSampler(train_dataset)

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.batch_size, shuffle=False,
    num_workers=args.workers, pin_memory=True, sampler=train_sampler)

Code. 4

Dataset은 평소와 같이 설정해주고,

DistributedSampler로 dataset을 감싸준다.

그 다음 DataLoader를 만들 때, sampler에 train_sampler를 입력하면 된다.

DistributedSampler 코드

class DistributedSampler(Sampler):
    def __init__(self, dataset, num_replicas=None, rank=None):
        num_replicas = dist.get_world_size()
        rank = dist.get_rank()
        self.dataset = dataset
        self.num_replicas = num_replicas
        self.rank = rank
        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
        self.total_size = self.num_samples * self.num_replicas
        
    def __iter__(self):
        g = torch.Generator()
        g.manual_seed(self.epoch)
        indices = torch.randperm(len(self.dataset), generator=g).tolist()
        indices = indices[self.rank:self.total_size:self.num_replicas]
        return iter(indices)

Code. 4

Bert 실험 결과는?

메모리가 골고루 분산되고, GPU-Util도 100%인 것을 확인할 수 있다.

 

대학원 연구실에 있을 당시엔 이런것을 할 필요가 없었지만, 회사에 들어가고 나니 대규모 데이터, 모델을 다뤄야 해서 공부해보았다.

비록, 이 글은 짜깁기(특히 당근개발자분 글 거의 복붙..)이지만 더 숙련되면 더 자세히 쓰도록 하겠다.