GPU 메모리·연산 효율 최적화 팁

GPU 메모리·연산 효율 최적화 팁

포스트 요약: 딥러닝 모델 학습·추론 과정에서 자주 병목이 되는 GPU 자원을 효율적으로 사용하는 방법을 소개합니다. Mixed Precision, Gradient Checkpointing, DataLoader 튜닝 등 실전 팁과 코드 예시를 수록했습니다.

---

1. 서론

딥러닝 모델 규모가 커질수록 GPU 메모리 사용량과 연산량이 주요 제약 요소로 떠오릅니다. 메모리 부족으로 인한 OOM(Out-of-Memory) 에러, 연산 시간 지연은 실험 속도를 크게 떨어뜨립니다. 이 글에서는 대표적인 최적화 기법들과 적용 예시를 다룹니다.

---

2. Mixed Precision 학습

2.1 개념

FP32 대신 FP16 연산을 활용해 메모리 사용량과 대역폭을 절반 이상으로 줄입니다. NVIDIA의 amp 라이브러리를 사용하면 손쉽게 적용할 수 있습니다.

2.2 구현 예시

import torch
from torch.cuda.amp import autocast, GradScaler

model, optimizer = MyModel().cuda(), torch.optim.Adam(...)
scaler = GradScaler()

for data, target in loader:
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

– autocast() 블록 내 연산은 자동으로 FP16/FP32 섞어 실행됩니다.
– GradScaler가 underflow/overflow를 방지하며 스케일링 관리.

---

3. Gradient Checkpointing

3.1 개념

중간 활성화(activation) 값을 저장하지 않고, backward 시 재계산해 메모리를 절약합니다. 트레이드오프로 계산 시간은 늘어납니다.

3.2 구현 예시 (PyTorch)

import torch.utils.checkpoint as cp

class MyModel(nn.Module):
    def forward(self, x):
        x = cp.checkpoint(self.layer1, x)
        x = cp.checkpoint(self.layer2, x)
        return self.classifier(x)

– checkpoint(fn, x) 호출 시 fn(x)의 활성화를 저장하지 않고 backward 단계에 재실행.

---

4. DataLoader 최적화

옵션	설명
num_workers	데이터 전처리 병렬 워커 수 조절. CPU 코어 수 대비 4~8배 권장.
pin_memory=True	GPU로 전송 시 페이지 잠금 메모리 사용. 데이터 로드 속도 개선.
persistent_workers=True	Epoch 간 워커 프로세스 유지. 반복 로딩 시 오버헤드 감소.

loader = DataLoader(
    dataset, batch_size=64, shuffle=True,
    num_workers=8, pin_memory=True,
    persistent_workers=True
)

---

5. 메모리 계산 및 모니터링

5.1 예상 메모리 계산

파라미터 수 \(P\)와 배치 크기 \(B\), 데이터 타입 크기 \(D\) (예: FP32=4바이트)로 대략 메모리 사용량을 예측:

// 메모리 ≈ P × D + B × activation_size

5.2 실시간 모니터링

torch.cuda.reset_peak_memory_stats()
# ... forward/backward ...
peak = torch.cuda.max_memory_allocated() / (1024**2)
print(f"Peak GPU Memory: {peak:.1f} MB")

---

6. 종합 비교 표

기법	메모리 절약	연산 오버헤드	적용 예시
Mixed Precision	40–50%	미미	amp.autocast
Checkpointing	30–40%	50–100%↑	torch.utils.checkpoint
DataLoader 튜닝	간접적	낮음	num_workers, pin_memory

---

7. 결론 및 다음 단계

GPU 자원을 효율화하면 대규모 모델 학습이 가능해지고, 실험 속도가 빨라집니다. 위 기법들을 조합하여 최적의 워크플로우를 구축해 보세요.

1. Mixed Precision + Checkpointing 병행 실험
2. Auto Mixed Precision(AMP) vs Native AMP 성능 비교
3. 데이터 파이프라인 병목 진단 및 개선

참고 문헌

1. NVIDIA Apex Documentation: Mixed Precision Training
2. PyTorch Documentation: torch.utils.checkpoint
3. PyTorch Data Loading and Processing Tutorial