[Pytorch] DDP 환경에서 load_state_dict 사용시 OOM 발생 상황을 해결하자

DDP 환경을 구성하고 scratch 부터 학습할때는 잘 동작하던 코드에서 

backbone에 pretrained weight를 load 하니 바로 OOM이 발생하는 상황을 만나버렸다.

초기 상황 구성은 아래와 같았다.

# OOM이 발생한 코드
model = get_model()

checkpoint_path = r"~~~"
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint)

model = torch.nn.parallel.DistributedDataParallel(~~~)

 

단순히 초기 weight 만 읽어왔을 뿐인데 OOM이 터져버리니 당황스럽긴 한데 

메모리 동작하는 꼴을 보니 0번 GPU 가 메모리 터져서 발생한 상황인 듯 했다. 

 

추정키로는 각 GPU에 복사되어야 할 정보들이 0번 GPU로 먼저 다 올라오면서 

감당이 안되고 OOM 이 발생한 것 같다. 

이런 상황을 해결하기 위해서는 각 GPU로 map_location을 지정하여 직접 load 해 주면 문제가 해결된다.

 

# 수정된 코드
model = get_model()

checkpoint_path = r"~~~"
checkpoint = torch.load(checkpoint_path, map_location='cuda:{}'.format(args.local_ranl))
model.load_state_dict(checkpoint)

model = torch.nn.parallel.DistributedDataParallel(~~~)