[pytorch] Multi-GPU Training | 다중 GPU 학습 예시| Distributed Data Parallel (DDP) | Data Parallel (DP)

Multi-GPU Training

 

Multi-GPU Training (다중 GPU 학습)은 여러 개의 GPU를 사용하여 딥러닝 모델을 학습하는 방법이다. 딥러닝 모델은 점점 커지고 데이터셋의 크기 또한 방대하기 때문에 다중 GPU를 사용하여 딥러닝 모델을 학습하는 것은 사실상 필수적인 기술이라 볼 수 있다. 

 

Pytorch에서는 multi-gpu 학습을 위한 몇 가지 방법을 제공한다.

 

Data Prarallel (DP)

# DataParallel 모듈로 모델 감싸기
model = nn.DataParallel(model)

 

torch.nn.DataParallel 모듈을 사용하는 방법으로, 굉장히 간단하게 동작하지만 몇 가지 치명적인 단점이 존재하는 방법이다.

 

• 장점
  • 아주 간단하다 (딥러닝 모델을 해당 함수로 감싸기만 하면 동작)
• 단점
  • 메모리 사용량 증가 : 각 GPU에서 모델의 복사본을 만들어 메모리를 사용하기 GPU의 수가 증가할 수록 메모리 사용량이 증가
  • 통신 병목 현상 발생 :  각 GPU에서 연산을 수행하고 연산 결과를 하나의 GPU로 모은 후에 모델을 업데이트하기 때문에 GPU 간에 데이터를 복사하고 통신하는 데 시간이 소요. 또한 하나의 GPU로 연산 결과를 모으기 때문에 GPU 수가 증가할 수록 하나의 GPU의 메모리 사용량이 증가해 효율적인 사용이 불가능.

 

 

Distributed Data Parallel (DDP)

 

torch.nn.parallel.DistributedDataParallel 모듈을 분산 학습 환경에서 여러 GPU들 간의 통신을 처리할 수 있는 기능을 제공하기 때문에 다중 GPU뿐만 다중 머신을 사용해서 딥러닝 모델을 학습시킬 수도 있다. 

 

여러가지 면에서 DDP가 DP보다 우수하고 특히 모델과 데이터셋의 크기가 클수록 DDP를 사용하는 것이 유리하다. 하지만 DDP의 경우 분산 학습을 위한 코드 환경을 세팅하는 것이 조금 복잡하다는 단점이 있다.

 

• DDP는 다중 프로세스 병렬 처리를 사용하기 때문에 모델 복제본 간의 GIL connection 이슈가 없음
• 단일 GPU 학습 코드에 비해 몇 가지 추가/수정 필요
• 학습 코드를 함수화하고 해당 함수를 멀티프로세싱 모듈로 실행하는 방식으로 분산 학습을 진행 가능

 

* 본 포스팅에서는 단일 머신 다중 GPU 환경의 분산 학습에 대해 설명한다. (다중머신(여러 대의 컴퓨터(서버)) X)

 

torch.distributed.init_process_group

torch.distributed.init_process_group(backend='nccl',
                            init_method='tcp://127.0.0.1:23456',
                            world_size=ngpus_per_node,
                            rank=process_id)

분산 학습을 위한 초기화 함수로 각 프로세스마다 호출되어야 하고, 분산 학습을 위해 필요한 모든 설정이 완료된 후에만 다음 단계로 진행할 수 있다. 따라서, 모든 프로세스가 init_process_group 함수를 호출하기 전까지는 실행이 차단된다.

• backend: 사용할 분산 처리 백엔드
  • GPU training : 'NCCL'
  • CPU training : 'Gloo'
• init_method: 초기화 방법으로 'NCCL' 백엔드에 단일 머신 다중 GPU 사용 시 'tcp://localhost:port'로 지정
• world_size: 전체 프로세스 개수 (단일 머신의 경우 GPU 개수)
• rank: 현재 프로세스 id. rank는 0부터 world_size - 1까지의 값을 가짐

 

 

DistributedSampler 

train_sampler = DistributedSampler(dataset=train_set, shuffle=True)
batch_sampler_train = torch.utils.data.BatchSampler(train_sampler, opts.batch_size, drop_last=True)
train_loader = DataLoader(train_set, batch_sampler=batch_sampler_train, num_workers=opts.num_workers)

• DistributedSampler는 분산 데이터 병렬학습(distributed data parallel training)시 각 프로세스가 미니배치를 나누어 학습할 데이터 샘플을 결정하는 역할
• 일반적으로 각 프로세스는 전체 데이터셋을 고루 나누어 학습하지만, 이렇게 나누어 학습하는 경우 다른 프로세스에서 학습하는 데이터 샘플과 중복되는 경우가 발생할 수 있음
• DistributedSampler는 데이터셋의 각 샘플에 대한 인덱스를 분산처리에 맞게 새로운 순서로 만들어주고, 해당 인덱스를 이용하여 프로세스들 간의 중복 없는 데이터 분배 가능

 

 

DistributedDataParallel

Collective Communication

model = DistributedDataParallel(module=model, device_ids=[local_gpu_id])

• DistributedDataParallel은 각각의 GPU에 데이터와 모델이 분배
• 각각의 GPU에서 계산된 그래디언트들이 전체적으로 동기화되며 합쳐지는 방식으로 학습
• 이를 통해 각 GPU에서 계산된 그래디언트가 Master GPU에서 처리되어 가중치를 업데이트(DP보다 훨씬 더 효율적)

 

torch.multiprocessing.spawn

import torch.multiprocessing as mp

def train(rank, world_size):
    # 분산 학습 코드 작성
    pass

if __name__ == '__main__':
    world_size = 4
    mp.spawn(train, args=(world_size,), nprocs=world_size)

PyTorch에서 제공하는 분산 학습을 위한 프로세스 그룹을 생성하는 함수로 분산 학습을 위해 여러 프로세스를 실행하는 데 사용

• fn : 실행할 함수. 함수의 첫 번째 파라미터는 rank로 지정.
• args : 함수에 전달할 인자를 지정. (함수의 첫 번째 파라미터인 rank는 제외)
• nprocs : 실행할 프로세스 개수

 

정리해보면 'torch.distributed.init_process_group'으로 분산 학습 환경을 초기화, 데이터가 프로세스 간 중복되지 않도록 'DistributedSampler'를 사용, 모델은 'DistributedDataParallel'로 wrapping한 학습 코드를 함수로 구성하고(첫 번째 파라미터는 rank), 구성된 함수를 'torch.multiprocessing.spawn'를 사용해서 실행시키면 된다는 뜻이다.

 

DDP 코드 예시

• 모델 : resnet18 모델
• 데이터셋 : cifar10
• 단일 머신 다중 GPU 예시
• torch.multiprocessing.spawn, torch.distributed.init_process_group, DistributedSampler,DistributedDataParallel 가 모두 적용된 예시로 각자 환경에 맞게 수정해서 사용 가능

import argparse

import torch
import torchvision.transforms as transforms
from torchvision.datasets.cifar import CIFAR10

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

from torch.nn.parallel import DistributedDataParallel
from torchvision.models import resnet18

def get_args_parser():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--epoch', type=int, default=3)
    parser.add_argument('--batch_size', type=int, default=256)
    parser.add_argument('--port', type=int, default=2033)
    parser.add_argument('--root', type=str, default='./cifar')
    parser.add_argument('--local_rank', type=int)
    return parser

def init_distributed_training(rank, opts):
    # 1. setting for distributed training
    opts.rank = rank
    opts.gpu = opts.rank % torch.cuda.device_count()
    local_gpu_id = int(opts.gpu_ids[opts.rank])
    torch.cuda.set_device(local_gpu_id)
    
    if opts.rank is not None:
        print("Use GPU: {} for training".format(local_gpu_id))

    # 2. init_process_group
    torch.distributed.init_process_group(backend='nccl',
                            init_method='tcp://127.0.0.1:' + str(opts.port),
                            world_size=opts.ngpus_per_node,
                            rank=opts.rank)

    # if put this function, the all processes block at all.
    torch.distributed.barrier()

    # convert print fn iif rank is zero
    setup_for_distributed(opts.rank == 0)
    print('opts :',opts)


def setup_for_distributed(is_master):
    """
    This function disables printing when not in master process
    """
    import builtins as __builtin__
    builtin_print = __builtin__.print

    def print(*args, **kwargs):
        force = kwargs.pop('force', False)
        if is_master or force:
            builtin_print(*args, **kwargs)

    __builtin__.print = print


def main(rank, opts):
    init_distributed_training(rank, opts)
    local_gpu_id = opts.gpu

    train_set = CIFAR10(root=opts.root,
                        train=True,
                        transform=transforms.ToTensor(),
                        download=True)

    train_sampler = DistributedSampler(dataset=train_set, shuffle=True)
    
    batch_sampler_train = torch.utils.data.BatchSampler(train_sampler, opts.batch_size, drop_last=True)
    train_loader = DataLoader(train_set, batch_sampler=batch_sampler_train, num_workers=opts.num_workers)

    model = resnet18(pretrained=False)
    model = model.cuda(local_gpu_id)
    model = DistributedDataParallel(module=model, device_ids=[local_gpu_id])

    criterion = torch.nn.CrossEntropyLoss().to(local_gpu_id)
    optimizer = torch.optim.SGD(params=model.parameters(),
                                lr=0.01,
                                weight_decay=0.0005,
                                momentum=0.9)

    print(f'[INFO] : 학습 시작')
    for epoch in range(opts.epoch):

        model.train()
        train_sampler.set_epoch(epoch)

        for i, (images, labels) in enumerate(train_loader):
            images = images.to(local_gpu_id)
            labels = labels.to(local_gpu_id)
            outputs = model(images)

            optimizer.zero_grad()
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

        print(f'[INFO] : {epoch} 번째 epoch 완료')

    print(f'[INFO] : Distributed 학습 테스트완료')

if __name__ == '__main__':

    parser = argparse.ArgumentParser('Distributed training test', parents=[get_args_parser()])
    opts = parser.parse_args()
    opts.ngpus_per_node = torch.cuda.device_count()
    opts.gpu_ids = list(range(opts.ngpus_per_node))
    opts.num_workers = opts.ngpus_per_node * 4

    torch.multiprocessing.spawn(main,
             args=(opts,),
             nprocs=opts.ngpus_per_node,
             join=True)