[Python, Pytorch] Distributed Data Parallel (DDP) 사용법

안녕하세요, 오늘은 Distributed Data Parallel (DDP) 사용하는 방법에 대해서 알아보겠습니다.

 

Distributed Data Parallel (DDP)는 PyTorch에서 멀티 GPU 학습을 효율적으로 처리하는 병렬 학습 프레임워크입니다.

 

각 GPU가 자신의 모델 복사본을 가지고 데이터를 나누어 학습하므로써, Backward 시에만 Gradient를 서로 통신하여 동기화합니다.

 

다음은, DPP를 활용하는 코드 예시 입니다.

 

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, DistributedSampler

def main():
    # 환경 변수로부터 rank/local_rank/world_size 읽기
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    # DDP 초기화
    dist.init_process_group(backend="nccl")

    # 데이터셋 및 DistributedSampler 설정
    transform = transforms.Compose([transforms.ToTensor()])
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    train_sampler = DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset, batch_size=128, sampler=train_sampler)

    # 모델, loss, optimizer
    model = CNN().to(local_rank)
    model = DDP(model, device_ids=[local_rank])
    criterion = nn.CrossEntropyLoss().to(local_rank)
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    num_epochs = 10
    import time
    start_time = time.time()
    for epoch in range(num_epochs):
        model.train()
        train_sampler.set_epoch(epoch)  # Epoch마다 셔플을 다르게
        total_loss = 0.0
        for images, labels in train_loader:
            images = images.to(local_rank)
            labels = labels.to(local_rank)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        if dist.get_rank() == 0:  # 0번 GPU만 출력
            print(f"[Epoch {epoch+1}] Loss: {total_loss:.4f}")

    print(time.time() - start_time)

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

!torchrun --nproc_per_node=2 ddm.py

 

위의 코드는 GPUs가 2개인 경우에 DDP를 활용한 모델 학습 코드입니다.

해당 코드는 다음과 같은 내용으로 구성되어 있습니다.

 

1. GPU ID : 환경변수 LOCAL_RANK를 통해 각 프로세스가 사용할 GPU를 설정
2. CUDA 디바이스 설정 : torch.cuda.set_device(local_rank)로 GPU 지정
3. DDP 초기화 : dist.init_process_group(backend="nccl")로 GPU 간 통신 그룹 생성
4. 데이터 로딩 : DistributedSampler로 GPU 마다 데이터 분할
5. 종료 : dist.destroy_process_group()으로 통신 정리

다음은 DataParallel과 Distributed Data Parallel (DDP) 학습 결과를 보여줍니다. (GPU - RTX 4090 x2EA 기준)

 

DataParallel 결과

Distributed Data Parallel (DDP)

 

지금까지, PyTorch를 활용해서 Distributed Data Parallel (DDP)를 적용하는 방법에 대해서 알아보았습니다.

감사드립니다.