PyTorch에서 멀티 GPU를 활용하여 모델을 학습시킬 때 가장 흔히 사용하는 도구가 torch.nn.DataParallel(이하 DP)입니다. DP는 단일 노드 내에서 여러 GPU에 데이터를 분산시켜 병렬 처리를 수행하는 직관적인 방법을 제공합니다. 이 글에서는 DP의 초기화 과정과 순전파(Forward Pass) 시 내부적으로 어떤 메커니즘이 작동하는지 소스 코드 레벨에서 분석합니다.

1. DataParallel 초기화 과정

DataParallel 인스턴스를 생성할 때 가장 먼저 호출되는 __init__ 메서드는 가용 리소스를 확인하고 실행 환경을 설정하는 역할을 합니다.

def __init__(self, module, device_ids=None, output_device=None, dim=0):
    super(DataParallel, self).__init__()

    # CUDA 사용 가능 여부 확인
    if not torch.cuda.is_available():
        self.module = module
        self.device_ids = []
        return

    # GPU 리스트 설정: 미지정 시 모든 가용 GPU 사용
    if device_ids is None:
        device_ids = list(range(torch.cuda.device_count()))
    
    # 결과 값을 모을 출력 장치 설정: 미지정 시 device_ids의 첫 번째 장치 사용
    if output_device is None:
        output_device = device_ids[0]

    self.dim = dim
    self.module = module
    self.device_ids = [_get_device_index(d, True) for d in device_ids]
    self.output_device = _get_device_index(output_device, True)
    self.src_device_obj = torch.device("cuda:{}".format(self.device_ids[0]))

    _check_balance(self.device_ids)

    # GPU가 하나만 있을 경우 모델을 해당 장치로 바로 이동
    if len(self.device_ids) == 1:
        self.module.cuda(device_ids[0])

여기서 주의할 점은 output_device입니다. 기본적으로 device_ids 리스트의 0번 인덱스에 해당하는 GPU가 마스터 GPU 역할을 수행합니다. 만약 모델이 cuda:1에 로드되어 있는데 device_ids가 [0, 1, 2]라면, src_device_obj는 cuda:0이 되어 런타임 에러가 발생할 수 있습니다. 즉, 모델의 파라미터가 위치한 장치와 device_ids[0]이 일치해야 합니다.

2. 순전파(Forward) 메커니즘

DP의 핵심 로직은 forward 메서드에 집중되어 있습니다. 데이터를 나누고, 모델을 복제하고, 연산을 수행한 뒤 다시 결과를 합치는 일련의 과정이 포함됩니다.

def forward(self, *inputs, **kwargs):
    if not self.device_ids:
        return self.module(*inputs, **kwargs)

    # 모델의 파라미터와 버퍼가 마스터 GPU(device_ids[0])에 있는지 확인
    for param in chain(self.module.parameters(), self.module.buffers()):
        if param.device != self.src_device_obj:
            raise RuntimeError(f"module must have its parameters and buffers on device {self.src_device_obj}")

    # 1. Scatter: 입력 데이터를 각 GPU로 분산
    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
    
    if len(self.device_ids) == 1:
        return self.module(*inputs[0], **kwargs[0])

    # 2. Replicate: 모델을 각 GPU로 복제
    replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
    
    # 3. Parallel Apply: 각 GPU에서 독립적으로 연산 수행
    outputs = self.parallel_apply(replicas, inputs, kwargs)
    
    # 4. Gather: 각 GPU의 출력값을 마스터 GPU로 수집
    return self.gather(outputs, self.output_device)

Step 1: Scatter (데이터 분산)

scatter 함수는 입력 텐서를 지정된 차원(기본값은 batch dimension인 0)을 기준으로 쪼개어 여러 GPU에 전달할 준비를 합니다. 텐서가 아닌 객체는 참조값만 복사되어 전달됩니다.

def scatter(inputs, target_gpus, dim=0):
    def scatter_map(obj):
        if isinstance(obj, torch.Tensor):
            return Scatter.apply(target_gpus, None, dim, obj)
        if isinstance(obj, tuple) and len(obj) > 0:
            return list(zip(*map(scatter_map, obj)))
        if isinstance(obj, list) and len(obj) > 0:
            return list(map(list, zip(*map(scatter_map, obj))))
        if isinstance(obj, dict) and len(obj) > 0:
            return list(map(type(obj), zip(*map(scatter_map, obj.items()))))
        return [obj for _ in target_gpus]

    try:
        return scatter_map(inputs)
    finally:
        scatter_map = None

Step 2: Replicate (모델 복제)

마스터 GPU에 있는 모델의 가중치를 다른 모든 GPU 장치로 복제합니다. 이 과정은 매 forward 호출 시마다 발생하므로, 모델이 매우 클 경우 오버헤드의 원인이 되기도 합니다.

Step 3: Parallel Apply (병렬 실행)

각 GPU에 분산된 데이터와 복제된 모델을 매핑하여 실제 연산을 수행합니다. 이 단계는 멀티 스레딩을 통해 각 장치에서 동시에 실행됩니다.

Step 4: Gather (결과 취합)

각 GPU에서 연산된 결과(Logits 등)를 다시 output_device(마스터 GPU)로 모으는 과정입니다. 흩어져 있던 텐서들이 하나의 배치로 다시 합쳐지게 됩니다.

3. 주요 요약

PyTorch DataParallel은 단일 프로세스에서 멀티 스레드를 사용하여 동작합니다. 데이터는 매번 쪼개져서 전달되고(Scatter), 모델 가중치 역시 매번 복사되며(Replicate), 결과값은 다시 한곳으로 모입니다(Gather). 이러한 구조 때문에 마스터 GPU의 메모리 사용량이 다른 GPU보다 상대적으로 높게 나타나는 현상이 발생하며, 연산 장치 간 통신 병목이 학습 속도에 영향을 줄 수 있습니다.