Pytorch DDP

들어가기 전 환경 세팅

들어가기 전에 PyTorch + cuda를 설치하면 되는데 이는 공식 홈페이지를 참고해서 그대로 진행하시면 됩니다. 파이토치 설치는 그대로 진행하면 되는데, 간혹 cuda를 manually 설치를 하는 분들이 있는 것 같습니다. 개인적으로는 추천하지 않습니다. 세팅할 것도 많고 특히 window에서는 설정해줘야하는 부분들이 많습니다. 그래서 *miniconda나 venv 등 가상환경을 꼭 만들어서 진행하는걸 *강력히 권고합니다. 아래에 간략히 제가 주로 세팅하는 방법들을 알려드립니다

*저 같은 경우 anaconda보다는 miniconda를 활용하는 편입니다. 가상환경을 만드는데 있어서 완전 필수적인 요소만 설치해서 훨씬 가볍습니다.

1. miniconda 설치:

Linux:

1. 링크 접속 및 Miniconda3-latest-Linux-x86_64.sh 다운
2. bash Miniconda3-latest-Linux-x86_64.sh 실행

Window:

1. 링크 들어가서 Miniconda3-latest-Windows-x86_64.exe 다운 후 실행

2. miniconda 가상환경 만들기:

Linux:

1. 터미널에서 conda create -n your_own_env_name python=3.9

Window:

1. 시작에서 Anaconda Prompt 실행
2. 터미널에서 conda create -n your_own_env_name python=3.9

*your_own_env_name에 원하는 형태의 가상환경 이름 설정해주시면 됩니다.

3. miniconda 내 패키지 설치 (Linux & Window 동일)

1. 터미널에서 conda activate your_own_env_name 실행
2. 링크에서 원하는 버전 및 OS 선택 후 터미널에서 실행 (pip / conda 둘다 상관 X)
3. *무조건 cuda를 포함한 명령어 실행 (ex. conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 pytorch-cuda=11.8 -c pytorch -c nvidia)

*추가로 최근에 numpy를 설치할 때 version 2가 설치가 될때가 있는데 그럴 때는 version 1을 새로 설치해주시면 됩니다. (pip install numpy==1.26. 등)

그러면 이제 가상환경 내에 PyTorch 및 cuda가 자동으로 설치가 됩니다. 혹시 cudnn 등 추가적인 cuda 환경을 설치하고 싶으신 분들은 conda install -c anaconda cudatoolkit==[원하는 버전] 및 conda install -c anaconda cudnn 을 실행해주시면 cudatoolkit 원하는 버전 및 cudatoolkit에 맞는 cudnn 버전이 설치가 됩니다.

DDP 적용하기

이제 바로 본론으로 들어가겠습니다. DDP 적용 관련해서 2가지로 나눠서 보겠습니다. 첫번째는 terminal 및 script 입력 방법이고, 두번째는 파이썬 코드 내 세팅입니다.

Terminal 및 Script 입력 방법

우선 파이썬 코드 상으로 DDP가 완료된 상황이라는 가정하에, 다음과 같이 입력하시면 됩니다.

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 --master_port 56789 main.py

하나씩 살펴보도록 하겠습니다.

• CUDA_VISIBLE_DEVICES=0,1,2,3: 이는 본인이 사용할 수 있는 GPU 중 몇번 GPU를 사용할지를 지정해줍니다. 즉, 위 예시에서는 본인 로컬 및 server에 8개의 GPU가 있다고 가정한다면 0~3번 GPU만 사용하겠다는 의미입니다. 만약 따로 지정을 안해주면 전부 사용하게 됩니다.
• torchrun: PyTorch version 2.0 이전에는 python -m torch.distributed.launch 사용했었는데 version 2.0 이후에는 python -m 등을 따로 입력할 필요 없이 DDP를 실행하게 해줍니다.
• --nproc_per_node=4: 해당 node에서 총 몇개의 processor를 사용할지를 정할 수 있습니다. 여기서 processor란 GPU를 의미합니다. 즉, gpu 개수와 동일하게 맞춰주면 됩니다.
  (위 예시 같은 경우는 4개의 GPU를 사용하므로 4; node는 해당 로컬 및 server)
• --master_port 56789: 사실 이 부분은 optional인데 DDP 실행 시 port 번호를 정해줍니다. 간혹 로컬 및 server에서 여러개의 DDP를 실행할 수도 있는데, 이 때 port 번호가 겹쳐서 실행이 안될 때가 있습니다. 그럴때는 본인이 좋아하는 번호를 아무거나 입력해주시면 됩니다.
• main.py: 실행하고자하는 파일 명입니다. (너무 당연하지만…)

적용할 수 있는 argument가 더 있긴 있습니다. 예를 들어, --master_addr 가 있는데, 이 경우는 multi node를 활용할 때, 즉 여러개의 server를 활용할 때 사용하게 됩니다. 다만, 이 posting의 경우 DDP를 처음 접하는 분들을 위한 내용 위주이고, 대부분 multi-node를 활용할 일이 없을 것 같아서 생략하도록 하겠습니다.

파이썬 코드 내 세팅

여기도 들어가기에 앞서 짧게 말씀드릴 부분이 있습니다. DDP는 위의 --master_addr 같은 argument들을 지정하는 것을 보면 알겠지만 GPU별로 각각 일종의 가상환경을 만들어주는 것과 비슷합니다. 즉, 각 GPU별로 가상환경처럼 설정을 하여 파이썬 파일을 돌리는 걸로 이해해주시면 쉽게 이해하실 듯 합니다. 그 과정에서 각 GPU별로 파이썬 내에 변수들이 설정이 되는데(os.environ['변수명']) 2가지 변수들에 관해서 짧게 짚고 넘어가겠습니다.

• os.environ['WORLD_SIZE']: torchrun을 통해 실행되는 node 번호를 말합니다. 즉, 실제로는 server 번호가 됩니다. 여기서 server 1개를 기준으로(single node) 설명하고 있기 때문에 대부분의 경우 os.environ['WORLD_SIZE']=0 이 되겠습니다.
• os.environ['LOCAL_RANK']: 여기는 torchrun을 통해서 실행되는 각 processor를 의미합니다. 즉, GPU 번호로 이해해주시면 될 것 같습니다. (현재 예시에서는 0~3까지).

DDP Initialization

이제 이를 바탕으로 파이썬 코드 내에서 어떻게 세팅하는지 살펴보겠습니다. 각자마다 코딩하는 스타일이 있겠지만 저는 아래와 같이 작성하곤 합니다.

args.device = 'cuda:0'
args.world_size = 1
args.rank = 0
args.local_rank = int(os.environ.get("LOCAL_RANK", 0))
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.local_rank = torch.distributed.get_rank()

일단 이건 제 코딩 스타일 방법인데, 저는 args.xxx을 자주 활용하곤 합니다. (따로 변수들을 설정할수도 있겠지만 argument 안에 넣어두면 코드 어느곳에서도 사용할 수 있어 편하더라구요.)

그럼 코드들을 하나하나 살펴보도록 하겠습니다.

• args.device='cuda:0' ~ args.rank=0: DDP에 필요한 변수들을 initialize 해주는 단계라고 이해해주시면 될 것 같습니다.
• args.local_rank = int(os.environ.get("LOCAL_RANK", 0)): 위에서 말씀드린 것처럼 os.environ['LOCAL_RANK']' 의 경우는 코드가 돌아가는 해당 GPU 번호를 의미하게 됩니다. 따라서, os.environ 내부에 있는 get() function을 통해 args.local_rank에 GPU 번호에 해당하는 변수를 받는 부분으로 이해해주시면 되겠습니다.
  • 예시 1): 0번 GPU에 해당한다면 args.local_rank = 0
  • 예시 2): 3번 GPU에 해당한다면 args.local_rank = 3
• torch.cuda.set_device(args.local_rank): 보통 PyTorch에서 GPU에 parameter를 올리는 방법이 2가지가 있는데 params.to('cuda:#')와 params.cuda()가 있습니다. 첫번째 방법은 다들 많이 사용하는 방식이라 설명은 생략하겠습니다. 2번째 방법은 현재 default GPU에 올리는 방법입니다. 이때, default GPU는 보통 실행하는 첫번째 GPU (0번째)로 지정이 됩니다. 여기서 변경을 할 때 set_device 사용할 수 있습니다. 즉, 이 경우는 위에서 지정한 args.local_rank를 통해 default GPU를 지정해주는 부분입니다.
• torch.distributed.init_process_group(backend='nccl', init_method='env://'): 이 부분이 결국 최종적으로 각각 GPU를 잘 동작시킬 수 있게 initalize 해주는 부분이라고 이해해주면 되겠습니다. 여기서 backend=’nccl’이 있는데 처리해주는 동작을 cuda로 지정해주는 부분이고, init_method=’env://’ 부분은 현재의 DDP를 위해 생성된 임의의 가상환경을 지정해주는 부분입니다. 대부분 이 형태를 사용하기때문에 굳이 변경할 필요는 없을 것 같습니다.
• args.world_size, args.local_rank: 저는 최종적으로 지정된 WORLD_SIZE 와 LOCAL_RANK 를 이 두개에 활용하곤 합니다. args.local_rank 의 경우 확인 차원에서 한번 더 initalization 하는 걸로 이해해주시면 됩니다.

Model에 DDP 적용하기

다음으로는 model에 DDP를 어떻게 적용 및 학습을 하는지 살펴보겠습니다.. (매우 간단합니다!!

from torch.nn.parallel import DistributedDataParallel as DDP
model = DDP(model,device_ids=[args.local_rank])
...
logits = model(x)
loss = loss_fn(logits, labels)
loss.backward() 

보시면 PyTorch에서 제공하는 DistributedDataParallel 를 통해 모델을 분배해줍니다. 그 과정에서 이전에 지정해뒀던 args.local_rank를 통해 각 GPU들에 분배를 해주는 형태입니다. (device_ids=[args.local_rank] , args.local_rank에 대괄호 필요).

Additional Tip - export NCCL_P2P_DISABLE=1

아쉽게도 가끔 원활히 동작하지 않을 수 있습니다… Configuration이 꼬이거나 충돌하는 등 다양한 원인으로 인해 동작하지 않을 때가 있습니다. 그래도 요즘은 LLM (ChatGPT, Claude, Gemini 등등)이 있어 오류 메세지만 잘 캐치해낸다면 디버깅이 훨씬 수월해지긴 했습니다. 코드 내부의 문제가 아니라면 적절히 활용하는 것을 강추합니다.

그럼에도… 오류 메시지가 뜨지 않고 무한 loading에 빠지는 등의 문제가 생길 때도 있습니다. 제 경우였는데, DDP함수에 딱 들어가면서 갑자기 freeze된 후 무한 loading되는 현상이였습니다. 정확한 원인은 모르겠으나 아마 DDP함수 내부 동작 중에 GPU끼리 모델을 처리할 때 문제가 생기는 상황인 것 같습니다… (PyTorch framework 내부까지 들어가서 일일이 print 해보면서 디버깅해봤던 적도 있는.. )

그럴 때 export NCCL_P2P_DISABLE=1 를 한번 시도해보셔도 될 것 같습니다. 터미널에서 한번 실행해주시고 코드를 돌리면 원활히 동작하더라구요ㅎㅎ

마치며….

처음으로 coding 관련해서 posting을 해봤는데 PyTorch 및 DDP를 처음 입문하시는 분들에게 유익한 정보이길 바랍니다.ㅎㅎ 저도 그랬고 주변을 보면, DDP를 처음 시도하실 때 꽤 오랜시간동안 struggling하곤 하더라구요.(ㅠㅠ) 블로그를 찾아보면서 적용해봐도 안되고 LLM에 물어봐도 불친절할 때가 종종 있는데, 그런 경우에 이 글이 큰 도움이 됐길 바랍니다!

그럼 다음 글에서 뵙도록 하겠습니다! (maybe 논문리뷰?)