GPU 분산 학습 - DDP (Distributed Data Parallel)

DDP를 알기 전에, 먼저 DP에 대해서 알아야 한다.

일단 이름에서 알 수 있듯이, 뭐 병렬적인 데이터 처리? 이런 느낌인 것 같다.

 

DP (Data Parallel)

한 프로세스에 대해서 여러 GPU를 사용하는 것.

프로세스 = 실행 중인 프로그램 하나 (e.g. 웹 브라우저)
쓰레드 = 프로세스 내에는 여러 개의 쓰레드가 존재, 같은 메모리를 공유 중 (e.g. UI 생성, 네트워크 요청, 자바 스크립트 등등)

 

큰 배치를 GPU 개수만큼 조각내서 각각의 GPU에 넘기고, 각 GPU가 순전파, 역전파를 수행한 뒤 다시 모아서 파라미터를 업데이트

 

나루토 식당 예시.

나루토 마을에 웨이터(메인 프로세스)가 손님인 나루토(batch)를 데려온다.

손님 접시를 4등분 실행 ------- (GPU0에서) batch 수를 4등분하고, 나머지 GPU에게 전송

접시마다 포크와 숟가락을 세팅 -------- GPU마다 GPU0 모델의 파라미터를 복사해놓음

손님이 분신술 써서 각 테이블에서 식사 -------- 각 GPU마다 Forwarding 수행

빈 접시를 다시 모아서 셰프에게 갖다주고, 셰프는 접시를 보고 요리법을 생각 -------- GPU0에서 loss를 계산

그리고 영수증을 다시 분신들에게 각각 나눠주고, 분신들은 자신이 요리를 맛있게 먹었나 생각하고 식사 매너를 수정함 ------ 각 GPU마다 loss를 알려주고 Backwarding 수행 & GPU0에게 gradient 전송 & GPU0에서 parameter 업데이트 후 다른 GPU들에게 업데이트된 parameter 복사

 

import torch
import torch.nn as nn

model = MyModel()               # CPU 메모리에 생성
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)  # 한 프로세스에서 다중 GPU 활용
model = model.cuda()            # 모델 파라미터를 GPU VRAM에 복사

# 학습 루프에서
for xb, yb in loader:
    xb, yb = xb.cuda(), yb.cuda()   # 입력도 한 번에 GPU VRAM에 복사
    optimizer.zero_grad()
    preds = model(xb)               # 내부적으로 batch를 각 GPU로 분할
    loss = criterion(preds, yb)
    loss.backward()                 # 각 GPU에서 gradient 계산 후 메인 GPU로 모임
    optimizer.step()                # 메인 GPU에서 파라미터 업데이트

DDP (Distributed Data Parallel)

여러 GPU가 탑재돼있을 때, 이 GPU를 모두 모아 동시에 하나의 모델을 동기적으로 학습하는 것

1 GPU당 1 프로세스, 각 프로세스가 자신의 GPU에서 학습이나 데이터를 처리하고, 매 스텝마다 gradient를 서로 교환해서 파라미터를 일치시킨다.

 

DP에는 기존에 GPU0에게 많은 부하가 일어난다는 점, 공유 자원에 접근하는 환경에선 lock 걸리는 일이 많다는 것.. 다양한 문제점이 있었다. 따라서 멀티 thread가 아니라 멀티 process를 사용하자는 목적으로 나타난 DDP.

 

이제 예시를 먼저 설명하고, 코드를 설명하겠다.

나루토 식당 예시. (GPU가 4대가 있다고 가정) 

현재 식당에는 4개의 주방이 존재하는 상태. 4명의 셰프가 존재하는 상태.

 

1) 식당 오픈

- 식당 4개의 주방을 키고, 각 주방마다 당일 번호(rank)를 부여받는다.

- 그리고 "우리는 한 몸이다"라는 의미로 world_size라는 것을 설정해서 서로 연결함 (world_size = 4, 통신 채널 NCCL)

                          = GPU 0 ~ 3에 대응하는 프로세스들이 모여서 하나의 분산 그룹을 구성

 

2) (준비 과정) 레시피 복사 

- 원본 레시피 (모델 파라미터) 복사본을 각 셰프(GPU 0 ~ GPU 3)에게 나눠준다.

- 이제 각 셰프들은 그 레시피대로 그대로 조리하고, 맛 평가 때 의견을 모으라고 안내를 받는다.

 

3) (오픈 후) 주문 분배

- 손님들이 64개의 주문서를 내밀면, 셰프들에게 16장씩 균일하게 배분.

- 각 주방은 매 라운드(batch)마다 다른 주문 묶음을 받게 되는데, 라운드가 바뀔 때마다 주문 순서를 섞어서 공평하게 나눈다.

 

4) 조리 및 맛 평가

- 라운드 1 (batch 1)

각 셰프들은 16개 요리를 만들게 된다 (GPU마다 forwarding하고 prediction)이때 재료를 꺼내고 버너를 켜는 등 요리에 필요한 것들 (메모리 접근, cuDNN 커널 실행)은 각자의 주방 (VRAM)에서 이뤄진다.- 첫 시식 (손실 계산)각 셰프들은 자신이 만든 16개 요리를 자기 기준(CrossEntropyLoss 등)으로 시식하며 음식에 점수를 매김 (loss)

- 레시피 개선

각 셰프들은 자신의 레시피에 대한 수정안(gradient)을 만들게 된다.

- 맛 평가 회의

각 셰프들이 각자 수정안을 담당 매니저(NCCL)에게 들고감

매니저가 4개의 의견을 모아서 평균 및 합산한 뒤, 동일한 수정안을 4명에게 뿌려줌

- 레시피 업데이트

각 셰프는 자기 레시피북에 수정안을 반영함 (weight update)

- 라운드 2 

...

 

 

이제 이것을 코드로 한번 작성해보자.

 

1. 식당 오픈

import torch.multiprocessing as mp

def main():
    world_size = 4  # GPU 개수
    mp.spawn(ddp_worker, args=(world_size,), nprocs=world_size, join=True)
    # 4개의 독립 프로세스를 띄워서 각각 rank=0,1,2,3을 인자로 전달하게 된다.

import torch.distributed as dist
import os

def ddp_worker(rank, world_size):
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(
        backend='nccl',
        rank=rank,
        world_size=world_size
    )
    # 이제 rank 0~3 의 4개 프로세스가 하나의 “그룹”을 이룬다.

 

2. 레시피 복사

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

def ddp_worker(rank, world_size):
    # (앞 단계 생략)
    torch.cuda.set_device(rank)
    device = torch.device(f'cuda:{rank}')

    model = MyModel().to(device)
    model = DDP(model, device_ids=[rank])

 

3. 주문 분배

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

def ddp_worker(rank, world_size):
    # (앞 단계 생략)
    dataset = MyDataset()  # 예: 64개 샘플
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
    loader = DataLoader(
        dataset,
        batch_size=64,
        sampler=sampler,
        num_workers=2,
        pin_memory=True
    )

 

4. 학습

for epoch in range(epochs):
    sampler.set_epoch(epoch)
    for xb, yb in loader:
        # ┌────────────────────────────────────────────────────────┐
        # │ (1) 배치 로드 & GPU 전송                              │
        # └────────────────────────────────────────────────────────┘
        xb = xb.to(device, non_blocking=True)  # CPU DRAM → GPU VRAM
        yb = yb.to(device, non_blocking=True)

        # ┌────────────────────────────────────────────────────────┐
        # │ (2) 순전파 (forward)                                  │
        # └────────────────────────────────────────────────────────┘
        preds = model(xb)  # cuBLAS/cuDNN 커널이 VRAM 상 데이터로 연산

        # ┌────────────────────────────────────────────────────────┐
        # │ (3) 손실 계산                                         │
        # └────────────────────────────────────────────────────────┘
        loss = criterion(preds, yb)

        # ┌────────────────────────────────────────────────────────┐
        # │ (4) 역전파 (backward) + 그래디언트 동기화             │
        # └────────────────────────────────────────────────────────┘
        optimizer.zero_grad()
        loss.backward()
        # DDP가 내부적으로 NCCL all-reduce 호출
        #    각 GPU VRAM에 있는 그래디언트를 서로 합산·평균 처리

        # ┌────────────────────────────────────────────────────────┐
        # │ (5) 파라미터 업데이트                                 │
        # └────────────────────────────────────────────────────────┘
        optimizer.step()