[Pytorch] Multi-GPU

파라미터 개수가 많은 경우 하나의 GPU로는 연산을 감당할 수 없다. 오늘은 여러개의 GPU를 사용하여 병렬적으로 학습하는 방법에 대해 알아볼 것이다.

1. Multi-GPU

Multi-GPU에 학습을 분산하는 방법에는ⓛ Model parallel(모델 나누기)와 ② Data parallel(데이터 나누기)가 있다.


■ Model parallel

Model parallel은 model내의 layer들을 여러 GPU에 올리는 방법이다. 아래의 Alexnet이 가장 유명한 Model parallel이다.

C1부터 두개의 GPU를 사용하여 parameter를 주고받으며 병렬적으로 학습하고 있는 모습을 볼 수 있다. Model parallel로 학습을 진행할 경우 아래 그림의 첫번째 상황을 조심해야 한다.

두번째 파이프라인처럼 양쪽의 GPU에서 동시에 진행될 수 있도록 구성해야하며, 첫번째처럼 구성하면 병렬화를 하는 이유가 없어지게 된다.

class ModelParallelResNEet50(ResNet): 
def __init__(self, *args, **kwargs): 
super(ModelParallelResNet50, self).__init__(bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs) 
self.seq1 = nn.Sequential( slef.conv1, self.bn1, self.relu, self.maxpool, self.layer1, self.layer2).to('cuda:0') 
self.seq2 = nn.Sequential( self.layer3, self.layer4, self.avgpool).to('cuda:1') 
self.fc.to('cuda:1') 

def forward(self, x): 
	x = self.seq2(self.seq1(x).to('cuda:1')) 
	return self.fc(x.view(x.size(0), -1))

위의 코드를 보면 self.seq1의 layer들은 cuda0에 할당하고 self.seq2와 fc의 layer들은 coda1에 할당한다. 그 후 forward에서 self.seq1(x)에 .to('cuda:1')을 붙여서 self.seq2를 돌려 GPU두 개를 사용하였다. 물론 이렇게 처리하면 GPU를 동시에 사용하지 않고 한 쪽이 끝나면 다른 한 쪽이 시작하는 형식으로 사용하기 때문에 병렬화를 하는 이유가 없다. 이 코드는 Multi-GPU를 사용하는 방식 정도를 참고하면 될 것 같다.



■ Data parallel

Data parallel은 Data들을 여러 GPU에 올리는 방식이다. 아래에서 그 과정을 살펴보자.

· Forward 과정

1. Scatter mini-batch inputs to GPUs
: GPU-1이 여러개의 Data를 GPU-1, GPU-2, GPU-3, GPU-4에 하나씩 넣어준다.

2. Replicate model on GPUs
: GPU에 동일한 model을 할당한다.

3. Parallel forward passes
: 각 GPU 모두 forwad pass를 실행한다.

4. Gather outputs on GPU-1
: 각 GPU로 계산된 output을 다시 GPU-1로 모으면 GPU-1에서 각각의 Loss를 구한다.


· Backward 과정

1. Compute loss gradient on GPU-1
: GPU-1에서 각 Loss에 대해 가장 상위에서 흐르는 기울기, 즉 Loss gradient를 계산한다.

2. Scatter gradients to GPUs
: GPU-1이 각 Loss gradient를 GPU-1, GPU-2, GPU-3, GPU-4에 넣어준다.

3. Parallel backward passes
: Loss gradient를 사용하여 각 GPU에서 bacwkard pass를 수행함으로써 gradient를 얻는다.

4. Reduce gradient to GPU-1
: gradient를 다시 GPU-1에 모아서 그것들의 평균을 낸다.



위의 과정을 보면 GPU-1이 중간 중간 GPU-1, GPU-2, GPU-3, GPU-4로부터 데이터를 한 번에 넘겨받아서 처리를 한 후 다시 각 GPU에 넘겨주는 것을 확인할 수 있다. 여러개의 GPU에서 진행하던 작업을 하나의 GPU로 모으는 과정에서 병목현상이 발생하곤 한다.


이러한 쏠림현상을 보완한 것이 아래의 Distributed 방식이다.



※ Distributed data parallel

Data parallel과 다르게 Distributed data parallel방식은 각 GPU에서 backward수행까지 한 후에 gradient 평균을 내어 weight를 갱신한다.

1) Data parallel 코드

#Data parallel사용 

parallel_model = torch.nn.DataParallel(model) 

#Multi-GPUs로 Foward pass 
predictions = parallel_model(inputs) 

#loss 계산 
loss = loss_function(predictions, labels) 

#loss 평균 + backward pass 
loss.mean().backward() 

#parameter 갱신 
optimizer.step() 

#갱신된 parameter로 다시 Forward pass 
predictions = parallel_model(inputs)

2) Distributed data parallel 코드

train_sampler = torch.utils.data.distributed.DistributedSampler(train_data) 
shuffle = False 
pin_memory = True 
trainloader = torch.utils.data.DataLoader(train_data, batch_size=20, shuffle=shuffle, pin_memory=pin_memory, num_workers=3, sampler=train_sampler) 

def main(): 
	n_gpus = torch.cuda.device_count() 
    torch.multiprocessing.spawn(main_worker, nprocs=n_gpus, args=(n_gpus, )) 
    
def main_worker(gpu, n_gpus): 
	image_size = 224 
    batch_size = 512 
    num_workers = 8 
    epochs = 20 
    
    #생략 
    
    batch_size = int(batch_size / n_gpus) 
    num_worker = int(num_worker / n_gpus) 
    torch.distributed.init_process_group( backend = 'nccl', init_method = 'tcp://127.0.0.1:2568', #멀티프로세싱 통신 규약 정의 world_size = n_gpus, rank = gpu) 
    model = MODEL torch.cuda.set_device(gpu) model = model.cuda(gpu) 
    
    #Distributed data parallel 정의 
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu]) 
    
    #멀티프로세싱 코드 
    from multiprocessing imnport Pool 
    def f(x): 
    	return x*x 
    
    if __name__ == '__main__': 
    	with Pool(5) as p: 
        	print(p.map(f, [1, 2, 3]))