[Pytorch] Multi-GPU 제대로 사용하기

 

Pytorch에서의 Distributed Package 사용하기

  - 규모가 큰 모델을 학습할 때는, 보통 분산 학습을 진행합니다.

     Multi-GPU 학습을 할 때, 분산 학습을 사용할 수 있습니다. 직접 구현할 수도 있지만, Pytorch에서 제공하는 기능을 사용합니다.

 

  - 아래의 코드는 ImageNet1K를 돌리는 경우를 보여줍니다.

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

 

main_worker에서 dist.init_process_group을 통하여, GPU마다 분산 학습을 위해서 초기화를 진행합니다. torch docs에 따르면,

multi-gpu시, backend를 nccl로 설정하라고 되어 있습니다. 

 

DistributedDataParallel은 DataParallel에서 언급한 입력을 분산하고, forward 연산을 수행하고,

다시 backward 연산을 수행하는 역할을 진행합니다.

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):
	args.gpu = gpu
    torch.cuda.set_device(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 = DstributedDataParallel(model, device_ids = [args.gpu])
    
    acc = 0
    for i in range(args.num_epochs):
    	model = train(model)
        acc = test(model, acc)

DataLoader가 입력을 각 프로세스에 전달하기 위해서, DistributedSampler를 사용합니다. DistributedSampler는

DistributedDataParallel과 함께 사용해야 합니다. dataset을 간단하게 DistributedSampler로 감싸기만 하면됩니다.

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)

하지만, 파라미터를 사용하지 않는 경우가 있으면 문제가 발생합니다. Nvidia에서 Apex라는 Mixed Precision 연산을 위한 패키지를 사용하면, 문제를 해결할 수 있습니다.

 

보통 딥러닝은 32비트 연산을 하는데, 16 비트 연산을 사용해서, 메모리를 절약하고, 학습속도를 높

https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py

GitHub - NVIDIA/apex: A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch

import torch.distributed as dist
from apex.parallel import DistributedDataParallel as DDP


def main():
    global args
    
    args.gpu = 0
    args.world_size = 1
    
    args.gpu = args.local_rank
    torch.cuda.set_device(args.gpu)
    torch.distributed.init_process_group(backend='nccl',
                                         init_method='env://')
    args.world_size = torch.distributed.get_world_size()
    
    model = Bert()
    model.cuda(args.gpu)
    model = DDP(model, delay_allreduce=True)

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

위 코드를 실행할 때, 주의할 점으로는 torch.distributed.lauch를 통하여, 프로그램을 실행합니다. main.py를 실행할 때, 노드에서 4개의 프로세스를 돌아가도록 합니다. 각 프로세스는 GPU 하나에서 학습을 진행합니다.

 

만약 GPU가 2개라면, nproc_per_node를 2개로 수정하면 됩니다. main.py에서 batch_size와 num_worker를 설정하는데, 각 GPU마다의 batch_size와 worker 수를 의미합니다.

python -m torch.distributed.launch --nproc_per_node=4 main.py \
    --batch_size 60 \
    --num_workers 2 \
    --gpu_devices 0 1 2 3\
    --distributed \
    --log_freq 100

Reference

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

🔥PyTorch Multi-GPU 학습 제대로 하기