Skip to main content

분산 학습 (Distributed Training)

학습 목표

  • 데이터 병렬(Data Parallel)과 모델 병렬(Model Parallel)의 차이를 이해한다
  • DataParallel(DP)과 DistributedDataParallel(DDP)의 차이를 설명할 수 있다
  • FSDP(Fully Sharded Data Parallel)의 개념을 이해한다
  • 실무에서 분산 학습 방식을 선택하는 기준을 안다

왜 중요한가

대형 딥러닝 모델은 하나의 GPU로 학습하기에 시간과 메모리가 부족합니다. 분산 학습은 여러 GPU에 연산을 분배하여 학습을 선형적으로 가속합니다.

병렬화 전략

전략분배 대상적합 상황
데이터 병렬데이터 (배치 분할)모델이 GPU 한 장에 들어갈 때
모델 병렬모델 (레이어 분할)모델이 너무 커서 GPU 한 장에 안 들어갈 때
FSDP데이터 + 모델 파라미터대형 모델의 메모리 효율적 학습

DataParallel (DP) — 가장 간단

import torch.nn as nn

model = MyModel()
model = nn.DataParallel(model)  # 한 줄 추가
model = model.to('cuda')
한계: GPU 0에 기울기가 모이는 병목(GIL + GPU 0 부하), 멀티노드 미지원. 프로토타입에서만 사용합니다.

DistributedDataParallel (DDP) — 실무 표준

각 GPU가 독립 프로세스로 동작하며, All-Reduce로 기울기를 동기화합니다.
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

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

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

    # 각 GPU에 다른 데이터 배치 할당
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # 셔플링을 위해 필수
        for data, labels in dataloader:
            data, labels = data.to(rank), labels.to(rank)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, labels)
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

# 실행: torchrun --nproc_per_node=4 train.py

FSDP (Fully Sharded Data Parallel)

모델 파라미터, 기울기, 옵티마이저 상태를 GPU 간에 분산(shard)하여 메모리를 절약합니다. 대형 모델 학습에 필수적입니다.
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(
    MyModel(),
    # 각 Transformer 블록 단위로 샤딩
    auto_wrap_policy=transformer_auto_wrap_policy,
)

비교

방법설정 난이도메모리 효율속도멀티노드
DP매우 쉬움낮음보통불가
DDP보통보통높음가능
FSDP높음매우 높음높음가능
단일 노드 멀티 GPU에서 모델이 GPU 메모리에 들어간다면 DDP를 사용하세요. 모델이 너무 크면 FSDP를 사용합니다. LLM Fine-Tuning에서는 FSDP + QLoRA 조합이 일반적입니다. DP는 프로토타이핑 외에는 권장하지 않습니다.

체크리스트

  • 데이터 병렬과 모델 병렬의 차이를 이해한다
  • DDP가 DP보다 권장되는 이유를 설명할 수 있다
  • FSDP가 대형 모델에 필요한 이유를 안다
  • torchrun으로 분산 학습을 시작하는 방법을 안다

다음 문서