[Pytorch] 분산학습(V2) DP, DDP 개념 정리

💡 **DataParallel(DP)**은 Pytorch 에서 제공하는 가장 기본적인 방법이지만, GPU 메모리 불균형이 발생하는 문제점이 있음. **Custom DataParallel**은 GPU 메모리 불균형 문제를 어느 정도 해소해주지만, GPU util이 낮다는 문제가 여전히 존재. **Distributed DataParallel(DDP)**는 분산 학습을 위해 만들어진 Pytorch 기능이지만, multi-gpu 학습에 사용할 수 있고, GPU 메모리 불균형과 Util 문제를 해결.

DP (=pytorch multi-gpu)

---

Learning을 여러 개 GPU에서 사용하기 위해서는 모델을 각 GPU에 복사해서 할당 한다.

매 iteration 마다 mini-batch의 크기를 GPU 개수만큼 나누는데 이 작업을 ‘scatter’라 하며, 실제로 scatter 함수를 이용해서 작업을 수행한다.

Scatter 수행 후, 각 GPU에서 forward를 진행하고 각 GPU의 출력 값(ouput)을 특정 한 GPU로 모아주며 이 작업을 ‘gather’라 한다.
한 GPU로 각 GPU에 있는 모델의 출력 값을 모으는(gather) 이유는 loss function을 사용하여 loss 를 얻기 위해서 이고, 얻은 loss를 가지고 각 GPU에서 Backpropagation을 통해 gradient를 구한 후, 모델을 업데이트 하기 위해 gradient를 다시 한 GPU에 모으고, 모델을 업데이트를 한다.
<Reprlicate model each GPU → make Output per GPU → Output Gather → Compute Loss → Loss Scatter → Make gradient on each GPU → Gradient Gather → Model update(=parameter update)>

기존 DP default code

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

위의 과정과 같이 특정 한 GPU에 연산 결과를 모아주는 작업을 수행함으로 메모리 불균형이 발생하게 됨. 특정 GPU로 연산 결과 값을 몰아주는 이유는 Loss를 계산하기 위해서 이므로 Loss를 각 GPU에서 계산을 하게 된다면 메모리 불균형이 발생하지 않는다는 말이 된다.

그러면 어떻게 Loss를 각 GPU에서 연산할 수 있을까? Pytorch-Encoding package가 있다.

Loss function을 병렬 연산 가능하게 하는 방법은 Model을 병렬 연산으로 만드는 방법과 동일하다.
Pytorch에서 Loss fucntion도 하나의 모듈이므로, 이 모듈을 각 GPU에 Replicate 한다.
후, 데이터 label 또한 각 GPU로 scatter 하게 되면 각 GPU에서 loss 계산을 할 수 있는 상태가 된다.

Loss function을 parallel하게 만들어서 하는 코드를 보면, target을 scatter 후, replicated module(loss function)이 각각 계산을 하고, 계산한 output과 Reduce.apply 를 사용하여 각 GPU에 Backward 연산을 수행한다.

class DataParallelCriterion(DataParallel):
	def forward():
		targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
		replicas = self.replicate(self.module, self.device_ids[:len(input)])
		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), targets

DataParallelCriterion을 사용할 경우, DataParallel로 모델을 감싸면 안된다.
DataParallel은 디폴트로 하나의 GPU로 출력을 모으기 때문이다. 따라서, Custom DataParallel 클래스인 DataParallelModel을 사용해야 한다.

사용하는 방법은 Pytorch-Encoding package에서 parallel.py 파일만 가져와 import하여 사용하면 된다.

import torch
import torch.nn as nn
from parallel import DataParallelModel, DataParallelCriterion

model = make_model(args)
model = DataParallelModel(model)
model.cuda()

criterion = make_criterion(args)
criterion = DataParallelCriterion(criterion)

for i, (inputs, labels) in enumerate(train_loader):
	outputs = model(inputs)
	loss = criterion(outputs, labels)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()

위와 같은 방법(loss calculate parallel)을 사용하면 메모리 사용량의 차이가 상당히 줄어들어 batch 크기를 기존에 비해 늘릴 수 있어 학습 시간을 줄일 수 있다. 다만, GPU util은 낮게 사용한다.

GPU 성능을 높게 사용하려면? ⇒ DistributedDataParallel 사용

DDP (Distributed Data Parallel)

---

main.py 를 실행하면 main 함수가 실행이 되고, 다시 main 함수가 main_worker 함수들을 multi-processing으로 실행하게 된다.

GPU 4개를 하나의 노드로 보고 world_size를 설정한 후, mp.spawn 함수가 4개의 GPU에서 따로 main_worker를 실행하면 되는데, 이 방법은 main_worker에서 dist.init_process_group을 통해 각 GPU마다 분산 학습을 위한 초기화를 실행.

Multi-GPU 학습을 할 경우, backend로 nccl을 사용해야 하며, init_method에서 freeport에 사용 가능한 port를 적으면 되고, DataParallel 대신 DistributedDataParallel을 사용해야 한다.

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):
    global best_acc1
    args.gpu = gpu
    torch.cuda.set_device(args.gpu)
    
    print("Use GPU: {} for training".format(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 = make_model(args)
    model.cuda(args.gpu)
    model = DistributedDataParallel(model, device_ids=[args.gpu])

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

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

DistibutedSampler는 DistributedDataParallel과 함께 사용해야 한다. Dataset을 DistributedSampler로 감싸주고 DataLoader에 sampler 인자로 넣어주며 되며 DataLoader를 사용하듯이 똑같이 사용하면 된다.

내부의 코드를 보면 각 Sampler는 전체 데이터를 GPU 개수로 나눈 부분 데이터에서만 데이터를 샘플링하며, 부분 데이터를 만들기 위해 전체 데이터 셋 인덱스 리스트를 셔플하고 인덱스 리스트를 쪼개서 각 GPU sampler에 할당한다.

매 학습 epoch 마다 train_sampler.set_epoch(epoch) 명령어를 꼭 호출 해야 한다.

💡 학습 시키는 모델에 따라 또한 optimizer에 따라 multi-GPU **학습하는 방법을 선택**해야 합니다.