Multi_GPU 학습하기

모델을 나누거나 데이터를 나눠서 병렬적으로 처리를 할 수 있는데 모델의 병렬화는 병목과 파이프라인 설계의 어려움 때문에 섣불리 건드리기는 쉽지 않은 과제다.

Data Parallel

데이터를 나눠서 GPU에 할당 후, 결과의 평균을 취하는 방식

DataParallel
단순히 데이터를 분배 -> 평균
GPU 사용의 불균형, 병목이 발생한다.

DistributedDataParallel
각 CPU마다 process를 생성, 개별 GPU에 할당 -> 개별적으로 연산의 평균을 낸다.

Parallel 모델 생성

단순히 모델은 Encapsulate하면 된다.

torch.nn.DataParallel(model)

DistributedDataParallel

이 경우는 조금 더 복잡하다. 먼저 데이터 로더의 sampler를 설정해야 한다.

sampler = torch.utils.data.distributed.DistributedSampler(train_data)
pin_memory = True(이걸 설정해야지 메모리가 더 빠르게 데이터를 읽을 수 있다.)

loader = torch.utils.data.DataLoader(data, batch_size, shuffle, pin_memory = True, sampler = sampler, num_workers = 3)

그리고 multiprocessing을 이용해 process를 생성해야한다.

def main():
	n_gpus = torch.cuda.device_count()
    torch.multiprocessing.spawn(main_worker, nprocs = n_gpus, args = (n_gpus, ))

이 함수는 gpu 개수만큼 process를 생성해서 각 process에 main_worker라는 함수를 할당해 실행시킨다.

def main_workder(gpu, n_gpus):
	...
    torch.distributed.init_process_group(backend = 'nccl', init_method = 'tcp://127.0.0.1:2568',
    world_size = n_gpus, rand = gpu)
 model = model.cuda(gpu)
 model = torch.nn.parallel.DistributedParallel(model, device_ids = [gpu])

사실 실제로 실행되는 함수는 이 부분으로 gpu와 gpu 개수를 받아 받은 gpu에 모델을 할당하고 모델은 distributed로 encapsulate한다.