2026-01-28 | 7 min read | 150 | PyTorch multi-gpu

9-4. PyTorch 고급: 멀티 GPU 학습

1. 데이터 병렬 처리의 기본 개념

딥러닝 모델의 크기가 커지고 데이터셋의 규모가 방대해짐에 따라, 단일 GPU만으로는 학습 시간이 지나치게 오래 걸리는 문제가 발생합니다. 이러한 문제를 해결하기 위해 여러 개의 GPU를 사용하여 학습 속도를 향상시키는 기술이 필요하며, 그중 가장 널리 사용되는 방법이 데이터 병렬 처리(Data Parallelism) 입니다. 데이터 병렬 처리는 대량의 데이터를 여러 GPU에 분산하여 각 GPU가 데이터를 처리하도록 하는 방식입니다.

데이터 병렬 처리의 핵심 아이디어는 모델의 복사본을 각 GPU에 생성하고, 서로 다른 데이터를 각 GPU에 할당하여 병렬적으로 학습을 진행하는 것입니다. 각 GPU는 할당된 데이터로 모델을 학습하고, 학습 과정에서 계산된 기울기(gradient)를 서로 주고받아 모델 파라미터를 동기화합니다. 이렇게 함으로써 전체 학습 시간을 단축하고, 더 큰 규모의 데이터셋과 모델을 효과적으로 학습할 수 있습니다.

데이터 병렬 처리는 다음과 같은 장점을 가집니다.

학습 속도 향상: 여러 GPU를 동시에 사용하여 학습 속도를 대폭 향상시킬 수 있습니다.
대용량 데이터 처리: 단일 GPU에서 처리하기 어려운 대용량 데이터를 처리할 수 있습니다.
메모리 효율성: 모델의 복사본을 여러 GPU에 분산하여 메모리 사용량을 효율적으로 관리할 수 있습니다.

데이터 병렬 처리 개념 설명 뒤

위 그림은 데이터 병렬 처리의 기본적인 구조를 보여줍니다. 여러 개의 GPU가 있고, 각 GPU는 모델의 완전한 복사본을 가지고 있습니다. 입력 데이터는 여러 GPU에 분산되어 각 GPU에서 처리됩니다. 각 GPU는 할당된 데이터에 대해 순전파(forward pass) 및 역전파(backward pass)를 수행하고, 계산된 기울기를 집계하여 모델 파라미터를 업데이트합니다.

2. 데이터 병렬 처리의 작동 원리

데이터 병렬 처리는 크게 세 단계로 나눌 수 있습니다.

데이터 분할(Data Splitting): 입력 데이터를 여러 개의 GPU에 분산하는 단계입니다. 데이터는 일반적으로 미니 배치(mini-batch) 단위로 나뉘어 각 GPU에 할당됩니다. 각 GPU는 할당된 미니 배치를 사용하여 모델을 학습합니다.
순전파 및 역전파(Forward and Backward Propagation): 각 GPU는 할당된 데이터를 사용하여 순전파를 수행하고, 손실 함수를 계산합니다. 그 후 역전파를 통해 기울기를 계산합니다.
기울기 집계 및 파라미터 동기화(Gradient Aggregation and Parameter Synchronization): 각 GPU에서 계산된 기울기는 한 곳으로 모여 집계(aggregation)됩니다. 집계된 기울기는 모든 GPU의 모델 파라미터를 업데이트하는 데 사용됩니다. 이 과정에서 모델 파라미터가 동기화되어, 각 GPU가 동일한 모델을 유지하도록 합니다. 기울기 집계 방식에는 평균(average), 합(sum) 등이 사용됩니다.

수학적으로 살펴보면, 각 GPU에서 계산된 기울기를 $g_i$라고 할 때, 기울기 집계 후의 기울기 $g_{aggregated}$는 다음과 같이 표현될 수 있습니다.

평균: $g_{aggregated} = \frac{1}{N} \sum_{i=1}^{N} g_i$ (N은 GPU의 개수)
합: $g_{aggregated} = \sum_{i=1}^{N} g_i$

PyTorch에서는 이러한 기울기 집계 및 파라미터 동기화를 자동으로 처리해주는 DistributedDataParallel 모듈을 제공합니다.

3. PyTorch를 이용한 멀티 GPU 학습

PyTorch에서는 torch.nn.DataParallel과 torch.nn.parallel.DistributedDataParallel 두 가지 주요 모듈을 통해 멀티 GPU 학습을 지원합니다.

1) `torch.nn.DataParallel`

torch.nn.DataParallel은 비교적 간단하게 사용할 수 있으며, 단일 머신(single machine) 내에서 여러 개의 GPU를 사용하는 경우에 적합합니다. 이 모듈은 모델을 여러 GPU에 복제하고, 데이터를 각 GPU에 분산하여 병렬 학습을 수행합니다.

DataParallel의 작동 방식은 다음과 같습니다.

모델을 DataParallel로 감쌉니다.
입력 데이터를 GPU로 옮깁니다.
모델에 입력을 전달합니다. DataParallel은 자동으로 데이터를 각 GPU에 분산하고, 각 GPU에서 순전파를 수행합니다.
각 GPU에서 계산된 출력을 다시 메인 GPU로 모아서 반환합니다.
역전파를 수행하여 기울기를 계산합니다. DataParallel은 기울기를 모든 GPU에서 집계하고, 모델 파라미터를 업데이트합니다.

import torch
import torch.nn as nn

# 모델 정의
model = nn.Linear(10, 1)

# 여러 개의 GPU 사용 가능 여부 확인
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    # DataParallel로 모델 래핑
    model = nn.DataParallel(model)

# GPU로 모델 이동
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

# 가짜 입력 데이터 생성
input_tensor = torch.randn(20, 10).to(device) # 배치 크기 20

# 순전파
output = model(input_tensor)

DataParallel을 사용하면 코드를 비교적 간단하게 작성할 수 있지만, 단일 머신 내에서만 작동하고, 통신 오버헤드가 발생하여 학습 속도 향상 효과가 제한적일 수 있다는 단점이 있습니다.

2) `torch.nn.parallel.DistributedDataParallel`

torch.nn.parallel.DistributedDataParallel은 여러 머신(multi-machine)에 걸쳐 여러 개의 GPU를 사용하는 경우에 적합하며, 더 높은 수준의 성능을 제공합니다. 이 모듈은 분산 학습을 위한 다양한 통신 백엔드(backend)를 지원하며, 모델 파라미터 동기화를 효율적으로 처리합니다. DistributedDataParallel은 데이터 병렬 처리 외에도, 모델 병렬 처리(model parallelism)를 지원하여 대규모 모델 학습에도 활용될 수 있습니다.

DistributedDataParallel을 사용하기 위해서는 먼저 분산 환경을 초기화해야 합니다. 분산 환경 초기화에는 주로 torch.distributed.init_process_group 함수를 사용하며, 통신 백엔드(예: nccl, gloo)를 지정하고, 각 프로세스의 랭크(rank)와 월드 사이즈(world size)를 설정해야 합니다.

DistributedDataParallel의 작동 방식은 다음과 같습니다.

분산 환경 초기화: 각 프로세스는 고유한 랭크를 할당받고, 통신 그룹을 형성합니다.
모델을 DistributedDataParallel로 감쌉니다.
입력 데이터를 각 프로세스의 GPU로 분산합니다.
각 GPU에서 순전파 및 역전파를 수행합니다.
기울기를 집계하고, 모델 파라미터를 동기화합니다.

import torch
import torch.nn as nn
import torch.distributed as dist
import os

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def main(rank, world_size):
    setup(rank, world_size)

    # 모델 정의
    model = nn.Linear(10, 1)
    model.to(rank)  # 각 프로세스에 할당된 GPU로 모델 이동
    model <mark class="highlight"><strong><u> nn.parallel.DistributedDataParallel(model, device_ids</u></strong></mark>[rank]) # DistributedDataParallel로 래핑

    # 옵티마이저
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

    # 가짜 데이터 생성
    input_tensor = torch.randn(20, 10).to(rank) # 배치 크기 20
    target_tensor = torch.randn(20, 1).to(rank)

    # 순전파, 역전파, 최적화
    for epoch in range(10):
        optimizer.zero_grad()
        output = model(input_tensor)
        loss = nn.MSELoss()(output, target_tensor)
        loss.backward()
        optimizer.step()
        if rank <mark class="highlight"> 0:
            print(f'Epoch {epoch}, Loss: {loss.item()}')

    cleanup()

if __name__ </mark> '__main__':
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(main,
                                args=(world_size,),
                                nprocs=world_size,
                                join=True)

위 코드는 DistributedDataParallel을 사용한 간단한 예제입니다.

setup() 함수는 분산 환경을 초기화합니다. 각 프로세스는 MASTER_ADDR와 MASTER_PORT를 통해 다른 프로세스와 통신하고, nccl 백엔드를 사용하여 GPU 간 통신을 수행합니다.
main() 함수는 각 프로세스에서 실행되는 학습 루프를 정의합니다. 모델을 생성하고, DistributedDataParallel로 감싼 후, 순전파, 역전파, 그리고 파라미터 업데이트를 수행합니다.
torch.multiprocessing.spawn() 함수는 여러 개의 프로세스를 생성하여 각 GPU에서 학습을 병렬적으로 수행합니다.

DistributedDataParallel을 사용하면 DataParallel에 비해 더 높은 성능을 얻을 수 있지만, 분산 환경을 설정하고 관리해야 하는 복잡성이 증가합니다.

4. 멀티 GPU 학습 시 고려사항

1) 배치 크기(Batch Size)

멀티 GPU 학습에서 가장 중요한 고려 사항 중 하나는 배치 크기입니다. 배치 크기는 각 GPU에 할당되는 데이터의 양을 결정하며, 학습 성능에 큰 영향을 미칩니다. 일반적으로, GPU의 개수가 증가하면 배치 크기를 늘려야 합니다.

배치 크기를 너무 작게 설정하면 각 GPU가 처리하는 데이터의 양이 적어 학습 효율이 떨어지고, 기울기의 추정치가 불안정해질 수 있습니다.
배치 크기를 너무 크게 설정하면 GPU 메모리 부족으로 인해 학습이 중단될 수 있습니다.

따라서 적절한 배치 크기를 선택하기 위해서는 GPU 메모리 사용량, 학습 속도, 그리고 모델의 복잡성을 고려해야 합니다.

2) 기울기 집계 방식

DataParallel과 DistributedDataParallel 모두 기울기 집계를 자동으로 처리하지만, 집계 방식(예: 평균, 합)을 이해하는 것은 중요합니다. 기울기 집계 방식은 배치 크기와 학습률(learning rate)에 영향을 미칩니다.

기울기를 평균하는 경우(default): 배치 크기가 증가하더라도 학습률을 변경할 필요가 없습니다.
기울기를 합하는 경우: 배치 크기가 증가하면 학습률을 1/GPU 개수만큼 줄여야 합니다.

3) 통신 오버헤드

멀티 GPU 학습에서는 GPU 간의 통신이 필요하며, 통신 오버헤드는 학습 속도에 영향을 미칠 수 있습니다. DataParallel은 모든 통신을 메인 GPU를 통해 처리하므로, 통신 오버헤드가 더 클 수 있습니다. DistributedDataParallel은 GPU 간에 직접 통신하므로, 통신 오버헤드를 줄일 수 있습니다.

통신 오버헤드를 줄이기 위해서는 다음과 같은 방법을 고려할 수 있습니다.

DistributedDataParallel 사용: DataParallel보다 통신 효율이 좋습니다.
torch.distributed.all_reduce 사용: 기울기 집계 방식을 최적화할 수 있습니다.
GPU 간의 고속 통신 네트워크 사용: InfiniBand와 같은 네트워크는 통신 속도를 향상시킬 수 있습니다.

4) 메모리 관리

멀티 GPU 학습은 GPU 메모리를 많이 사용하므로, 메모리 관리가 중요합니다.

배치 크기를 적절하게 설정하여 GPU 메모리 사용량을 조절합니다.
불필요한 변수를 제거하고, 메모리 사용량을 최소화합니다.
torch.cuda.empty_cache()를 사용하여 사용하지 않는 GPU 메모리를 해제합니다.
梯度累积(Gradient Accumulation): 작은 배치 크기로 여러 번 기울기를 계산하고, 기울기를 누적하여 한 번에 업데이트하는 기법입니다. 이를 통해 GPU 메모리 사용량을 줄이면서, 더 큰 배치 크기의 효과를 낼 수 있습니다.

5) 난수 시드 고정

재현 가능한 학습 결과를 얻기 위해서는 난수 시드를 고정해야 합니다.

import torch
import random
import numpy as np

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)  # GPU를 사용하는 경우
    torch.cuda.manual_seed_all(seed)  # 멀티 GPU를 사용하는 경우
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

위 코드는 난수 시드를 고정하는 함수를 보여줍니다. torch.cuda.manual_seed와 torch.cuda.manual_seed_all은 GPU에서 사용되는 난수 시드를 고정하는 데 사용됩니다. torch.backends.cudnn.deterministic True는 cuDNN 라이브러리의 동작을 결정적으로 만들어 재현성을 확보하고, torch.backends.cudnn.benchmark False는 cuDNN의 자동 튜닝 기능을 비활성화합니다.

5. 결론

멀티 GPU 학습은 딥러닝 모델의 학습 속도를 획기적으로 향상시키는 중요한 기술입니다. PyTorch는 DataParallel과 DistributedDataParallel 두 가지 모듈을 제공하여 멀티 GPU 학습을 쉽게 구현할 수 있도록 지원합니다.

DataParallel은 단일 머신 내에서 여러 개의 GPU를 사용하는 경우에 적합하며, DistributedDataParallel은 여러 머신에 걸쳐 여러 개의 GPU를 사용하는 경우에 적합합니다. DistributedDataParallel은 더 높은 성능을 제공하지만, 분산 환경을 설정하고 관리해야 하는 복잡성이 있습니다.

멀티 GPU 학습 시에는 배치 크기, 기울기 집계 방식, 통신 오버헤드, 메모리 관리 등을 고려해야 합니다. 적절한 설정과 관리를 통해 멀티 GPU 학습의 효과를 극대화하고, 딥러닝 모델의 학습 효율성을 향상시킬 수 있습니다.