[딥러닝]Pytorch multi GPU 학습하기(Horovod)

 

알아서 GPU를 전부 인식해주고 , 병렬처리를 지원해주면 좋겠지만 기본적으로 pytorch는 single gpu를 사용합니다. 이를 위해서는 추가적인 작업을 통해 병렬 처리를 하도록 설정해 주어야합니다. 

 

 

현재 제가 사용하고 있는 서버의 nvidia-smi 화면입니다. 

0번 gpu는 현재 학습에 활용되고 있지만 1번은 학습에 참여하고 있지 않습니다.  

 

데이터 병럴처리를 위한 방법은 찾아본 바에 의하면 크게 세가지가 있다고 합니다.

첫번째로 torch.nn의 DataParallel이 있습니다. 정의한 model을 wrap 해주기만 하면 되서 굉장히 사용하기 편합니다.

 

 

model = 아무 모델 

model = nn.DataParallel(model).to(device) 혹은  .cuda()

학습과정에서 바뀌는건 전혀 없이 이부분만 바꿔주면 됩니다.

 

그런데 아래 그림의 4번을 보면, 중간에 데이터를 전부 하나의 gpu로 모으는 과정이 있습니다. 이에 따라 메모리의 불균형이 일어나게 됩니다. 결국 남는 메모리 = 자원 낭비기 때문에 효율적이지는 못합니다. 

(사실 이 부분은 모델의 출력값이 큰 BERT같은 경우에 해당됩니다.이미지 분류 같은 경우에는 이 방법을 쓰는게 제일 좋을 거 같습니다) 

 

 

torch.nn.DataParallel의 과정

 

다른 방식으로는 Cutomized DataParallel이 있습니다. 

https://github.com/zhanghang1989/PyTorch-Encoding/blob/master/encoding/parallel.py

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

 

위에서 출력값을 통해 기울기를 계산하기 위해 gpu 하나로 모으는 과정을 조금 틀어서 gather하지 않고, 기울기도 각각의 gpu에서 병렬 계산하게 구현하였다고 합니다. 그래서 조금 다른점은 model뿐만 아니라 loss까지 패키지에서 제공하는 함수로 wrap을 해줍니다. 그런데, 호환의 문제인지 저 같은 경우에는 동작을 잘안하는 문제가 있었습니다. 코드는 원활하게 돌아가는데 multi- gpu가 아니라 단일로 돌아가는 문제가 있었습니다. . 

 

 

결국 시도했던 방법 중 마지막 방법만 정상적으로 multi-gpu 학습이 가능했는데요 오늘은 마지막 Horovod를 중심으로 설명을 해보겠습니다. 

 

 

 

주의!   Horavod는 리눅스 환경에서만 지원합니다. windows의 경우 앞서 얘기드린 다른 방법을 시도하는걸 권장드립니다.

 

Horovod 설치

 

HOROVOD_GPU_OPERATIONS=NCCL pip3 install horovod # gpu installation

이걸 통해 설치해봅시다. CPU 버전도 있는데 그냥 pip install horovod를 통해 설치하면 cpu 버전이 설치되므로 

주의하도록 합시다. 

 

Horovod 설정 

 

import horovod.torch as hvd
hvd.init()


torch.cuda.set_device(hvd.local_rank())

처음에는 horovod의 초기화를 진행해줘야합니다. 그리고 torch에 device를 할당해 줍시다 

 

추가로 datasampler와 이를 이용해서 dataset, dataloader를 적용해주자 

이 과정에서 shuffle = False로 해줍시다 

 

train_dataset = ...

# Partition dataset among workers using DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

공식 가이드에서는 이런식으로 표현해놨네요 

 

train_dataset = CustomDataset(train['img_path'].values, train_labels, train_transform)
train_sampler = torch.utils.data.distributed.DistributedSampler( train_dataset, num_replicas=hvd.size(), rank=hvd.rank())

train_loader = DataLoader(train_dataset, batch_size = CFG['BATCH_SIZE'], shuffle=False, num_workers=CFG['NUM_WORKERS'],sampler =train_sampler)


val_dataset = CustomDataset(val['img_path'].values, val_labels, test_transform)
val_sampler = torch.utils.data.distributed.DistributedSampler( val_dataset, num_replicas=hvd.size(), rank=hvd.rank())
val_loader = DataLoader(val_dataset, batch_size = CFG['BATCH_SIZE'], shuffle=False, num_workers=CFG['NUM_WORKERS'],sampler =val_sampler)

저는 이미지 처리에 사용했는데 이런식으로 적용했습니다 

 

model.cuda()

optimizer = optim.SGD(model.parameters())

# Add Horovod Distributed Optimizer
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())

뭐든 간에 정의해준 model을 cuda로 보내고 optimizer를 hvd에서 이용가능한 객체로 바꾸어주겠습니다 

 

이 부분은 dataloader에 관여하는 친구들은 비슷하게 적용되는 느낌이네요(ex accelerate)

 

hvd.broadcast_parameters(model.state_dict(), root_rank=0)

 

요건 어느 프로세스를 기준으로 학습을 할 건지 정해줍니다 

 

뒤에 과정은 일반적인 학습과정과 동일합니다 

 

horovodrun -np 4(사용할 GPU 숫자) python 실행파일.py

 

실행은 다음과 같이 진행해주겠습니다. 여기까지 진행하셨으면 gpu 를 모두 사용하는 걸 확인하실 수 있습니다.