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

DDP를 사용하려면 다음과 같은 설정단계를 거쳐야 합니다.

 

---

첫 번째 단계는 초기화 단계입니다.

 

전체 GPU의 개수가 몇 개인지를 설정합니다.

그리고 현재 프로세스가 사용하는 GPU 번호를 정합니다.

이것을 위해서 아래 코드를 이용하면 됩니다.

 

# 1번
    dist_url = 'env://'
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    local_rank = int(os.environ['LOCAL_RANK'])

# 2번
    torch.distributed.init_process_group(backend='nccl', init_method=dist_url, world_size=world_size, rank=rank)
    torch.cuda.set_device(local_rank)
    torch.distributed.barrier()

# 3번
    if torch.distributed.get_rank() == 0:
        print(f'RANK {rank}    WORLD_SIZE {world_size}     LOCAL_RANK {local_rank}')

 

1번에서 world_size는 사용되는 전체 GPU의 개수입니다.

 

local_rank는 이 프로세스가 사용하는 GPU의 번호입니다.

DDP에서는 1개의 프로세스가 1개의 GPU를 사용하는 구조입니다.

예를 들어, 4개의 GPU가 있다면, 모두 4개의 프로세스가 각자 자신의 GPU를 사용합니다.

이 때, 프로세스 입장에서 자신이 사용하는 GPU의 번호를 local_rank라고 합니다.

 

local_rank를 환경변수 LOCAL_RANK에서 받아옵니다.

프로그램이 실행되면서 이 변수가 자동으로 설정되기 때문입니다.

 

2번에서는 그룹을 초기화 하고, 프로세스가 사용할 GPU 번호를 실제로 설정하는 과정입니다.

 

torch.distributed.barrier()는 여기까지 수행된 후, 더 진행하기 전에 다른 프로세스들이 여기까지 수행하기를 기다릴 수 있도록 합니다.

이것을 이용하면 1개 프로세스가 먼저 독주하여 동기화가 흐트러지는 것을 막을 수 있습니다.

 

3번에서는 메인 프로세스만 출력하도록 하는 방법을 보여줍니다.

여러  프로세스들을 수행되기 때문에 1개의 print 문 일지라도 각 프로세스가 실행하면 그 개수만큼 출력됩니다. 

이렇게 되면 출력이 혼잡스러워질 것입니다.

이를 해결하기 위해 메인 프로세스만 출력하도록 하면 됩니다.

자신이 메인 프로세스일 때만 print문을 수행하는 것입니다.

메인 프로세스의 rank = 0로 확인할 수 있습니다.

 

---

두 번째 단계는 데이터 샘플러 설정단계입니다.

여러 프로세스들이 1개의 데이터세트를 공유하기 때문에, 샘플러가 필요합니다.

데이터세트는 이전과 동일한 방법으로 만들면 됩니다.

데이터 로더에서 sampler만 새로 생성된 것을 지정해주면 됩니다.

 

    sampler = DistributedSampler(dataset=dataset_train, shuffle=True)

    dataloader_train = DataLoader(  dataset_train, 
                                    num_workers=4, 
                                    batch_size=128//4, 
                                    collate_fn=collater, 
                                    pin_memory=True,
                                    sampler=sampler)

---

세 번째 단계는 모델을 DDP 모드로 설정하는 것입니다.

 

    # 1번
    retinanet = torch.nn.SyncBatchNorm.convert_sync_batchnorm(retinanet)
    
    # 이 환경변수는 torch.run으로 실행하면 자동적으로 설정되는 것으로 판단됨
    local_rank = int(os.environ['LOCAL_RANK'])  
    
    # 2번
    retinanet = torch.nn.parallel.DistributedDataParallel(retinanet, 
                       device_ids=[local_rank], 
                       find_unused_parameters=True)

1번은 batch normalization 모드를 DDP에 맞춰서 설정하는 것입니다.

 

2번에서는 모델을 DDP 모드로 맞추고, 이 모델이 어느 GPU에 올라가야 하는지 설정합니다.

find_unused_parameters=True 는, 훈련 중 사용되지 않는 파라미터가 있을 때, 에러가 발생하지 않도록 설정해 줍니다.

 

---

네 번째 단계는 훈련단계에서 sampler에게 epoch을 알려줍니다.

 

    for epoch_num in range(parser.epochs):
        dataloader_train.sampler.set_epoch(epoch_num)

 

훈련에서 매번 새로운 epoch을 시작할 때 마다, 데이터로더는 자신의 샘플러에게 epoch을 알려줍니다.

아마도, 여러 프로세스들이 데이터세트를 공유하다보니 동기화와 관련된 것 같습니다.

 

---

다섯 번째 단계는 훈련된 모델을 저장하는 단계입니다.

 

모델을 저장하는 것은 하나의 프로세스만 하면 됩니다.

메인프로세스가 하는 것이 적절합니다.

따라서, 메인프로세스를 확인한 후에 모델 저장 명령을 내리면 됩니다.

 

if torch.distributed.get_rank() == 0:

    torch.save(retinanet.state_dict(), f'./{parser.saved_path}/model_{epoch_num}.pt')

---

마지막 단계는 훈련을 시작하는 쉘명령어 입니다.

 

Pytorch 1.9 이상에서는 torchrun 명령어를 이용해서 실행합니다.

 

torchrun --standalone \
    --nnodes=1 \
    --nproc_per_node=4 \
    _12_train_DDP.py \

--standalone은 다른 서버없이 단독 서버로 훈련하는 것을 의미하는 것 같습니다.

 

--nnodes는 훈련에 사용되는 서버의 전체 개수입니다.

 

-- nproc_per_node는 서버에 설치된 GPU의 개수입니다.

 

마지막에는 수행한 파이썬 코드를 지정합니다.