Pytorch에서 Multi GPU를 사용하는 방법(nn.DataParallel)

해당 게시글은 자연어처리 리뷰 모임인 집현전 팀의  고현웅님의 [Large-scal LM에 대한 얕고 넓은 지식들(part1)] 부분의 Data Parallelism 강의 영상을 기반으로 작성되었습니다.

Intro.

딥러닝은 기본적으로 매트릭스 연산 기반이기 때문에 모델 학습 시 GPU 연산이 필수적입니다. 하지만 딥러닝 모델의 사이즈가 너무 크고 데이터 양도 많다면, GPU 하나로는 연산이 힘든 경우가 생깁니다. 이 경우 배치 사이즈가 클 경우 자주 나타나는 CUDA OUT OF MEMORY 에러가 뜨게 됩니다. 만약 여러 개의 GPU를 가지고 있다면, 우리는 여러 개의 GPU를 이용해 병렬 처리를 하여 이를 해결할 수 있습니다. Pytorch에서 이러한 기능을 편리하게 사용할 수 있는 모듈을 제공합니다

 

nn.DataParallel()

아래 코드는 한 개의 노드(컴퓨터)에서 네 대의 GPU를 이용하는 경우에 nn.DataParallel을 이용해 데이터 병렬화를 적용한 예시입니다.

import os
import torch.nn as nn

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3' # 네 개 device를 사용
model = nn.DataParallel(model)

코드 실행 시 동작과정은 아래와 같이 Forward Pass와 Backward Pass 두 단계로 진행됩니다.

 

Foward Pass

1) 입력된 미니배치(mini-batch)를 scatter하여 각 디바이스로 전송

2) GPU0에 저장된 모델의 파라미터를 GPU 1,2,3으로 replicate하여 전송

3) 각 device로 복사된 모델로 forward하여 출력값을 구함.

4) 출력값들을 gather하여 GPU0에 모음

Forward Pass(1),(2): Scatter and  Replicate

Forward Pass(3),(4): forwad and gather

Backward Pass

1) GPU0의 gather된 출력값과 label을 이용하여 각각의 loss를 계산

2) 계산된 각각의 loss를 각각의 device에 scatter함

3) 전달받은 loss를 이용해서 각 device에서 backward를 수행

4) 모든 gradient를 GPU0으로 reduce하여 GPU0의 모델 파라미터를 업데이트

Backward Pass(3),(4): compute loss and scatter

Backward Pass(3),(4): backward and reduce

 

*위 그림에서는 모두 0번 GPU로 output을 gather하므로 실험 진행 시 0번 GPU의 사용량이 가장 많게 됩니다. 만약 0번 디바이스 이외에 다른 GPU로 gather를 진행하고 싶은 경우에는 아래의 코드처럼 nn.DataParallel에 두 번째 인자에 output device를 지정할 수 있습니다.

import os
import torch.nn as nn

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3' # 네 개 device를 사용
model = nn.DataParallel(model, output_device=1) # gradient를 1번 디바이스에 gather함

 

nn.DataParallel()의 문제점

1. 멀티쓰레드 모듈이므로 python에서 비효율적임

2. GPU0에서 업데이트된 모델을 매 스텝마다 모든 device로 replicate해야 함

3. 메모리 불균형이 일어나서 GPU를 100%활용할 수 없음