GPU연산 DP, DDP 공부해보기

이번에 아는 분들이랑 대회를 참여하고 있는데 나는 개인 GPU가 없다보니 base line에서 어떤 변화가 더 좋은지 알아보는 역할을 맡았는데 다른 분들이 코드 관련해서 이야기하는데 DP, DDP 이야기를 하더라. 무의식적으로 병렬처리 관련된 이야기인줄은 알았지만 제대로 알 수 없어 이번 기회에 한번 공부해보려고 한다. 해당 내용에 대한 참조는 글 밑에 있습니다.

 

DP(Data Parallelism)란?

텐서플로는 모르겠지만 torch에서는 torch.nn.DataParallel 로 작동이 가능하다. 이러한 과정에서도 Forward Pass(데이터를 처리 및 loss발생)와 Backward Pass(gradient계산)으로 나뉜다.

 

Forward pass

이를 한번 단계별로 설명해보자면

1. mini batch를 gpu개수 만큼 분할해 Scatter.

2. model의 파라미터를 각각의 gpu에게 Broadcast.

3. 각각의 gpu안에 있는 모델들은 forward 하여 logit계산.

4. 각각 계산된 logit을 Gather하여 하나의 gpu에 모은다.

5. logit에 대하여 loss계산

 

Backward pass

1. Loss에 대해 각 gpu에 Scatter.

2. 전달받은 Loss에 대해 각 gpu에서 Backward 하여 gradient 계산

3. 계산된 gradient 를 Reduce 하여 하나의 gpu에 더한다.

4. 총합된 gradient 이용하여 모델 업데이트

 

문제점은 각각의 gpu에서 계산된 gradient에 대해 하나의 gpu로 합친 후 업데이트 한 후(각 gpu에서 gradient를 모으고 하나의 gpu에게 모으는 과정역시 bottleneck이 걸릴 듯 하다), 업데이트된 모델을 다른 gpu에게 다시 복제(Broadcast)해줘야 하는데, 이 과정이 비싸다고 합니다. 따라서 해당 문제점에 대한 해결책으로는 Gradient 를 Gather하지 않고 각각의 gpu자체적으로 step() 을 수행하는 방법을 고안하게 됩니다. 

 

해당 문제점에 대한 해결은 각각의 gradient에 대해 하나의 gpu에 모아서 업데이트하는 방식 대신, 각각의 gradient들을 모든 gpu에게 다 뿌려주고 자체적으로 업데이트하면은 각각의 gpu는 업데이트된 모델을 가지게 됩니다. 그러나 이러한 All-reduce방법은 비싸다고 합니다(바로 밑의 그림에서는 마스터 gpu에 대한 그림이 없는데 저도 잘 모르지만 이렇게 단순하게 gradient를 분배할 수 없고, 밑의 두번째 그림처럼 마스터 gpu를 통해 gradient를 합한 후 분배하는 방식으로 해야하나 봅니다).

'

 

 

따라서 효율적인 All reduce 방식으로 DDP(Distributed Data Parallel)방법이 있습니다.

 

DDP(Distibuted Data Parralel)

해당 방법은 multi process 모듈이며, All-reduce를 활용하면서 마스터 프로세스의 개념이 없어 학습과정이 심플하다고 합니다.

 

해당 이미지를 보면은 각각의 multi process안에서 gradient가 발생하고 모델이 업데이트 된다는 것을 알 수 있습니다. 핵심이 되는 효율적인 all reduce방식은 다음과 같은 알고리즘으로 수행되는데 ring 무슨 방식이었던 거 같다. 이러한 방식을 통해 모든 gpu들이 다른 기기에서의 gradient도 순차적으로 얻을 수 있으며 한 gpu에 모을 필요 없이 parameter들을 업데이트할 수 있다(해당 방식은 한 곳에 gradient를 모아 업데이트 하는 방식을 사용하지 않을 뿐더러 업데이트 된 파라미터들을 broadcast하는 과정도 필요없으니 매우 효율적으로 사용할 수 있을 것 같다) 

DDP를 활용한 코드는 다음과 같다.

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)
    
  model = Bert()
  model.cuda(args.gpu)
  model = DistributedDataParallel(model, device_ids=[args.gpu])

하나의 데이터셋을 여러 gpu들이 공유하므로 sampler가 필요하다. 

 

그렇다면 기존 DP에 대한 첫번째 문제점인 single process에서 multi process로 만들었는데 All-reduce는 어디서 일어나는 걸까?

backward() 와  step() 연산을 비교해보면 backward()하는 연산이 훨씬 무거운 연산이라고 한다. 따라서 이러한 무거운 연산을 중첩시킬 수록 학습과정을 수행하는 시간이 짧아진다고 한다(이 말은 무슨 뜻인지 모르겠다). 그래서 backward()가 끝날 때까지 기다리는 것보다 All-reduce를 수행하는 것이 더 빠르다고 한다.

 

---

References

https://nbviewer.org/github/tunib-ai/large-scale-lm-tutorials/blob/main/notebooks/05_data_parallelism.ipynb

Jupyter Notebook Viewer

https://swprog.tistory.com/entry/%ED%8C%8C%EC%9D%B4%ED%86%A0%EC%B9%98Pytorch-Distributed-Data-Parallel-DDP%EC%82%AC%EC%9A%A9%ED%95%98%EA%B8%B0

파이토치(Pytorch) Distributed Data Parallel (DDP)사용하기

 

깊게는 안들어가고 기본적인 개념에 대해서 한번 공부해보았습니다. 헷갈리는 부분도 많고 아직 쓸 일은 없을 거 같지만 나중에 필요할 경우 다시 와서 공부할 거 같습니다.