DistributedDataParallel: 다중 GPU 학습

다중 GPU 학습의 표준

DistributedDataParallel(DDP)은 PyTorch의 표준 다중 GPU 학습 방식이야. GPU마다 프로세스 하나를 띄우고, 각 프로세스가 모델 복사본을 하나씩 들고 서로 다른 데이터 조각에서 순전파와 역전파를 실행해. 역전파 중에는 프로세스 사이에서 기울기를 all-reduce로 동기화해. 각 옵티마이저가 같은 평균 기울기를 자신의 가중치에 적용하므로 모든 모델 복사본은 같은 상태를 유지해.

단일 GPU 스크립트에 추가할 것

dist.init_process_group("nccl")로 프로세스 그룹을 초기화해.
torch.cuda.set_device(rank)로 프로세스마다 사용할 장치를 정해.
model = DDP(model, device_ids=[rank])로 모델을 DDP로 감싸.
DistributedSampler를 사용해 프로세스마다 서로 다른 데이터 조각을 보게 해.
기존 torch.distributed.launch 대신 현대적인 실행기인 torchrun으로 실행해.

DistributedSampler의 함정

DistributedSampler가 없으면 모든 프로세스가 전체 데이터셋을 순회해 같은 작업을 중복하므로 계산을 낭비해. 또 에포크가 시작될 때마다 샘플러의 set_epoch(epoch)를 반드시 호출해야 해. 빼먹으면 매 에포크의 섞기 순서가 같아져 학습이 이유 없이 정체될 수 있어.

torchrun: 현대적인 실행기

torchrun --nproc_per_node=4 train.py는 프로세스 네 개를 띄우고 LOCAL_RANK, WORLD_SIZE, RANK 환경 변수를 설정해 각 프로세스가 자신의 순위를 알게 해. 기존 torch.distributed.launch도 작동하지만 새 코드에는 torchrun을 사용해.

Code

DDP 학습 스크립트: 기준 모양·python

import os
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler

def setup():
    dist.init_process_group("nccl")
    rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(rank)
    return rank

def train(model, dataset, num_epochs):
    rank = setup()
    model = model.to(rank)
    model = DDP(model, device_ids=[rank])

    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler,
                        num_workers=4, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    criterion = nn.CrossEntropyLoss()

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)               # MANDATORY for shuffle
        for x, y in loader:
            x, y = x.to(rank), y.to(rank)
            optimizer.zero_grad()
            loss = criterion(model(x), y)
            loss.backward()                    # gradients sync here (all-reduce)
            optimizer.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    model = MyModel()
    dataset = MyDataset()
    train(model, dataset, num_epochs=10)

torchrun으로 실행·python

# Single-node, 4 GPUs
# torchrun --nproc_per_node=4 train.py

# Single-node, 8 GPUs, with extra args to your script
# torchrun --nproc_per_node=8 train.py --lr 1e-4 --batch_size 32

# Multi-node — each node runs its own torchrun
# Node 0:
#   torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
#            --master_addr=NODE0_IP --master_port=29500 train.py
# Node 1:
#   torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 \
#            --master_addr=NODE0_IP --master_port=29500 train.py

import os
print(os.environ['LOCAL_RANK'], os.environ['RANK'], os.environ['WORLD_SIZE'])

0번 랭크 전용 기록: 보편적인 방식·python

import os
import torch.distributed as dist

def is_main():
    return not dist.is_initialized() or dist.get_rank() == 0

def log(*args, **kwargs):
    if is_main():
        print(*args, **kwargs)

# Use throughout your training
log(f"epoch {epoch}: train_loss={train_loss:.4f}")

# Same trick for checkpoint saving — only rank 0 writes
if is_main():
    torch.save(model.module.state_dict(), 'best.pt')   # .module to unwrap DDP

# And for any setup that should happen ONCE, not N-times
if is_main():
    os.makedirs('outputs', exist_ok=True)

Exercise

단일 GPU 학습 스크립트를 DDP로 바꾸고 torchrun --nproc_per_node=N으로 실행해. N은 GPU 수야. 다음을 확인해. (a) is_main()으로 제한해 손실이 N번이 아니라 에포크마다 한 번만 출력되는가, (b) GPU가 두 개 이상일 때 학습이 실제로 빨라지는가, (c) 샘플러의 .set_epoch가 알맞은 위치에 있어 에포크마다 섞기 순서가 달라지는가.