Multi-GPU 분산 학습 완전 가이드: DDP, FSDP, DeepSpeed

# NCCL 디버깅 로그 활성화
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL

# 통신 인터페이스 지정 (Multi-Node 환경에서 중요)
export NCCL_SOCKET_IFNAME=eth0

# InfiniBand 사용 비활성화 (필요시)
export NCCL_IB_DISABLE=1

# P2P(Peer-to-Peer) 통신 레벨 설정
export NCCL_P2P_LEVEL=NVL  # NVLink 사용

GPU 0: Model Copy + Data Shard 0 → Gradient 0 ─┐
GPU 1: Model Copy + Data Shard 1 → Gradient 1 ─┤── AllReduce ──→ Averaged Gradient
GPU 2: Model Copy + Data Shard 2 → Gradient 2 ─┤
GPU 3: Model Copy + Data Shard 3 → Gradient 3 ─┘

GPU 0: Layer의 Weight 행렬 상반부  ─┐
                                     ├── 결합 → Layer Output
GPU 1: Layer의 Weight 행렬 하반부  ─┘

GPU 0: Layers 1-8   │ Micro-batch 1 → │ Micro-batch 2 → │ ...
GPU 1: Layers 9-16  │                  │ Micro-batch 1 → │ Micro-batch 2 → │ ...
GPU 2: Layers 17-24 │                  │                  │ Micro-batch 1 → │ ...
GPU 3: Layers 25-32 │                  │                  │                  │ Micro-batch 1 → │

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

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

def cleanup():
    dist.destroy_process_group()

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

    # 모델 생성 및 DDP 래핑
    model = nn.Linear(1024, 512).to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    # Dataset 및 DataLoader (DistributedSampler 사용 필수)
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    optimizer = optim.Adam(ddp_model.parameters(), lr=1e-3)
    loss_fn = nn.MSELoss()

    for epoch in range(10):
        sampler.set_epoch(epoch)  # 매 epoch마다 호출하여 데이터 셔플 보장
        for batch in dataloader:
            inputs, targets = batch
            inputs = inputs.to(rank)
            targets = targets.to(rank)

            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = loss_fn(outputs, targets)
            loss.backward()  # 여기서 자동으로 gradient AllReduce 수행
            optimizer.step()

    cleanup()

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

accumulation_steps = 4

for i, batch in enumerate(dataloader):
    inputs, targets = batch[0].to(rank), batch[1].to(rank)

    # 마지막 accumulation step에서만 gradient 동기화
    context = ddp_model.no_sync if (i + 1) % accumulation_steps != 0 else nullcontext
    with context():
        outputs = ddp_model(inputs)
        loss = loss_fn(outputs, targets) / accumulation_steps
        loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 단일 노드, 4 GPU
torchrun --nproc_per_node=4 train.py

# 멀티 노드 (2 노드, 각 8 GPU)
# 노드 0에서:
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
    --master_addr="192.168.1.1" --master_port=29500 train.py
# 노드 1에서:
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=1 \
    --master_addr="192.168.1.1" --master_port=29500 train.py

import torch
import torch.distributed as dist
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

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

    # Transformer 모델 생성
    model = Transformer(model_args).to(rank)

    # Mixed Precision 설정
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,    # 파라미터를 bfloat16으로 저장
        reduce_dtype=torch.float32,     # gradient reduce는 float32로 수행
    )

    # 각 Transformer 레이어에 FSDP 적용 (레이어 단위 sharding)
    for layer in model.layers:
        fully_shard(layer, mp_policy=mp_policy)

    # 최상위 모델에도 FSDP 적용
    fully_shard(model, mp_policy=mp_policy)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    for epoch in range(num_epochs):
        for batch in dataloader:
            inputs = batch["input_ids"].to(rank)
            labels = batch["labels"].to(rank)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = loss_fn(outputs, labels)
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

{
  "zero_optimization": {
    "stage": 1,
    "reduce_bucket_size": 5e8
  }
}

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  }
}

{
  "zero_optimization": {
    "stage": 3,
    "contiguous_gradients": true,
    "stage3_max_live_parameters": 1e9,
    "stage3_max_reuse_distance": 1e9,
    "stage3_prefetch_bucket_size": 1e7,
    "stage3_param_persistence_threshold": 1e5,
    "reduce_bucket_size": 1e7,
    "sub_group_size": 1e9,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "offload_param": {
      "device": "cpu",
      "pin_memory": true
    }
  }
}

{
  "train_batch_size": 64,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "zero_optimization": {
    "stage": 2,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "reduce_bucket_size": 2e8,
    "allgather_bucket_size": 2e8
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": 1e-4,
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": 0.01
    }
  },
  "scheduler": {
    "type": "WarmupDecayLR",
    "params": {
      "warmup_min_lr": 0,
      "warmup_max_lr": 1e-4,
      "warmup_num_steps": 1000,
      "total_num_steps": 50000
    }
  }
}

from accelerate import Accelerator

accelerator = Accelerator()

# 기존 코드에서 변경하는 부분은 이것뿐
model, optimizer, dataloader, scheduler = accelerator.prepare(
    model, optimizer, dataloader, scheduler
)

for batch in dataloader:
    optimizer.zero_grad()
    outputs = model(batch["input_ids"])
    loss = loss_fn(outputs, batch["labels"])
    accelerator.backward(loss)  # loss.backward() 대신
    optimizer.step()
    scheduler.step()

$ accelerate config

# 질문에 답하면 자동으로 설정 파일이 생성된다
# - 분산 학습 유형 (multi-GPU, multi-node, TPU 등)
# - GPU 개수
# - Mixed precision 사용 여부
# - DeepSpeed / FSDP 사용 여부 및 설정

compute_environment: LOCAL_MACHINE
distributed_type: MULTI_GPU
num_machines: 1
num_processes: 4
mixed_precision: bf16
use_cpu: false

compute_environment: LOCAL_MACHINE
distributed_type: DEEPSPEED
deepspeed_config:
  zero_stage: 2
  gradient_accumulation_steps: 4
  offload_optimizer_device: none
  offload_param_device: none
mixed_precision: bf16
num_machines: 1
num_processes: 8

compute_environment: LOCAL_MACHINE
distributed_type: FSDP
fsdp_config:
  fsdp_sharding_strategy: FULL_SHARD
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch: BACKWARD_PRE
  fsdp_state_dict_type: SHARDED_STATE_DICT
mixed_precision: bf16
num_machines: 1
num_processes: 8

accelerate launch train.py

# 기본 GPU 상태 확인
nvidia-smi

# 1초 간격 자동 갱신 모니터링
watch -n 1 nvidia-smi

# 특정 GPU만 모니터링
nvidia-smi --id=0,1

# 프로세스별 GPU 사용량 모니터링
nvidia-smi pmon -i 0 -s um -d 1

# 연속 디바이스 모니터링 (1초 간격)
nvidia-smi dmon -d 1

# CSV 형태로 주요 지표 출력 (스크립트에서 파싱하기 용이)
nvidia-smi --query-gpu=index,name,temperature.gpu,utilization.gpu,utilization.memory,memory.used,memory.total,power.draw --format=csv,noheader,nounits

# 출력 예시:
# 0, NVIDIA A100-SXM4-80GB, 45, 98, 72, 65536, 81920, 285
# 1, NVIDIA A100-SXM4-80GB, 43, 95, 68, 61440, 81920, 278

dataloader = DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,           # CPU 코어 수에 맞게 조정
    pin_memory=True,          # GPU 전송 속도 향상
    prefetch_factor=2,        # 미리 로드할 batch 수
    persistent_workers=True,  # worker를 epoch 간 유지
)

import torch

# 현재 GPU 메모리 사용량 확인
print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"Cached: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
print(f"Max Allocated: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB")

# 메모리 snapshot 저장 (상세 분석용)
torch.cuda.memory._record_memory_history()
# ... 학습 코드 실행 ...
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# 1단계: Batch size 줄이기
batch_size = 16  # → 8, 4, 2로 단계적으로 줄여본다

# 2단계: Gradient accumulation으로 effective batch size 유지
gradient_accumulation_steps = 4  # batch_size * 4 = effective batch size

# 3단계: Mixed precision 사용
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

with autocast(dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 4단계: Gradient checkpointing 활성화
from torch.utils.checkpoint import checkpoint
# 또는 HuggingFace 모델:
model.gradient_checkpointing_enable()

# 5단계: 메모리 할당 설정 최적화
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

Batch size 줄이기 → Mixed Precision → Gradient Checkpointing
→ ZeRO Stage 2 → ZeRO Stage 3 → CPU Offloading → NVMe Offloading

# 원인 1: 타임아웃 값이 너무 짧음
# 해결: 타임아웃 증가
export NCCL_BLOCKING_WAIT=1

# Python에서 타임아웃 설정
import datetime
dist.init_process_group(
    backend="nccl",
    timeout=datetime.timedelta(minutes=30)  # 기본값은 30분
)

# 원인 2: GPU 간 P2P 통신 문제
# 해결: P2P 비활성화
export NCCL_P2P_DISABLE=1

# 원인 3: 네트워크 인터페이스 선택 오류 (Multi-Node)
# 해결: 올바른 네트워크 인터페이스 지정
export NCCL_SOCKET_IFNAME=eth0
export GLOO_SOCKET_IFNAME=eth0

# 원인 4: Docker 환경에서 Shared Memory 부족
# 해결: Docker 실행 시 SHM 크기 증가
# docker run --shm-size=16g ...

# 디버깅을 위한 NCCL 로그 활성화
export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=ALL
export TORCH_DISTRIBUTED_DEBUG=DETAIL

import torch.autograd.profiler as profiler

with profiler.profile(
    activities=[
        profiler.ProfilerActivity.CPU,
        profiler.ProfilerActivity.CUDA,
    ],
    schedule=profiler.schedule(wait=1, warmup=1, active=3, repeat=1),
    on_trace_ready=profiler.tensorboard_trace_handler('./log/profiler'),
    record_shapes=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(dataloader):
        if step >= 7:
            break
        train_step(model, batch)
        prof.step()

# 1. Gradient compression 사용 (DDP)
from torch.distributed.algorithms.ddp_comm_hooks import (
    default_hooks as default,
    powerSGD_hook as powerSGD,
)
ddp_model.register_comm_hook(
    state=powerSGD.PowerSGDState(process_group=None, matrix_approximation_rank=1),
    hook=powerSGD.powerSGD_hook,
)

# 2. Bucket size 조정 (DDP)
ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=50)  # 기본값 25MB

# 3. overlap_comm 활성화 (DeepSpeed)
# ds_config.json에서:
# "zero_optimization": { "overlap_comm": true }

모델이 단일 GPU에 들어가는가?
├── YES → DDP (가장 빠르고 간단)
│         ├── Batch size를 키우고 싶다면 → ZeRO Stage 1/2 or FSDP SHARD_GRAD_OP
│         └── 충분하다면 → DDP만으로 충분
└── NO → 모델이 전체 GPU 메모리 합에 들어가는가?
          ├── YES → FSDP (FULL_SHARD) 또는 ZeRO Stage 3
          │         ├── PyTorch 네이티브 선호 → FSDP
          │         └── 더 많은 옵션/커스터마이징 → DeepSpeed
          └── NO → ZeRO Stage 3 + CPU/NVMe Offloading
                    또는 Pipeline Parallelism + Tensor Parallelism 조합