[PYTORCH] DistributedDataParallel이란? - Nvidia APEX로 구현하기

Introduction

---

행렬(이하 Tensor)연산을 기반으로 하는 Deep Neural network는 CPU에서 연산을 할때보다 GPU를 통해서 그래픽 연산을 할때 더 높은 퍼포먼스를 보인다. Deeplearning이 점차 발달하면서 신경망을 더욱 깊고 크게 설계하는 것이 트렌드가 되면서 이제는 GPU 한장으로는 학습하기에 벅차다. (hpyer-scale...) 또한 Deeplearning 을 점목하는 task가 점점 늘어나고, 이미지보다 더 큰 동영상, 음원 등을 학습데이터로 사용하기 때문에 한개의 GPU로 학습시키기에는 성능/메모리의 제한이 있다.

그래서 이제는 하나의 GPU만을 사용해서 학습을 시키는 것이 아닌 여러개의 GPU를 사용하는 multi-GPU 학습이 필요하다. Pytorch에서는 DataParallel, DistributedDataParallel 기법을 통해 multi-GPU를 사용할 수 있도록 할 수 있다. Nvidia에서도 apex라는 라이브러리를 사용하여 (물론 16-bit precision연산이 핵심이지만) DistributedDataParallel(DDP)를 할수 있는 기능이 포함되어있다.

일단 DDP를 적용하게되면 (동일한 그래픽카드를 여러개 보유하고 있다는 가정 하에, 동일한 그래픽카드가 아니면 고난 시작이기 때문에 그냥 따로 쓰는걸 추천...) 배치사이즈도 더 크게, 학습시간은 더 짧게 할 수 있다. 필자는 환경만 받쳐준다면 multi-GPU 환경을 한번쯤은 구축해보는 것을 추천한다. 

그래서 Pytorch에서 제공해주는 DataParallel, DistributedDataParallel을 사용할 수 있는데, 사용하다보면 여러모로 세팅해야할 것도 많고 간간히 문제도 많이 발생한닫고 한다. 그래서 Nvidia의 APEX 라이브러리를 사용해서 DistributedDataParallel을 사용해볼 것이다. 

 

NVIDIA APEX SETUP

---

사실 Nvidia Apex는 Mixed precision 연산을 위해 만들어진 라이브러리인데 여기에 Distributed 기능도 포함되어있다. 아 Mixed precision은 딥러닝에서 연산속도를 높이기 위해 사용한다. 보통 딥러닝은 32bit 연산을 수행하는데 bit수를 낮춰서 속도를 높이는 것이라고 한다. 

 

apex를 사용하기 위해서는 당연히 GPU환경이 구축회더있어야 하고 몇개 라이브러리를 설치해야한다. 필자는 리눅스 + 도커컨테이너를 사용한다.

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

위 코드만 터미널에서 실행시키면 끝 (파이썬은 라이브러리 설치는 초간단이지만 한번 꼬이면 이만한 지옥도 없다)

 

 

QUICK START

---

이제 초간단으로 DDP를 어떻게 적용하는지 알아보도록 한다.

# FOR DISTRIBUTED: (can also use torch.nn.parallel.DistributedDataParallel instead)
from apex.parallel import DistributedDataParallel

우리는 apex의 DDP 기능을 사용해서 분산학습을 진행할 것이라서 apex의 DistributedDataParallel을 import 해준다.

 

parser = argparse.ArgumentParser()
# FOR DISTRIBUTED:  Parse for the local_rank argument, which will be supplied
# automatically by torch.distributed.launch.
parser.add_argument("--local_rank", default=0, type=int)
args = parser.parse_args()

argparse를 세팅한다. 이때 --local_rank는 필수로 추가해야하는 argument이다. 내가 사용하는 프로젝트가 argparse를 사용하지 않는다 하더라도 (일반적으로는 많이 사용하긴 하지만) 필수로 세팅되어있어야한다. 이것은 나중에 실행할때 명령어로 torch.distributed.launch를 포함하는데 이것과 연관이 있다.

 

# FOR DISTRIBUTED:  If we are running under torch.distributed.launch,
# the 'WORLD_SIZE' environment variable will also be set automatically.
args.distributed = False
if 'WORLD_SIZE' in os.environ:
    args.distributed = int(os.environ['WORLD_SIZE']) > 1

DDP 즉 분산시스템에서는 gpu가 꽂혀있는 컴퓨터를 node라고 부르고 (일단 필자는 하나의 컴퓨터에 여러개 GPU가 있는 상태로 사용할 것이기 때문에 node는 0) 총 gpu의 개수는 world size라고 한다. 그래서 여기서는 gpu의 개수가 한개 이상이면 args.distributed를 True로 바꿔주며 밑에 distributed 관련 코드가 실행될 수 있도록 하는 것이다.

 

if args.distributed:
    # FOR DISTRIBUTED:  Set the device according to local_rank.
    torch.cuda.set_device(args.local_rank)

    # FOR DISTRIBUTED:  Initialize the backend.  torch.distributed.launch will provide
    # environment variables, and requires that you use init_method=`env://`.
    torch.distributed.init_process_group(backend='nccl',
                                         init_method='env://')

torch.backends.cudnn.benchmark = True

 위에서 args.distributed가 true가 되면 코드가 실행될 것이다. 여기서 torch.cuda.set_device로 local_rank가 세팅된다. 일단 우리는 한개의 컴퓨터에서 연산을 진행하기 때문에 local이라는 키워드가 붙는 것이고 (만약 연산하는  컴퓨터가 여러개라면 global) rank는 컴퓨터의 process id라고 생각하면 된다. 기본적으로 여러개 컴퓨터를 사용하는 구조가 포함되어있는 패키지이기 때문에 한개의 컴퓨터에서 쓴다면 그냥 0이라고 쓰면 된다.

두번째로 process_group은 어떤 것을 정하면 되는 것이냐면 학습을 위해 AllReduce라는 작업을 사용하는데, 이것은 gpu간 통신을 통해서 gradient 계산, 업데이트 등을 하는 것이다. 여러개의 gpu에서 학습을 진행하기 때문에 각각의 gpu에서 연산된 결과가 있을 것이고 이것을 통합하여 gradient를 연산해야하기 때문에 backend가 필요하다. backend 종류는 nccl, mpi 등이 있는데 무난한게 nvidia의 통신 라이브러리인 nccl을 사용할 것이다.

 

N, D_in, D_out = 64, 1024, 16

# Each process receives its own batch of "fake input data" and "fake target data."
# The "training loop" in each process just uses this fake batch over and over.
# https://github.com/NVIDIA/apex/tree/master/examples/imagenet provides a more realistic
# example of distributed data sampling for both training and validation.
x = torch.randn(N, D_in, device='cuda')
y = torch.randn(N, D_out, device='cuda')

model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

이번 예제에서는 정말 어떻게 돌아가는지만 보는 것이기 때문에 임의의 input data와 임의의 traget data를 사용한다. 그냥 랜덤하게 x, y를 세팅한다. 그리고 모델도 초간단으로 linear 한개만 사용할 것이다. optimizer도 그냥 만들어주면 된다. 

그 다음에 amp.initialize 함수를 통해서 amp 기능을 초기화시켜주면 된다. 이것은 apex에서 mix precision을 사용할 수 있는 amp(automatic mixed precision)툴이다. 안쓸꺼면 안써도되는것 중에 하나이다. 아 apex는 a pytorch extension의 약자이다. 여기서 opt_level은 16bit를 얼마나 적용할지 크기를 결정하는 것이다.

 

if args.distributed:
    # FOR DISTRIBUTED:  After amp.initialize, wrap the model with
    # apex.parallel.DistributedDataParallel.
    model = DistributedDataParallel(model)
    # torch.nn.parallel.DistributedDataParallel is also fine, with some added args:
    # model = torch.nn.parallel.DistributedDataParallel(model,
    #                                                   device_ids=[args.local_rank],
    #                                                   output_device=args.local_rank)

그리고 마지막으로 모델은 DistributedDataParallel으로 감싸기만 하면 된다. 

 

loss_fn = torch.nn.MSELoss()

loss도 어차피 토이프로젝트이기 때문에 MSELoss 적용하였다.

 

for t in range(100000):
    if t % 100 == 0:
        print("epoch: ", t)
    optimizer.zero_grad()
    y_pred = model(x)
    loss = loss_fn(y_pred, y)
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_loss.backward()
    optimizer.step()

if args.local_rank == 0:
    print("final loss = ", loss)

그리고 기존에 학습하는 방식돠 동일하게 학습하면 된다. 만약 amp를 위에서 initalize하였다면 학습때 Loss 계산할때에도 amp.scale_loss를 사용해서 backward 해주어야 한다.

 

이렇게 해서 코드 작성은 완료가 되었다.

이제 실행을 해보아서 여러개의 gpu에 잘 할당이 되고 돌아가는지 테스트해보면 될 것이다.

python -m torch.distributed.launch --nproc_per_node=2 distributed_data_parallel.py

터미널에서 아래와 같은 명령어를 적고 실행하면 된다. 여게서 nproc_per_node에 컴퓨터에 박혀있는 gpu의 개수를 적으면 된다 (만약 꽂혀있는것보다 작게 적으면 0번 gpu부터 순서대로 할당된다).

 

 

Source Code

---

https://github.com/waverDeep/DistributedDataParallel/blob/master/example02.py

GitHub - waverDeep/DistributedDataParallel

 

Reference

---

 

https://pytorch.org/docs/stable/distributed.html#basics

Distributed communication package - torch.distributed — PyTorch 1.10.1 documentation

https://github.com/NVIDIA/apex

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

https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py

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

https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed

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