[네이버 부스트 캠프 AI Tech] Multi-GPU에 간단한 고찰

본 글은 네이버 부스트 캠프 AI Tech 기간동안

개인적으로 배운 내용들을 주단위로 정리한 글입니다

 

본 글의 내용은 새롭게 알게 된 내용을 중심으로 정리하였고

복습 중요도를 선정해서 정리하였습니다

 

이번 부스트캠프 과정주의 저의 가장큰 목표는

'경험하지 않은 어려운 길을 가자'입니다

 

평소 GPU 1대도 구하기 힘들어 파라미터가 큰 모델은 학습을 돌려보지 못했던 저에게

Multi - GPU는 너무나 꿈같은 이야기 처럼 들렸습니다

 

그럼에도 불구하고 

개인 단위가 아니라 회사 단위로 넘어 갔을 때 Multi - GPU에 대한 개념이 잡혀있지 않다면

큰 부족함으로 느껴질 것 같아서 

pytorch에서 Multi-GPU를 학습하는 방법에 대해서 조사해보고 정리하였습니다

 

추가적인 정보나 잘못된 정보가 있다면 댓글 부탁드리겠습니다 :) 

 

 

본 글은

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

블로그의 글을 참고하여 작성되었습니다

자세한 정보나 확실한 정보를 얻고 싶은신 분들은 위의 링크를 참고해주시기를 바랍니다

---

✅ Week 3

목차

1. PyTorch Data Parallel
2. Custom Data Paralle
3. PyTorch Distributed
4. Apex

---

✅ Intro

현재 딥러닝의 한계는 없는 것일까?

사실 최근 딥러닝의 한계는 명확하다

GPT, llama와 같이 초초초? 거대 모델이 요즘 트렌드이다

또한 빅데이터셋과 같이 초초초? 거대 데이터 셋으로 학습한 위의 모델들을 

일반적인 GPU 사양으로 돌릴 수 가 있을까?

 

이러한 문제점을 다룰 수 있는 또는 해결 할  수 있는 방법은 크게 3가지가 있다

---

1. GPU 성능 개선

이 부분은 명확하다

현재 나와있는 GPU의 성능이 기하급수적으로 발전하면 위에서 표현된 딥러닝의 한계는 없어질 것이다

하지만 이부분은 소프트웨어 개발자나 딥러닝 연구자의 입장에서 어떻게 할 수  있는 부분이 아니다

 

2. 모델 경량

이 부분은 최근 딥러닝 연구에서도 활발하게 진행되고 있는 분야있다

MobileNet, 지식증류 기법과 같이 기존의 모델 보다 파라미터 수를 급격하게 줄이면서도

어느 정도의 성능을 보장하는 그러한 방법들이 이 부분에 들어간다고 할 수 있다

 

3. GPU 여러대 활용하기

오늘 다루어볼 부분은 이 부분이다 바로 현재 존재하는 GPU를 여러대 활용해서 문제를 해결하는 것이다

실제로 Pytorch와 CUDA에서 이런 환경을 제공해주고 있다

따라서 이러한 부분의 기술 익히고 연습한다면 문제를 해결하는데 큰 도움이 될 수 있을 것이다

 

---

✅1. PyTorch Data Parallel

딥러닝 모델을 학습하다보면 항상 마주치는 공포의 단어가 있다

OOM - out of memory이다

 

그럴때 마다 항상 드는 생각은 

아.. Batch size나 줄이자... 이다

 

하지만 이러한 방법으로 항상 OOM을 해결 할 수는 없다

 

모델은 갈수록 커지고 있고

데이터 셋도 갈 수록 증가하고 있다

 

그리고 좋은 성능을 얻기 위해서는 큰 모델을 학습하는 것은 좋은 선택지 중 하나이다

 

이러한 OOM을 해결하기 위해서는 

여러대의 GPU를 활용하는 방법을 고려해보아야한다

 

물론 예산적인 문제 때문에 비싼 GPU를 여러대 보유하는 것은 개인에게 힘든일이다

 

하지만 회사단위나 클라우드 시스템을 활용한다면 

Multi GPU를 활용해 볼 수 있을 것이다 

 

그런상황에서 Pytorch는 Data Parallel을 제공해주고 있다

 

https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255

사진을 보면 복잡해 보이지만 

간단하다

 

우선 우리가 활용하는 Model을 사용하는 GPU 만큼 복사한다 

그리고 학습 iteration 마다 batch 사이즈를 다시 더 작은  batch 사이즈로 줄여서 GPU에게 나눠주는 방법이다 - scatter

나눠준 데이터를 통해서 GPU에서 forward가 진행되고 각 GPU가 만든 output을 출력하면 다시 하나의 GPU로 모이게 된다 - gather

 

그리고 loss를 구해서 이제 backward과정을 진행하는데 이경우도 각 GPU서 각각 실행된다

parameters같은 경우는 각 GPU에 존재하기 때문에 Gradient를 구할 수가 있다

 

다소 복잡해 보일 수 있는 과정을 Pytorch에서는 한 개의 코드만 추가해주면 간단하게 실행된다

from torch import nn #근데 왜 import torch.nn as nn으로 할까? 그냥 torch에서 nn을 import하는 코드가 더 깔끔하지 않나?

#아무튼
#우리가 활용하는 모델을 변수 model이라고 하자
model = Resnet18() #가정하면
model = nn.DataParallel(model) #이렇게 한줄 추가해주면 이제 우리 모델은 병렬처리된다

정말간단하다

 

그리고 앞에서 살펴본 과정 scatter, gather을 직접 함수로 호출해서 사용할 수있다

def data_parallel(module, input, device_ids, output_device):
    replicas = nn.parallel.replicate(module, device_ids)
    inputs = nn.parallel.scatter(input, device_ids)
    replicas = replicas[:len(inputs)]
    outputs = nn.parallel.parallel_apply(replicas, inputs)
    return nn.parallel.gather(outputs, output_device)

이렇게 nn에 구현이 되어 있고 자세한 설명은 공식 튜토리얼 링크에서 가져와서 아래에 적어두었다

• 복제(replicate): 여러 기기에 모듈을 복제합니다.
• 분산(scatter): 첫번째 차원에서 입력을 분산합니다.
• 수집(gather): 첫번째 차원에서 입력을 수집하고 합칩니다.
• 병렬적용(parallel_apply): 이미 분산된 입력의 집합을 이미 분산된 모델의 집합에 적용합니다.

멀티-GPU 예제

 

그리고 굳이 전체 모델을 병렬처리하지 않더라도 모델 안에 부분적인 Layer를 병렬 처리할 수도 있다

아래 코드에 예시로 작성해 보았다

import torch
import torch.nn as nn


class DataParallelModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.block1 = nn.Linear(10, 20)

        # wrap block2 in DataParallel
        self.block2 = nn.Linear(20, 20)
        self.block2 = nn.DataParallel(self.block2) ## 이제 block2 레이어는 병렬처리된다

        self.block3 = nn.Linear(20, 20)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        return x

---

그러나 이러한 병렬 학습을 진행 하다보면 'GPU의 불균형이' 발생할 수 도 있다

이게 무슨 말인가하면 

예를 들어서 GPU 4대로 병렬 처리를 하고 있는데

한개의 GPU가 다른 GPU보다 많은 사용량을 보인다고한다면 - 사용랑은 nvidia-smi로 체크할 수 있다!!

 

그 한개의 GPU 때문에 더 많은 Batch_size로 학습 할수가 없다 

이는 엄청난 시간의 차이를 만들 수 있다

 

이런 경우 해결 방법이 없을까?

바로 출력 GPU의 위치를 조정해주는 것이다

 

앞에서 gather 과정에서 출력을 하나의 GPU로 모아준다고 했었다

이러한 gather는 디폴트 값이 정해져 있어서 정해진 GPU로 보내지지만 이를 조정해서

출력이 모이는 GPU를 조정한다면 어느정도 해결 할 수 도 있을 것이다

import os
import torch.nn as nn

os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1, 2, 3'
model = nn.DataParallel(model, output_device=1) #output  GPU 지정

간단하게 output_device를 파라미터로 추가해주기만 하면 조절할 수 있다

 

그러나 여기서 의문이 생길 수가 있다

단지 출력 GPU만 조절해준다고 불균형이 사라질까?

 

사실상 조삼모사에 가까운걸 아닐까?

 

사실 위의 방법은 임시적인 방법이지 이상적인 방법은 아니다

따라서 이를 해결하기 위해서는 Custom으로 조정을 해주어야한다

이제부터 Custom으로 Data 병렬을 하는 방법을 알아보도록 하자

---

✅ 2. Custom DataParallel 

사실 GPU의 불균형이 일어나는 이유는 간단하다

gather 과정에서 출력을 하나로 모아주기 때문이다

 

앞에서 살펴 보았듯 이 하나의 GPU로 출력을 모으는 이유는 loss를 계산하기 위해서이다

그렇다면 loss 계산을 병렬적으로 처리하면 어떨까?

 

고맙게도 

GitHub - zhanghang1989/PyTorch-Encoding: A CV toolkit for my papers.

loss를 병렬처리해서 출력하는 코드가 이미 다른 누군가에 의해서 만들어져있다

 

from torch.nn.parallel.data_parallel import DataParallel

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

 

과정은 앞에서 데이터가 병렬로 나누어지고 모아지는 과정과 비슷하게 흘러간다

 

이를 적용해서 학습하는 코드는 아래와 같다

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

model = ResNet18()
model = DataParallelModel(model)
model.cuda()

criterion = nn.NLLLoss()
criterion = DataParallelCriterion(criterion)  ##앞에서 작성한 코드 적용

...

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

---

✅ 3. Pytorch Distributed

Pytorch는 DataParallel과 함께 Distributed도 지원해준다

쉽게 말해서 분산처리를 하는 것인데 이는 여러 컴퓨터에서 모델을 학습하는 것과 같다

 

아래 pytorch 튜토리얼을 확인하면 분산 처리의 과정을 이해할 수 있을 것이

Writing Distributed Applications with PyTorch — PyTorch Tutorials 2.0.0+cu117 documentation

 

이를 학습에 적용하고 싶다면 아래의 예시 코드를 참고하기를 바란다

GitHub - pytorch/examples: A set of examples around pytorch in Vision, Text, Reinforcement Learning, etc.

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 = ResNet18()
    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)

솔직히 이 부분은 코드를 봐도 직관적으로 이해가 잘 되지 않는다

조금더 공부를 해봐야겠다..

 

참고한 블로그에서 설명을 가져오면 다음과 같다

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 연산을 수행하는 역할을 합니다.


<출처 : 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>

 

그리고 DataLoader의 경우 아래와 같이 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)

이런식으로 분산학습을 진행하면 GPU의 최대 성능을 활용해서 분산 처리를 할 수 있다

 

그러나 이러한 분삭학습의 경우 문제점이 하나 발생할 수 도 있다고 한다

모델에서 학습에 사용하지 않는 parameter가 있는 경우에 문제가 발생한다고 한다

 

이를 해결하기 위해서 다양한 패키지를 활용해 볼 수 있다

---

✅ 4. Apex

NDIVIA에서 만든 Apex의 경우 모델을 효율적으로 학습하게 해준다

 

이부분은 조금 CS적인 지식이 필요하기 때문에

개인적으로 자세하게 이해는 하지 못했지만

대충 설명하자면

 

보통 딥러닝의 학습은 32비트를 통해서 이루어진다 

그러나 Apex를 활용하면 16비트로 학습을 할 수있기 때문에

효율적인 학습이 가능하다 이를  Mixed Precision 연산 이라고 하는데

사실 이부분은 우리가 보고 있는 분산처리와는 관계가 없다

 

그러나 Distributed DataParallel, 일명 DDP 기능도 존재하는데

이부분이 우리가 공부하고 있는 부분이다

 

간다하게 아래 링크로 들어가면 예시 코드를 볼 수 있다

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

구현 방법은 크게 어렵지 않은것 같다

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 = ResNet18()
    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)

Pytorch와 비슷하게 DDP를 import 해서 mdoel을 덮어주면 된다

그리고 아래 코드를 입력해서 실행해주면 된다

 

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

이 부분에서 nproc_per_node는 gpu 갯수에 맞춰주면되고 

아래 입력되는 값들은 각 GPU에 들어가는 batch size와 num_workers 수이다

 

---

대부분 자연어 task와 같이 거대 모델을 활용할 때 분산처리나 Apex를 많이 활용한다고 한다

하지만 img같은 분야에서는 Data Parallel만 해도 충분히 처리할 수 있을 것이라고한다

 

아직 GPU가 1대라 분산처리를 실습해볼 수는 없지만

언젠간 마주치면 다시 이 글을 읽으면서 방법들을 정리해 보아야겠다

 

---

⭐Reference

 

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

 

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

 

💥 Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups

 

Nvidia Apex를 이용한 모델 학습 최적화