GPUTraining & View

[yongggg's] pytorch Multi GPU 사용하기

Yonggg 2022. 4. 27. 14:07

안녕하세요 오늘은 제가 일하면서 공부를 했던, pytorch 모델을 학습할 때, Multi GPU를 사용하는 방법을 소개해드리겠습니다.

연구실에서는 GPU 밖에 받지 못해서, 하나만 사용해도 불편함이 없었는데, 회사에서는 4장이 있는 서버를 받아습니다. pretrained model을 사용할 시, 적은 batch를 가져감에도 불구하고 memory issue가 났고 batch 수를 낮추더라도 너무 오래걸린다는 문제가 있었습니다. 이를 해결하기 위해 4장(혹은 N장 중에, N-n장)의 GPU를 모두 사용하는 방법을 소개하겠습니다.

 

먼저, [그림 1]은 huggingface에 공개된 Thomas Wolf의 "Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups" 글에 설명된 그림이다.

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

Multi GPU를 이용하여 딥러닝을 수행할 때, 모델을 각 GPU에 복사(replicate)하여 할당해야한다. 그리고 dataloader의 iteration 마다 batch를 GPU의 개수만큼 나누어(scatter) 각 GPU에 batch/Number of GPU 만큼 입력을 나누어 갖도록한다. 이렇게 나누어진 입력들은 각 GPU에서 forward 과정을 진행하고, 모델이 입력을 받아 출력 값을 뱉어내면, 이 출력 값들을 이용하여 각각의 GPU에서는 Loss를 계산한다. 또한, 이 Loss의 gradients를 backward 하기 위해 하나의 GPU로 모은다(gather). 각 GPU에서는 출력에 대한 N개의  gradient를 갖고 있을 것이며, 이를 1개의 GPU에 보내주어 backward 과정을 수행한다.

지금까지 Multi GPU를 사용할 때 Forward->Loss->Backward 진행 과정을 설명했으며,  본격적으로 Muti GPU를 사용하는 방법을 설명하겠다.

1. GPU 환경 확인

먼저, 자신의 GPU가 몇 개인지, 또 어떤 GPU를 어느정도 쓰고있는지를 확인하기 위해, 다음의 코드로 GPU 환경을 확인한다.

$nvidia-smi

다음 그림은 위의 코드를 실행하여 필자의 서버를 기준으로 TiTAN V GPU가 4장이 있는 것을 확인할 수 있다.

 

[그림 2] GPU 환경 확인

이 곳에 하나 혹은 여러개의 GPU가 잡힌다면, 다음 병렬처리 소개를 읽으면 된다. 하지만, 명령어가 동작하지 않거나 GPU가 잡히지 않는다면, 버전에 맞는 nvidia graphic driver, cuda, cudnn을 설치하고 다시 위를 확인해야한다.

2. DataParallel

먼저, pytorch framework로 학습 시, 제일 간단하게 병렬처리를 할 수 있는 방법은 pytorch에서 지원하는 torch.nn의 DataParallel 이다.  다음과 같은 torch 코드로 쉽게 multi gpu를 사용할 수 있다.

 

from torch.nn import DataParallel

model = DataParallel(model)
model.to(device, device_ids=[device_num1, device_num2, ...])

 

이렇게 모델을 DataParallel의 인자로 넣어주면, 위의 설명대로 GPU에 replicate->scatter->forward->loss->backward

->gather 순으로 과정이 진행된다. 이 때, gather 과정에서 모델의 각 출력이 하나의 GPU에 모이기 때문에, 특정 GPU의 메모리 사용량이 많아진다. 이는, default로 설정되어있는 GPU에 다른 GPU에서 계산된 gradient 까지 모이기 때문에 다른 GPU보다 메모리 사용량이 많아진다. 따라서 다음 코드를 통해 default GPU가 아닌 다른 GPU에 모델의 출력을 모아 GPU사용량의 차이를 줄일 수 있다. 

 

import os
import torch.nn as nn
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,..."
model = nn.DataParallel(model, output_device=1)

 

3. Custom DataParallel

위처럼 gradient 혹은 모델의 출력을 한 GPU를 모으더라도 GPU 사용량에 불균형은 역시 존재한다. 그 이유는 모델의 출력을 계산하여 loss function을 계산해야하기 때문이다. 다시 말해, loss function을 계산하려면, gradient 및 출력을 모을 수 밖에 없는데, 이 때문에 GPU 간에 불균형이 생기는 것이다.

이를 해결하기 위하여, loss function을 각 GPU에서 계산하도록 만들어 GPU 메모리 불균형 문제를 어느정도 해결할 수 있도록 하는 코드가 있다.

zhanghan의 PyTorch-Encoding 코드를 이용하여 각 GPU에서 loss를 계산할 수 있도록 할 수 있다.

 

https://github.com/zhanghang1989/PyTorch-Encoding/blob/master/encoding/parallel.py

 

이 코드를 이용하여, DataParallelCriterion(): 부분의 코드를 다음과 같이 수정하고

 

    def forward(self, inputs, *targets, **kwargs):
        # input should be already scatterd
        # scattering the targets instead
        if not self.device_ids:
            return self.module(inputs, *targets, **kwargs)
        targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            return self.module(inputs, *targets[0], **kwargs[0])
        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        targets = tuple(targets_per_gpu[0] for targets_per_gpu in targets)
        outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs)
        return Reduce.apply(*outputs) / len(outputs)

 

_criterion_parallel_apply(): 부분의 코드를 다음과 같이 수정한 뒤, 

 

    def _worker(i, module, input, target, kwargs, device=None):
        if torch_ver != "0.3":
            torch.set_grad_enabled(grad_enabled)
        if device is None:
            device = get_a_var(input).get_device()
        try:
                with torch.cuda.device(device):
                        output = module(input, target)
                with lock:
                        results[i] = output
        except Exception as e:
            with lock:
                results[i] = e

 

다음 코드와 같이 모델과 loss를 DataParallelModel, DataParallelCriterion으로 감싸주면, 지정한 GPU의 개수를 모두 사용하여 병렬적으로 학습이 진행된다.

 

model = models.UttClsModel(config=config, freeze_transformer=args.freeze_transformer, num_labels=3)
if device.type != 'cpu':
    model = DataParallelModel(model)
model.to(device)

optim = AdamW(model.parameters(), lr=5e-5)
criterion = nn.CrossEntropyLoss()
criterion = DataParallelCriterion(criterion)
criterion.to(device)

 

지금까지 torch framework에서 Multi GPU를 사용해서 학습할 수 있는 skill을 소개했습니다. 다음에는 기회가 된다면, Nvidia의 Apex 모듈을 활용하여 Multi GPU 사용 방법을 소개해보겠습니다. 오늘 글을 통해 많은 도움이 되시길 바랍니다^^!!