PyTorch를 사랑하는 당근마켓 머신러닝 엔지니어 Matthew 입니다. PyTorch를 사용해서 Multi-GPU 학습을 하는 과정을 정리했습니다.
포스트는 다음과 같이 진행합니다.
- 딥러닝과 Multi-GPU
- PyTorch Data Parallel 기능 사용하기
- Custom으로 Data Parallel 사용하기
- PyTorch에서 Distributed 패키지 사용하기
- Nvidia Apex를 사용해서 학습하기
- Multi-GPU 학습 방법 비교하기
이 글을 끝까지 다 읽고 나면 다음 nvidia-smi 사진과 같이 4개의 GPU를 full로 멋지게 사용할 수 있습니다. 😀
딥러닝과 Multi-GPU 🥕
딥러닝은 기본적으로 GPU 에서 학습을 합니다. Deep Neural Network는 기본적으로 매트릭스 연산을 하기 때문에 GPU를 사용할 경우 연산처리 속도가 상당히 빨라집니다. 딥러닝이 발전함에 따라 점점 네트워크의 크기도 커졌습니다. 비전 분야에서 Neural Network를 깊게 쌓는 방식이 성공하면서부터 그 이후의 딥러닝 연구에서는 대부분 큰 모델을 사용했습니다. 다음 그림에서보면 ResNet이 152개의 층을 쌓은 것을 알 수 있습니다. 비전 분야에서는 ResNet 이후로 꾸준히 큰 데이터셋과 큰 모델로 성능을 높이는 연구를 해왔습니다. 상대적으로 가벼운 모델을 사용하던 NLP 분야에서도 2018년 BERT를 기점으로 큰 데이터셋에 큰 모델로 성능을 높이는 방향으로 연구가 되고 있습니다.
대부분의 경우 Nvidia의 GPU를 사용하는데 GPU 마다 메모리의 양이 다릅니다. 보통 개인이 집이나 연구실에서 딥러닝으로 학습을 할 경우에 GTX 1080 TI 같은 게임용 GPU를 많이 사용합니다. 이런 GPU의 경우 그래픽 작업용이나 딥러닝 연산용 GPU에 비해 가성비가 뛰어나다는 장점이 있습니다. 대부분의 경우에 GTX 1080 TI, TITAN XP와 같은 GPU 하나로도 딥러닝을 하는 데 큰 문제가 없습니다. 하지만 기업이나 연구실에서 큰 데이터셋에 대해 모델을 학습시킬 경우에 하나의 GPU로는 할 수 있는 것이 한정적입니다. 딥러닝에서 batch size는 성능에 영향을 주는 경우가 많습니다. 하나의 GPU 특히 게임용 GPU는 메모리의 한계가 있습니다. 예를 들어 메모리가 12G인 TITAN XP에서는 BERT base model을 batch size 30 이하로 돌릴 수 있습니다. BERT 논문에서 batch size 256으로 학습시킨 것과는 상당한 차이가 있습니다. 이러한 경우 multi-GPU 학습을 합니다. 말 그대로 여러 개의 GPU 에서 하나의 모델이 학습하는 것입니다.
멀티 GPU를 사용하는 경우 다음 그림과 같이 워크스테이션을 구축합니다. 현재 당근마켓에서는 4 개의 TITAN XP 로 워크스테이션을 구축했습니다. 워크 스테이션에서 multi-GPU 환경을 세팅했더라도 여러 개의 GPU를 제대로 활용하는 것은 생각보다 쉽지 않습니다. 여러 개의 GPU를 사용하는데 각 GPU 마다 memory 사용량이 다른 문제가 발생할 수 있습니다. 또한 하나의 GPU를 사용해서 학습하는 경우보다 그닥 학습이 빨라지지 않는 경우도 많습니다. 이러한 환경에 익숙하지 않다면 multi-GPU를 제대로 활용하는데까지 많은 시간이 소요될 수 있습니다. 딥러닝 장비를 사용해서 학습을 하는 것은 공짜가 아닙니다. 자체적으로 워크 스테이션을 만들어서 사무실이나 연구실에서 학습하는 경우는 부담이 덜하겠지만 클라우드에서 모델을 학습하는 경우에는 부담이 큽니다. Multi-GPU로 학습하도록 코드를 디버깅하는 동안 비용이 계속 발생하기 때문입니다. 따라서 이 글에서는 PyTorch으로 multi-GPU 학습하는 동안 겪은 문제들과 해결 방법에 대해 소개하려 합니다.
PyTorch Data Parallel 기능 사용하기 🥕
PyTorch에서는 기본적으로 multi-gpu 학습을 위한 Data Parallel이라는 기능을 제공합니다. Data Parallel이 작동하는 방식을 보여주는 것이 다음 그림입니다.
딥러닝을 여러 개의 GPU에서 사용하려면 일단 모델을 각 GPU에 복사해서 할당해야 합니다. 그리고 iteration을 할 때마다 batch를 GPU의 개수만큼 나눕니다. 이렇게 나누는 과정을 ‘scatter’ 한다고 하며 실제로 Data Parallel에서 scatter 함수를 사용해서 이 작업을 수행합니다. 이렇게 입력을 나누고 나면 각 GPU에서 forward 과정을 진행합니다. 각 입력에 대해 모델이 출력을 내보내면 이제 이 출력들을 하나의 GPU로 모읍니다. 이렇게 tensor를 하나의 device로 모으는 것은 ‘gather’ 이라고 합니다.
보통 딥러닝에서는 모델의 출력과 정답을 비교하는 loss function이 있습니다. Loss function을 통해 loss를 계산하면 back-propagation을 할 수 있습니다. Back-propagation은 각 GPU에서 수행하며 그 결과로 각 GPU에 있던 모델의 gradient를 구할 수 있습니다. 만약 4개의 GPU를 사용한다면 4개의 GPU에 각각 모델이 있고 각 모델은 계산된 gradient를 가지고 있습니다. 이제 모델을 업데이트하기 위해 각 GPU에 있는 gradient를 또 하나의 GPU로 모아서 업데이트를 합니다. 만약 Adam과 같은 optimizer를 사용하고 있다면 gradient로 바로 모델을 업데이트하지 않고 추가 연산을 합니다. 이러한 Data Parallel 기능은 코드 한 줄로 간단히 사용 가능합니다.
nn.DataParallel로 model을 감싸면 학습을 할 때 다음과 같은 작업을 하는 것입니다. 위에서 언급한 대로 replicate → scatter → parallel_apply → gather 순서대로 진행합니다. Gather가 하나의 gpu로 각 모델의 출력을 모아주기 때문에 하나의 gpu의 메모리 사용량이 많을 수 밖에 없습니다.
일반적으로 DataParallel을 사용한다면 다음과 같이 학습 코드가 돌아갑니다.
PyTorch의 DataParallel 을 테스트하기 위해 사용한 코드는 김준성님의 BERT 코드입니다(링크: https://github.com/codertimo/BERT-pytorch). BERT 논문의 모델 사이즈보다 작은 사이즈로 multi GPU 학습을 테스트 했습니다. 모델에 입력으로 들어가는 sequence의 길이는 163, layer 수는 8층, attention head의 수 8개 그리고 hidden unit의 수는 256으로 사용했습니다. 0번, 1번, 2번, 3번 GPU를 사용해서 multi-gpu 학습을 시작한 다음에 nvidia-smi로 GPU 사용 현황을 체크했습니다. 다음과 같이 0번 GPU가 1, 2, 3번 GPU에 비해 6G 정도 더 많은 메모리를 사용하고 있습니다. 이렇게 하나의 GPU가 상대적으로 많은 메모리를 사용하면 batch size를 많이 키울 수 없습니다. 이 실험을 할 때는 200까지 batch size를 키울 수 있었습니다. 딥러닝에서 batch size는 학습 성능에 영향을 주는 경우가 많기 때문에 메모리 사용 불균형은 꼭 해결해야할 문제입니다. 또한 학습을 더 빨리하고 싶어서 multi-GPU를 쓰는 경우도 많습니다. 학습이 오래 걸릴 경우 batch size 차이로 1주일을 더 학습시켜야 하는 상황이 될 수도 있습니다.
메모리 불균형 문제를 제일 간단히 해결하는 방법은 단순히 출력을 다른 GPU로 모으는 것입니다. 디폴트로 설정되어있는 GPU의 경우 gradient 또한 해당 GPU로 모이기 때문에 다른 GPU에 비해 메모리 사용량이 상당히 많습니다. 따라서 출력을 다른 GPU로 모으면 메모리 사용량의 차이를 줄일 수 있습니다. 다음 코드와 같이 간단하게 출력을 모으고 싶은 GPU 번호를 설정하면 됩니다.
output_device를 설정하고 다시 학습을 시작하면 GPU 사용량이 달라진 것을 알 수 있습니다. 0번 GPU의 메모리 사용량은 줄고 1번 GPU의 메모리 사용량은 늘었습니다. 하지만 여전히 균형하게 사용하지 않는 것을 볼 수 있습니다. 모델 출력의 크기는 batch size에 따라 달라집니다. 이대로 batch size를 늘리면 1번 GPU의 메모리 사용량은 점점 늘어나게 됩니다. 따라서 이 방법은 일시적으로 문제를 해결하는 것 같아 보여도 적절한 해결 방법은 아닙니다. 게다가 GPU-Util을 보면 GPU를 제대로 활용하지 못하는 것을 확인할 수 있습니다.
Custom으로 DataParallel 사용하기 🥕
DataParallel을 그대로 사용하면서 메모리 불균형의 문제를 해결할 수 있는 방법에 대한 힌트는 PyTorch-Encoding이라는 패키지에 있습니다(패키지 링크: https://github.com/zhanghang1989/PyTorch-Encoding). 하나의 GPU의 메모리 사용량이 늘어나는 것은 모델의 출력을 하나의 GPU로 모은 것 때문입니다. 왜 하나의 GPU로 모델의 출력을 모을까요? 왜냐하면 모델의 출력을 사용해서 loss function을 계산해야하기 때문입니다. 모델은 DataParallel을 통해 병렬로 연산할 수 있게 만들었지만 loss function이 그대로이기 때문에 하나의 GPU에서 loss를 계산합니다. 따라서 loss function 또한 병렬로 연산하도록 만든다면 메모리 불균형 문제를 어느정도 해결할 수 있습니다.
PyTorch-Encoding 중에서도 다음 파이썬 코드에 loss function을 parallel하게 만드는 코드가 들어있습니다.
Loss function을 병렬 연산 가능하게 만드는 방법은 모델을 병렬 연산으로 만드는 방법과 동일합니다. PyTorch에서는 loss function 또한 하나의 모듈입니다. 이 모듈을 각 GPU에 replicate 합니다. 그리고 데이터의 정답에 해당하는 tensor를 각 GPU로 scatter 합니다. 그러면 loss를 계산하기 위한 모델의 출력, 정답, loss function 모두 각 GPU에서 연산할 수 있도록 바뀐 상태입니다. 따라서 각 GPU에서 loss 값을 계산할 수 있습니다. 각 GPU에서는 계산한 loss로 바로 backward 연산을 할 수 있습니다.
Loss function을 parallel 하게 만들어서 연산하는 과정을 코드로 보자면 다음과 같습니다. 데이터의 정답에 해당하는 target을 scatter 한 다음에 replicate한 module에서 각각 계산을 합니다. 계산한 output와 Reduce.apply를 통해 각 GPU에서 backward 연산을 하도록 만듭니다.
DataParallelCriterion을 사용할 경우에 일반적인 DataParallel로 모델을 감싸면 안됩니다. DataParallel은 기본적으로 하나의 GPU로 출력을 모으기 때문입니다. 따라서 Custom DataParallel 클래스인 DataParallelModel을 사용합니다. DataParallelModel과 DataParallelCriterion을 사용해서 학습하는 과정은 다음과 같습니다. 사용하는 법은 상당히 간단합니다. Pytorch-Encoding 패키지에서 parallel.py 파일만 가져와서 학습 코드에서 import 하도록 만들면 됩니다.
이렇게 학습을 할 경우에 Nvidia-smi 출력 결과는 다음과 같습니다. batch size 는 200으로 동일합니다. DataParallel 만 사용할 때에 비해 1번 GPU와 2번 GPU의 메모리 사용량의 차이가 상당히 줄었습니다. batch size를 기존에 비해 늘릴 수 있기 때문에 학습 시간도 전체적으로 1/3 정도가 줄었습니다. 하지만 GPU-Util의 수치로 확인할 수 있듯이 GPU 성능을 여전히 제대로 활용 못하고 있습니다. GPU 성능을 100 %로 끌어 올리려면 어떻게 해야할까요?
PyTorch에서 Distributed 패키지 사용하기 🥕
딥러닝을 하시는 분들은 분산 학습에 대해 들어보신 적이 있을 겁니다. DeepMind에서 알파고나 알파스타를 발표할 때 어떤 식으로 학습 했는지 설명하는데 이렇게 규모가 큰 모델을 학습할 때는 보통 분산 학습을 합니다.
분산 학습 자체는 하나의 컴퓨터로 학습하는게 아니라 여러 컴퓨터를 사용해서 학습하는 경우를 위해 개발된 것입니다. 하지만 multi-GPU 학습을 할 때도 분산 학습을 사용할 수 있습니다. 분산 학습을 직접 구현할 수도 있지만 PyTorch에서 제공하는 기능을 사용할 수도 있습니다.
PyTorch에서는 DataParallel과 함께 분산 학습과 관련된 기능을 제공합니다. PyTorch에서 분산 학습을 어떻게 하는지 궁금하다면 다음 PyTorch Tutorial을 보는 것을 추천합니다.
단순히 분산 학습을 사용해서 multi-GPU 학습을 하고 싶다면 PyTorch에서 공식적으로 제공하는 example을 보는 것이 좋습니다. 비전 분야에서 큰 데이터셋 중에 유명한 것이 ImageNet 입니다. 다음 링크가 ImageNet에 딥러닝 모델을 학습시키는 코드 예제입니다. 이 예제에서 여러 머신에서 분산 학습을 하는 방법을 소개하는데 하나의 머신에서 여러 GPU 학습하는 방법도 소개합니다.
ImageNet 예제의 main.py 에서 multi-GPU와 관련된 주요 부분을 다음과 같이 정리해 봤습니다. 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 마다 분산 학습을 위한 초기화를 실행합니다. PyTorch의 docs를 보면 multi-GPU 학습을 할 경우 backend로 nccl을 사용하라고 나와있습니다. init_method에서 FREEPORT에 사용 가능한 port를 적으면 됩니다. 이렇게 분산 학습을 위한 초기화를 하고 나면 분산 학습이 가능합니다. 28번째 줄을 보면 model에는 DataParallel 대신에 DistributedDataParallel을 사용하는 것을 볼 수 있습니다. DataParallel에서 언급한 입력을 분산하고 forward 연산을 수행하고 다시 backward 연산을 수행하는 역할을 합니다.
DataLoader가 입력을 각 프로세스에 전달하기 위해서 다음처럼 DistributedSampler를 사용합니다. DistributedSampler는 DistributedDataParallel과 함께 사용해야 합니다. 사용 방법은 간단하게 정의해놓은 dataset를 DistributedSampler로 감싸주고 DataLoader에서 sampler에 인자로 넣어줍니다. 그 다음엔 평소에 DataLoader를 사용하듯이 똑같이 사용하면 됩니다.
DistributedSampler의 내부를 살짝 보자면 다음 코드와 같습니다(많은 부분을 생략했습니다). 각 Sampler는 전체 데이터를 GPU의 개수로 나눈 부분 데이터에서만 데이터를 샘플링합니다. 부분 데이터를 만들기 위해 전체 데이터셋 인덱스 리스트를 무작위로 섞은 다음에 그 인덱스 리스트를 쪼개서 각 GPU Sampler에 할당합니다. epoch 마다 각 GPU sampler에 할당되는 인덱스 리스트는 다시 무작위로 달라집니다. 그러기 위해서는 train_sampler.set_epoch(epoch) 명령어를 매 epoch 마다 학습 전에 실행해야 합니다.
PyTorch Distributed 패키지를 사용해서 BERT 작은 모델을 학습해봤습니다. Nvidia-smi 를 통해 확인한 GPU 메모리 사용 현황은 다음과 같습니다. GPU 메모리 사용량이 완전 동일한 것을 볼 수 있습니다. 또한 GPU-Util의 수치도 99%로 상당히 높은 것을 볼 수 있습니다. 여기까지 왔다면 multi-GPU 학습을 제대로 할 준비가 됐습니다.
하지만 Distibuted DataParallel의 경우 학습을 시작하려 할 때 간간히 문제가 발생할 수 있습니다. 다음 github issue 글이 여러 문제 중에 하나를 보여줍니다. BERT 코드를 돌릴 때도 에러가 발생했는데 모델에서 학습에 사용하지 않는 parameter가 있을 경우에 Distributed DataParallel이 문제를 일으킬 수 있다는 의견이 있습니다. 이러한 문제를 신경쓰지 않고 학습을 하기 위해서 찾아보다가 Nvidia에서 만든 Apex라는 패키지를 발견했습니다.
Nvidia Apex를 사용해서 학습하기 🥕
Nvidia에서 Apex라는 Mixed Precision 연산을 위한 패키지를 만들었습니다. 보통 딥러닝은 32 비트 연산을 하는데 16 비트 연산을 사용해서 메모리를 절약하고 학습 속도를 높이겠다는 의도로 만든 것입니다. Apex에는 Mixed Precision 연산 기능 말고도 Distributed 관련 기능이 포함합니다. 이 포스트에서는 Mixed Precision에 대한 내용은 다루지 않습니다.
Apex의 Distributed DataParallel 기능을 하는 것이 DDP 입니다. Apex에서 ImageNet 학습을 위해 만든 예제에 관련 내용이 있습니다. Apex 사용법은 Docs에 잘 나와있으니 살펴보시면 됩니다.
다음 코드 2번 줄에서 보듯이 apex에서 DistributedDataParallel을 import 해서 사용합니다. 위 PyTorch 공식 예제에서와는 달리 코드 내에서 멀티 프로세싱을 실행하지 않습니다. 19 줄에서 보듯이 DDP로 model을 감싸줍니다. 그 이외에는 PyTorch DistributedDataParallel과 동일합니다.
이 코드를 실행할 때는 다음 명령어를 사용해서 실행합니다. Torch.distributed.launch를 통해 main.py를 실행하는데 노드에서 4개의 프로세스가 돌아가도록 설정합니다. 각 프로세스는 GPU 하나에서 학습을 진행합니다. 만약 GPU가 2개라면 nproc_per_node를 2로 수정하면 됩니다. main.py에 batch_size와 num_worker를 설정하는데 각 GPU 마다의 batch_size와 worker 수를 의미합니다. batch size가 60이고 worker의 수가 2라면 전체적으로는 batch size가 240이며 worker의 수는 8입니다.
Nvidia Apex를 사용해서 multi-GPU 학습을 했습니다. GPU 사용 현황은 다음과 같습니다. GPU 메모리 사용량이 모든 GPU에서 일정합니다.(3번 GPU는 다른 작업이 할당받고 있기 때문에 잡혀있습니다). GPU-Util을 보면 99% 아니면 100 %인 것을 알 수 있습니다.
Multi-GPU 학습 방법 선택하기 🥕
지금까지 살펴본 PyTorch로 multi-GPU 학습하는 방법은 3가지 입니다.
- DataParallel
- Custom DataParallel
- Distributed DataParallel
- Nvidia Apex
DataParallel은 PyTorch에서 제공하는 가장 기본적인 방법이지만 GPU 메모리 불균형 문제가 생겼습니다. Custom DataParallel의 경우 GPU 메모리 문제를 어느정도 해결해주지만 GPU를 제대로 활용하지 못한다는 문제가 있었습니다. Distributed DataParallel은 원래 분산학습을 위해 만들어진 PyTorch의 기능이지만 multi-GPU 학습에도 사용할 수 있고 메모리 불균형 문제와 GPU를 활용하지 못하는 문제가 없었습니다. 하지만 간간히 문제가 발생하기 때문에 Nvidia에서 만든 Apex를 이용해서 multi-GPU 학습하는 것을 살펴봤습니다.
그렇다면 Apex를 사용하는 것이 항상 좋을까요? 제가 살펴본 이런 문제들이 딥러닝 학습을 할 때 항상 발생하지 않습니다. 만약 이미지 분류를 학습한다면 DataParallel 만으로 충분할 수 있습니다. BERT에서 GPU 메모리 불균형 문제가 생기는 이유는 모델 출력이 상당히 크기 때문입니다. 각 step마다 word의 개수만큼이 출력으로 나오기 때문에 이런 문제가 생깁니다. 하지만 이미지 분류의 경우 모델 자체가 클 수는 있어도 모델 출력은 그렇게 크지 않습니다. 따라서 GPU 메모리 불균형은 거의 없습니다.
이를 확인하기 위해 CIFAR-10에 PyramidNet을 학습해봤습니다. 학습에 사용한 코드 링크는 다음과 같습니다. CIFAR-10은 10개의 카테고리를 가진 32x32의 이미지 사이즈를 가지는 데이터셋 입니다. 또한 PyramidNet은 CIFAR-10 에서 최근까지 가장 높은 성능을 냈던 모델입니다. PyramidNet은 모델의 크기를 조절할 수 있습니다. Multi-GPU에서 학습 성능을 비교하려면 사이즈가 큰 모델을 쓰는 것이 좋습니다. 따라서 파라메터의 개수가 24,253,410인 모델을 실험에 사용했습니다. 다음 표에서 PyramidNet(alpha=270)에 해당하는 모델입니다. 학습에는 K80 4개를 사용했습니다.
PyramidNet Single GPU (batch size: 240)
우선 Single GPU로 PyramidNet을 학습시켰습니다. 하나의 GPU만 사용하면 batch size가 240 정도가 가장 크게 사용할 수 있는 한계입니다.
Batch size 240으로 학습을 할 경우에 하나의 batch를 처리하는데 6~7초 정도가 걸립니다. 총 학습시간(1 epoch을 train 하는데 걸린 시간)은 22분 정도가 걸렸습니다.
PyramidNet DataParallel (batch size: 768)
PyTorch의 기본적은 DataParallel 모듈을 사용해서 PyramidNet을 학습했습니다. Batch size는 768 정도까지 키울 수 있습니다. 기존에 하나의 GPU만 사용할 때보다 상당히 큰 batch size를 사용할 수 있습니다. 아래 사진을 보면 DataParallel 모듈만 사용해도 모든 GPU의 메모리가 거의 동일한 것을 볼 수 있습니다. 게다가 BERT의 경우 Adam을 사용하지만 PyramidNet에서는 일반 SGD를 사용하기 때문에 더더욱 메모리 불균형 문제가 없습니다.
학습 시간은 원래 22분에서 5분 정도로 확 줄었습니다. 또한 batch size가 768이 됐음에도 batch time이 오히려 6초에서 5초로 빨라진 것을 볼 수 있습니다. 따라서 Image Classification 을 학습할 때는 DataParallel 만 사용해도 충분한 것을 알 수 있습니다. 만약 더 큰 batch size로 학습하거나 더 빠르게 학습하고 싶은 경우(ImageNet 같은 데이터셋의 경우 학습이 훨씬 오래걸립니다) Distributed 학습을 사용할 수 있습니다. 단, 이 경우는 하나의 컴퓨터에서 multi-GPU 학습을 하는 것이 아니라 여러 대의 컴퓨터에서 학습을 하는 경우입니다.
따라서 학습시키는 모델에 따라 또한 optimizer에 따라 multi-GPU 학습하는 방법을 선택해야 합니다. 비전 분야보다는 자연어처리 분야에서 많이 발생하는 문제라고 볼 수 있습니다.
이 글을 마치며
딥러닝은 논문을 읽고 구현하는 것도 힘든 일인데 이렇게 자원을 제대로 활용하기 위해 들어가는 노력도 많습니다. 딥러닝 관련 논문에 대한 리뷰나 논문의 구현에 대한 자료는 많지만 어떻게 해야 제대로 자원을 활용해서 학습할 수 있는지에 대한 자료는 별로 없어서 글을 작성하게 되었습니다. PyTorch를 사용해서 딥러닝을 하시는 분들께 도움이 되었으면 좋겠습니다.
