torchrun과 Process Group 시작하기

PyTorch distributed training은 보통 GPU마다 하나의 process를 띄우는 방식으로 시작한다.

GPU 0 -> process 0 -> rank 0
GPU 1 -> process 1 -> rank 1
GPU 2 -> process 2 -> rank 2
GPU 3 -> process 3 -> rank 3

torchrun은 이 여러 process를 실행하고, 각 process에 rank와 world size 정보를 전달한다.

torchrun --nproc_per_node=4 train.py

각 process는 시작할 때 process group을 초기화한다.

import torch.distributed as dist

dist.init_process_group(backend="nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

이 단계의 목표는 아직 DDP가 아니다. 먼저 각 process가 자기 GPU를 잡고, 서로 통신할 수 있는 group에 들어가는 것을 확인한다.

확인

• PyTorch distributed에서 보통 GPU마다 process를 하나씩 띄우는 이유는 무엇인가?
• rank와 local rank는 어떻게 다른가?
• init_process_group은 어떤 준비를 해주는가?