Pytorch - DataParallel

제가 딥러닝을 처음 시작한 2017년 초반에 대학원에서 사용한 GPU는 GTX1080Ti 였습니다. 메모리가 대충 11 GB 정도 였던 것으로 기억하는데 서버당 GPU 4기씩 설치되어 있었으니 그 당시로서는 학교에서 사용할 수 있는 최선의 인프라를 제공받았었죠 ㅎㅎ. 물론 서버 하나를 제가 독점하지는 않았지만요. 

제가 처음 작업한 딥러닝 프로젝트는 100 GB가 넘는 대용량 음성 데이터를 이용한 multi-gpu training 이었습니다. 그 당시에는 Pytorch가 나오기 직전이었고 (아마 lua로 된 Torch만 있었던 것으로 기억합니다.) 딥러닝 프레임워크로 사용한 tensorflow 에는 gpu 병렬구현 api가 존재하지 않았습니다. 즉, multi-gpu에 필요한 1) gpu 별 모델 복사, 2) gpu 별 데이터 분산 및 파라미터 기울기 계산, 3) gpu 별 계산된 기울기를 통합하여 최종 모델 업데이트, 4) 다시 gpu 별 모델 복사를 직접 구현했어야 했습니다. 아마 2,3 일 정도 밤샘 작업해서 완성했었습니다.

GitHub - yjhong89/MultiGPU_wavenet

---

시간이 흘러 자연스럽게 Pytorch를 사용하게 되었고 여러 딥러닝 프레임워크들이 multi-gpu training을 손쉽게 구현할 수 있는 api들을 제공하기 시작했습니다. 이번 포스트에서는 Pytorch의 가장 기본적인 병렬 연산을 지원하는 torch.nn.DataParallel 함수에 대해 알아보도록 하겠습니다. (여기서 말하는 병렬이란 model을 gpu 별로 쪼개는 model parallel이 아닌 같은 모델을 여러 gpu에 복사하고 데이터를 gpu 별로 쪼개는 data parallel를 의미합니다.)

 

torch.nn.DataParallel

torch.nn.DataParallel 은 정말 간단한 함수입니다. 같은 모델을 지정한 여러 gpu device에 복사하고 입력 배치를 gpu 별로 쪼개 각 gpu 별로 forward / backward 연산을 각각 수행합니다. 아규먼트로는 다음과 같습니다.

• module: nn.Module 클래스를 상속하여 구성한 딥러닝 모듈이 들어갑니다.
• device_id: multi-gpu training에 사용할 gpu device id입니다. 디폴트 값은 None이고 직접 리스트 형식으로 지정할 수 있고 (e.g. device_id=[0,1,2]) 디폴트 값을 사용할 경우 cuda에서 잡히는 모든 device를 자동으로 사용합니다.
• output_device: 각 gpu에서 계산된 기울기를 합칠 때 사용할 gpu device id 입니다. 따로 지정하지 않았을 경우 첫 번째 gpu device가 할당됩니다.
• output_dim: 입력 텐서를 쪼갤 축입니다. 당연히 디폴트는 배치 축인 0이겠죠.

Operation

torch.nn.DataParallel 동작은 1) 모듈을 각 gpu로 복사하는 replicate, 2) 입력 텐서를 output_dim (dim=0) 을 기준으로 할당한 gpu 개수 별로 쪼개는 scatter, 3) 각 gpu 별로 분산된 입력을 이용해 연산하는 parallel_apply, 4) output_device 에서 연산 결과를 합치는 gather 동작으로 구성되며 이를 코드로 표현하면 다음과 같습니다.

def data_parallel(module, input, device_ids, output_device=None):
    if not device_ids:
        return module(input)

    if output_device is None:
        output_device = device_ids[0]

    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)

 

Cautions

먼저 매 forward iteration 마다 backward 연산이 끝난 모델이 각 device로 복사되는 구조이다보니 forward 함수 안의 특정한 업데이트는 반영되지 않습니다. 예를 들어 forward 함수를 거칠 때마다 1씩 증가하는 couter 속성이 있다고 하면 forward 함수가 끝난 후 기존 복사본은 제거되기 때문에 항상 초기값으로 유지되는 겁니다. 하지만 Batch normalization, spectral norm 같은 parameter, buffer 들은 output_device 에서 in-place로 업데이트되며 공유됩니다. 

또한, torch.nn.DataParallel로 모듈을 감쌀 경우에는 DataParallel 모듈이 새로운 멤버 속성을 생성하므로 기존 module의 속성에 그대로 접근할 수가 없습니다. 이럴 경우에는 getattr 함수를 이용하거나 .module 속성을 통해 접근할 수 있습니다.

class MyDataParallel(nn.DataParallel):
    def __getattr__(self, name):
        return getattr(self.module, name)

 

Examples

torch.nn.DataParallel을 사용하려면 먼저 사용할 gpu device id를 정의해야되겠죠.

if torch.cuda.is_available() and len(args.deviceIds) > 0:
    # remove any device which doesn't exists
    args.deviceIds = [int(d) for d in args.deviceIds if 0 <= int(d) < torch.cuda.device_count()] 
    # set args.deviceIds[0] (the master node) as the current device
    torch.cuda.set_device(args.deviceIds[0])
    args.device = torch.device("cuda")
else:
    args.device = torch.device('cpu')
    
# Retrieve model
model = loadOrCreateModel(...)  
if torch.cuda.is_available() and len(args.deviceIds) > 1:
    model = torch.nn.DataParallel(model, device_ids=args.deviceIds).to(device=device)

혹은 다음과 같이 구현해도 됩니다.

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0, 1,2,3"

device = ("cuda" if torch.cuda.is_available() else "cpu" )

이후에 간단한 실험을 위한 dataloader 및 module을 정의합니다.

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Parameters and DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100

class Model(nn.Module):
    # Our model

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())

        return output
        
for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())

현재 배치 사이즈는 30입니다. 이를 2 gpu를 이용하면 하나의 gpu 당 15개의 입력이 들어가게 되겠죠. 3 gpu를 이용하면 하나의 gpu 당 10개의 입력이 들어가게 될겁니다.

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])
        In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

---

만약 module의 forward 함수가 배치 단위로 이루어진 텐서가 아니라 scalar를 리턴한다면 어떨까요? 이런 경우에는 할당된 device 개수만큼의 길이를 가진 벡터를 리턴하게 됩니다.

 

참조

• https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html
• https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html
• https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html