분산 학습

from pyspark.ml.torch.distributor import TorchDistributor

def train_function():
    """각 GPU에서 실행되는 학습 함수"""
    import torch
    import torch.distributed as dist
    from torch.nn.parallel import DistributedDataParallel as DDP

    # 1. 분산 환경 초기화 (TorchDistributor가 자동 설정)
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    device = torch.device(f"cuda:{local_rank}")

    # 2. 모델 정의 및 DDP 래핑
    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])

    # 3. 분산 데이터 로더 설정
    sampler = torch.utils.data.distributed.DistributedSampler(dataset)
    dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=32, sampler=sampler
    )

    # 4. 학습 루프
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 에포크마다 셔플링
        for batch in dataloader:
            inputs, labels = batch[0].to(device), batch[1].to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = torch.nn.functional.cross_entropy(outputs, labels)
            loss.backward()
            optimizer.step()

        if rank == 0:  # 메인 프로세스에서만 로깅
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

    # 5. 모델 저장 (메인 프로세스만)
    if rank == 0:
        torch.save(model.module.state_dict(), "/tmp/model.pt")

    dist.destroy_process_group()

# 4개 GPU에서 분산 학습 실행
distributor = TorchDistributor(
    num_processes=4,       # 총 프로세스(GPU) 수
    local_mode=False,      # False: 여러 노드에 분산 / True: 단일 노드
    use_gpu=True           # GPU 사용
)
result = distributor.run(train_function)

{
    "train_batch_size": 64,
    "gradient_accumulation_steps": 4,
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true
    }
}

from pyspark.ml.torch.distributor import TorchDistributor

def train_with_deepspeed():
    import deepspeed
    import torch

    model = MyLargeModel()

    model_engine, optimizer, _, _ = deepspeed.initialize(
        model=model,
        config="ds_config.json"
    )

    for epoch in range(10):
        for batch in dataloader:
            inputs = batch["input_ids"].to(model_engine.device)
            labels = batch["labels"].to(model_engine.device)

            outputs = model_engine(inputs, labels=labels)
            loss = outputs.loss
            model_engine.backward(loss)
            model_engine.step()

# 8개 GPU에서 DeepSpeed 분산 학습
distributor = TorchDistributor(
    num_processes=8,
    local_mode=False,
    use_gpu=True
)
distributor.run(train_with_deepspeed)

import horovod.torch as hvd
import torch

# Horovod 초기화
hvd.init()
torch.cuda.set_device(hvd.local_rank())

# 모델 및 옵티마이저
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001 * hvd.size())

# Horovod 분산 옵티마이저로 래핑
optimizer = hvd.DistributedOptimizer(
    optimizer,
    named_parameters=model.named_parameters()
)

# 모델 파라미터 동기화 (초기화)
hvd.broadcast_parameters(model.state_dict(), root_rank=0)
hvd.broadcast_optimizer_state(optimizer, root_rank=0)

# 학습 루프
for epoch in range(10):
    for batch in dataloader:
        inputs, labels = batch[0].cuda(), batch[1].cuda()
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = torch.nn.functional.cross_entropy(outputs, labels)
        loss.backward()
        optimizer.step()

# 4노드 x 4GPU = 16 GPU 클러스터
new_cluster:
  spark_version: "15.4.x-gpu-ml-scala2.12"
  driver_node_type_id: "g5.12xlarge"   # A10G x 4
  node_type_id: "g5.12xlarge"          # A10G x 4
  num_workers: 3                        # Worker 3 + Driver 1 = 4노드

def train_function():
    import mlflow
    import torch.distributed as dist

    dist.init_process_group("nccl")
    rank = dist.get_rank()

    # 메인 프로세스에서만 MLflow 로깅
    if rank == 0:
        mlflow.start_run()
        mlflow.log_param("num_gpus", dist.get_world_size())
        mlflow.log_param("batch_size", 32)

    # ... 학습 루프 ...

    if rank == 0:
        mlflow.log_metric("final_loss", loss.item())
        mlflow.pytorch.log_model(model.module, "model")
        mlflow.end_run()

[상황]
- 모델: ResNet-50 (25M 파라미터)
- 데이터: 이미지 10만 장 (총 5GB)
- 배치 크기: 32

[실험 결과]
- GPU 1개 (A10G): 45분
- GPU 4개 (A10G × 4, DDP): 38분  ← 예상: 11분
- GPU 8개 (A10G × 8, 2노드): 42분  ← 오히려 거의 같음!

[단일 GPU 학습 사이클]
Forward Pass (20ms) → Backward Pass (30ms) → Optimizer Step (5ms)
= 총 55ms/step

[분산 GPU 학습 사이클 (4 GPU)]
Forward Pass (20ms) → Backward Pass (30ms) → 🔴 All-Reduce 통신 (40ms) → Optimizer Step (5ms)
= 총 95ms/step

🔴 All-Reduce 통신: 각 GPU의 그래디언트를 동기화하는 시간
- 단일 노드 내 (NVLink): 10~40ms
- 노드 간 (네트워크): 50~200ms

실제 Speedup = N / (1 + 통신_시간/연산_시간 × (N-1))

예시 (N=4, 통신_시간=연산_시간의 80%):
Speedup = 4 / (1 + 0.8 × 3) = 4 / 3.4 = 1.18배
→ GPU 4개를 써도 1.18배밖에 빨라지지 않음!

# 팁 1: Mixed Precision (FP16/BF16) — 통신량 50% 감소
# BF16은 A100/A10G에서 지원, 학습 안정성이 FP16보다 우수
from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()
with autocast(dtype=torch.bfloat16):
    output = model(input)
    loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 팁 2: Gradient Accumulation — 통신 빈도 감소
# 4 step마다 한 번 통신 → All-Reduce 횟수 1/4
accumulation_steps = 4
for i, batch in enumerate(dataloader):
    loss = model(batch) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 팁 3: 데이터 로딩 최적화 — GPU가 놀지 않게
dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=4,          # CPU 워커로 병렬 데이터 로딩
    pin_memory=True,        # GPU 전송 속도 향상
    prefetch_factor=2,      # 미리 2 배치 준비
    persistent_workers=True  # 워커 재사용 (시작 오버헤드 제거)
)