Pytorch 9 : Distributed training

이번 글은 여러 개의 gpu를 사용해서 모델을 학습시키는 코드를 살펴보고자 한다.

이 글에서는 하나의 머신에서 여러 개의 GPU를 사용하는 상황을 가정한다.

 

하나의 모델을 학습시킬 때 여러 개의 GPU를 활용하는 것을 우리는 'Distributed training'이라 하며 이는 상당히 자주 쓰이는 학습법 중 하나이다. 특히 mini-batch size가 상당히 큰 경우에는 하나의 gpu에서 연산하고자 하면 학습하는 데 걸리는 시간이 상당히 길어질 뿐만 아니라 'CUDA : Out of Memory'를 만날 수도 있다.

이 경우에는 Mini-batch를 여러 개의 (sub) Mini-batch로 나누어서 각각의 gpu에서 병렬적으로 연산이 이뤄지도록 하는 방법을 사용해야 한다.

이를 위한 코드를 살펴보자.

 

import torch
import torch.nn as nn

import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

기본적으로 매우 큰 사이즈의 Mini-batch를 여러 개의 (sub) Mini-batch로 나누어 여러 개의 GPU에서 독립적으로 연산을 진행한 이후에 나온 값들 (eg : gradient or buffer etc) 을 통합하여 모델 파라미터를 업데이트하는 것이 Distributed training의 가장 일반적인 방법이다. 이때 각 gpu에는 동일한 모델(과 파라미터)이 복제되어져 있고 독립적으로 gradient 연산이 이뤄진 이후에 이 gradient들을 서로 합쳐서 (이를 communication이라 한다.) 복제된 모델들의 파라미터가 업데이트된다. 

이를 우리는 Distributed Data-Parallel training (DDP) 라 부른다.

 

이를 위해서는 크게 3개의 패키지/모듈이 필요한데 torch.distributed, torch.nn.parallel.DistributedDataParallel, torch.multiprocessing이 그것이다.

 

torch.multiprocessing

 

Multi-GPU training 맥락에서 멀티프로세싱은 여러 GPU에서 병렬적으로 계산을 수행하는 작업을 의미한다.

이 패키지는 각 gpu에 프로세싱을 할당하여 독립적으로 작동하도록 해준다. 

그렇기 때문에 각 gpu는 메모리/리소스를 공유하지 않는다. 또한 파이토치의 multiprocessing 패키지는 tensor를 서로 다른 프로세스 간에 공유하여 처리하도록 하는 기능이 있다.

 

torch.distributed

 

torch.distributed는 서로 다른 gpu들 사이의 communication을 수행해준다. 

우리는 단일한 하나의 모델 파라미터를 최적화하는 것이고 (sub) Mini-batch data들을 여러 개의 gpu에 forward -> backward 연산을 처리하고 여기서 나온 gradient 값들을 활용해서 파라미터를 업데이트하려면 각 머신들 사이의 communication, 정보 교환이 이뤄져야 한다. 즉, gradient 값을 동기화해주어야 한다.

이 작업을 해주는 패키지가 torch.distributed이다. 추가적으로 communication을 위해서는 별도의 backend가 필요한데 파이토치에서 자체적으로 지원해주는 backend에는 gloo, nccl, mpi 등이 있다. 

일반적으로 CPU에서는 gloo, cuda는 nccl을 사용한다.

def setup(rank, world_size):
	os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    dist.init_process_group('nccl', rank=rank, world_size=world_size)

def cleanup():
	dist.destroy_process_group()

torch.nn.parallel.DistributedDataParallel

 

DistributedDataParallel은 단일 기기 (single Machine)에서 여러 개의 GPU (Multi-GPU)를 사용하는 작업, 또는 여러 기기 (multi Machine)를 사용하는 작업 등에서 활용된다.

DataParallel이란 모듈도 있지만 이는 추천되지 않는다. 

 

즉, 우리의 모델이 torch.nn.parallel.DistributedDataParallel을 통해서 여러 gpu에 복제되어 학습이 될 수 있도록 하고, torch.multiprocessing은 각각의 GPU에서 독립적으로 프로세스가 처리될 수 있도록 하고, 마지막으로 torch.distributed는 각 gpu 사이의 communication을 위한 작업을 지원해준다.

 

이제 학습 (train)을 위한 코드를 보도록 하자.

Mini-batch size는 2048이고 gpu 4개를 활용하여 학습시킨다고 가정하자.

 

# Assuming you have prepared your dataset, let's call it "dataset"
train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=512,
                                         shuffle=(train_sampler is None),
                                         num_workers=0,
                                         pin_memory=True,
                                         sampler=train_sampler)

def train(gpu, args):
    rank = args.nr * args.gpus + gpu
    setup(rank, args.world_size)

    # create model and move it to GPU with id gpu
    model = SimpleModel()
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    model = DistributedDataParallel(model, device_ids=[gpu])

    # define loss function and optimizer
    criterion = nn.CrossEntropyLoss().cuda(gpu)
    optimizer = torch.optim.SGD(model.parameters(), 1e-4)
    
    # adjust the way we load data
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    train_loader = torch.utils.data.DataLoader(dataset, batch_size=512,
                                         shuffle=(train_sampler is None),
                                         num_workers=0,
                                         pin_memory=True,
                                         sampler=train_sampler)

    num_epochs = 10
    for epoch in range(num_epochs):
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            # ensure we're doing this on the right device
            inputs, targets = inputs.cuda(gpu), targets.cuda(gpu)
            
            # forward pass and loss computation
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            # backward pass
            loss.backward()

            # parameter update
            optimizer.step()

            # zero the gradients ready for the next step
            optimizer.zero_grad()
        
    cleanup()

 

핵심이 되는 부분만 살펴보자.

모델 클래스를 인스턴스화한 후, 이를 gpu로 보낸다. (model.cuda(gpu))

이후 "DistributedDataParallel" 패키지로 이 모델을 감싸준다

 

train dataset도 torch.distributed 패키지를 활용해주어 우리의 데이터셋 (sub mini-batch)이 각각의 gpu에 load될 수 있도록 해준다. 

 

다음으로 forward 연산을 수행해주고 loss.backwrad(), optimizer.step() 부분에서 gradient 연산 및 gradient communication이 이뤄진다.

 

마지막으로 cleanup()은 우리가 프로그램을 종료하기 전에 (위 기준에서는 학습이 최종적으로 종료되기 전) 모든 프로세스가 communication을 종료하기 위해 명시적으로 호출해준다.

 

 

이상으로 Distributed training을 위한 가장 기본적인 패키지/모듈과 코드를 살펴보았다.

사실 이를 실제 작업에서 구현하기 위해서는 더 깊은 내용들이 필요하지만 본 글에서 그것까지 다루진 않겠다.