[팁] multi-gpu 환경에서 대형 모델에 대해 (distributed) data parallel을 사용하고자 할 때

[cth127]

지난 주말 LLaMA 6B를 Distributed data paralle (DDP)를 적용하여 훈련하며 겪은 시행착오를 정리해보았다.

1. DDP는 각 GPU마다 별개의 Process를 생성하는 방식으로, 하나의 Process에서 모든 device를 다루는 Data parallel(DP)와 차이가 있음. 후자의 경우 GPU 간 VRAM 사용에 불균형이 발생하므로 일반적으로 DDP가 더 추천된다. 참고 링크
   

2. 다음 튜토리얼에 따라 DDP를 사용할 경우 다음과 같이 backend 모드를 설정해야 하는데, gloo, nccl, mpi 세 가지 옵션이 있으나 multi-gpu 상황에선 nccl이 가장 권장된다. (gloo의 경우 속도가 5-20배 느려짐을 확인)

dist.init_process_group("nccl") # nccl, mpi or gloo

3. 다만 nccl을 사용했지만 위의 튜토리얼 코드도 안 돌아가는 경우가 있었는데(freeze 돼버리고 무한 time out), 도커 외부에서 다음 커멘트를 따라 스크립트를 변경해주면 된다. (도커 바깥의 스크립트가 도커 내부 nccl 작동에 영향을 주는 것 같다?!)

4. (내 경우) 훈련 자체가 freeze 돼버리는 건 torch나 transformers의 문제가 아니라 nccl 자체의 문제였다. nccl 자체의 문제인지 확인하고 싶으면 NCCL Test를 돌려볼 것. 거기서 나오는 다양한 에러 메시지에 대한 대응 방법이 다양하게 적혀있다. (참고로 현재 NCCL Test repo는 ambers:/DATA/nccl-tests 에 clone 되어있으니 참고.)

5. DDP를 돌릴 경우 각 process 별로 완벽히 독립적인 processing이 가능하도록, 모든 코드 실행 전단계에 다음과 같이 해당 process가 볼 수 있는 gpu device를 설정해주는 것이 좋다. 참고로 아래 방식은 torchrun을 통해 DDP를 쓸 때 코드고, basic use case처럼 torch.multiprocessing 모듈을 사용할 땐 알맞게 변형이 필요.

local_rank = os.environ['LOCAL_RANK']
os.environ['CUDA_VISIBLE_DEVICES'] = local_rank

6. 튜토리얼엔 안 써있지만, train dataset에 해당하는 것을 torch.utils.data.distributed.DistributedSampler로 한번 감싸줘야 data가 각 device로 나뉘어 들어간다. 참고 링크

7. 모든 설정을 마친 뒤, 1 GPU에선 잘 돌아가는 코드가 multi-gpu에선 OOM이 나는 것을 발견했다! 확인 결과, DDP로 모델을 래핑하면 parameter 수만큼 VRAM을 더 쓰게 된다는데, 일반적으로 모델은 작고 데이터는 많은 경우엔 문제가 생기지 않는다. 하지만 이번처럼 모델이 엄청 큰 경우 batch size 1도 안 먹히게 된다. 그 경우 위의 링크에 써진 것처럼 DDP를 call할 때 gradient_as_bucket_view=True를 설정해주면 된다.

8. 참고로 7번같이 특수한 경우는 DDP를 end-to-end로 제공하는 여러 모듈들 (transformers trainer, accelerate)에선 지원을 하지 않(는 것으로 보이)기 때문에 training loop와 DDP 코드를 직접 짜서 들어가는 것이 속 편하다.

9. 현재 서버에 있는 GPU로 A6000과 A100이 혼합되어 있는데, Device별로 VRAM이 달라서 똑같은 설정을 일괄적으로 적용하는 건 무리. LLaMA 6b의 경우 DDP 모드로 A6000에 올라가지도 않는다.

10. 참고로 현재 CUDA_VISIBLE_DEVICES는 A100 두 대가 0, 1번, 나머지 A6000들이 2-7번으로 들어가 있다. nvidia-smi로 볼 수 있는 순서(A100: 3, 5번)과는 차이가 있으니 주의!

11. 현재 A100 간 연결이 약한 상태로, NVLink를 통해 업데이트 될 예정 [서버 관련] Multi-GPU support needed #114 !

그 결과 1000시간 걸리는 훈련을 600시간 정도로 축소할 수 있었다. 만세~

ps. LLaMA 13B부터는 batch size 1도 OOM이 뜨는 것을 확인했다. Model Parallel에 대한 별도의 가이드가 필요할 것 같은데, 나중에 시도해보고 작성해보겠음.