[PYTHON] GPU 메모리 부족(OOM) 해결을 위한 7가지 전략과 성능 최적화 방법

GPU 메모리 부족(OOM) 해결

 

 

딥러닝 모델을 학습시키다 보면 누구나 마주하게 되는 가장 절망적인 순간이 있습니다. 바로 "RuntimeError: CUDA out of memory (OOM)" 에러입니다. 강력한 GPU 자원을 보유하고 있더라도 거대해지는 모델 파라미터와 고해상도 데이터를 처리하다 보면 비디오 메모리(VRAM)의 한계는 금방 찾아옵니다. 본 가이드에서는 파이썬 환경에서 GPU 메모리 점유 구조를 분석하고, 실무 개발자가 즉시 적용하여 학습 중단 문제를 해결할 수 있는 7가지 핵심 기법을 상세히 다룹니다.

---

1. GPU 메모리 부족(OOM)의 근본 원인 분석

OOM 에러는 단순히 데이터가 커서 발생하는 경우보다, 메모리 할당과 해제의 비효율성 때문에 발생하는 경우가 많습니다. 특히 파이토치(PyTorch)나 텐서플로우(TensorFlow)는 속도 향상을 위해 메모리 캐싱 전략을 사용하는데, 이 과정에서 '파편화(Fragmentation)'가 발생하여 가용 메모리가 있음에도 할당에 실패하게 됩니다.

주요 원인 체크리스트

• Batch Size 과다: 한 번에 너무 많은 데이터를 GPU에 올리려고 할 때.
• 모델 복잡도: 레이어가 너무 깊거나 파라미터 수가 GPU 용량을 초과할 때.
• 연산 그래프 누적: 필요 없는 기울기(Gradient) 정보가 메모리에 계속 남아 있을 때.
• 캐시 파편화: 메모리 할당과 해제가 반복되며 연속된 빈 공간이 부족해질 때.

---

2. CPU 메모리 vs GPU 메모리 관리 차이점

파이썬의 일반적인 메모리 관리와 딥러닝 프레임워크의 GPU 메모리 관리는 구조부터 다릅니다.

항목	CPU 시스템 메모리 (RAM)	GPU 비디오 메모리 (VRAM)
주요 목적	범용 데이터 처리 및 프로세스 관리	대규모 병렬 연산 및 텐서 저장
에러 발생 시	Swap 영역(가상 메모리) 활용 가능	즉시 OOM 발생 및 프로세스 중단
할당 방식	운영체제가 동적으로 관리	프레임워크 내부 캐시 알고리즘 사용
접근 속도	상대적으로 느림	매우 빠름 (대역폭이 높음)
해결 방법	Garbage Collection 자동 수행	수동 캐시 비우기 및 그래프 최적화 필요

---

3. GPU OOM 해결을 위한 7가지 실무 Python 코드 예제

성능 저하를 최소화하면서 메모리 부족 문제를 해결하는 핵심 코드 패턴들입니다.

Example 1: Gradient Accumulation (경사 누적) 방법

실제 배치 사이즈를 줄이되, 가중치 업데이트 주기를 조절하여 큰 배치 사이즈를 사용하는 것과 동일한 효과를 냅니다.

# PyTorch 예시
accumulation_steps = 4  # 4번의 미니 배치를 모아서 업데이트
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps # 손실 스케일링
    loss.backward()

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

Example 2: 정밀도 조정 (Mixed Precision Training)

FP32(32비트) 대신 FP16(16비트) 정밀도를 사용하여 메모리 점유율을 절반으로 줄이는 방법입니다.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, labels in dataloader:
    with autocast(): # 16비트 연산 적용
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()
    

Example 3: 주기적인 CUDA 캐시 정리 및 메모리 해제

파이썬 객체는 삭제되어도 GPU 캐시에 남아 있는 경우가 많습니다. 이를 강제로 비워줍니다.

import torch
import gc

# 모델 학습 루프가 끝난 후 혹은 에러 발생 직후 사용
def clear_gpu_memory():
    torch.cuda.empty_cache() # 캐시 비우기
    gc.collect()             # 파이썬 가비지 컬렉터 가동
    print("GPU 메모리 캐시 정리가 완료되었습니다.")

# 사용 예
# try: train() except RuntimeError: clear_gpu_memory()
    

Example 4: Gradient Checkpointing (체크포인팅) 전략

중간 레이어의 활성화 값을 저장하지 않고 재계산하여 메모리를 획기적으로 아끼는 방법입니다.

from torch.utils.checkpoint import checkpoint

# 모델 내부 forward 정의 시 사용
def forward(self, x):
    # 레이어 연산을 체크포인트로 감싸기
    x = checkpoint(self.very_heavy_layer, x)
    return x
    

Example 5: 추론 시 'no_grad' 모드 강제 적용

검증(Validation)이나 테스트 단계에서 기울기 계산 그래프가 생성되는 것을 차단하여 메모리를 보호합니다.

with torch.no_grad():
    for val_inputs, val_labels in val_loader:
        val_outputs = model(val_inputs)
        # 기울기를 계산하지 않으므로 메모리 점유가 매우 낮음
    

Example 6: 대용량 텐서의 CPU 오프로딩 (Offloading)

학습에 즉시 필요 없는 파라미터나 텐서를 잠시 시스템 메모리(RAM)로 옮기는 해결책입니다.

# 연산 시에만 GPU로 옮기고 다시 CPU로 반환
tensor_cpu = heavy_tensor.to('cpu')

# 필요할 때만 다시 GPU로
# res = model(tensor_cpu.to('cuda'))
    

Example 7: 배치 사이즈 동적 조절 및 에러 캡처

OOM 발생 시 자동으로 배치 사이즈를 줄여서 재시도하는 견고한(Robust) 코드 구조입니다.

batch_size = 64
while True:
    try:
        train_with_batch_size(batch_size)
        break
    except RuntimeError as e:
        if 'out of memory' in str(e):
            batch_size //= 2
            clear_gpu_memory()
            print(f"OOM 발생! 배치 사이즈를 {batch_size}로 줄여 재시도합니다.")
        else:
            raise e
    

---

4. 차이점 해결: 왜 내 GPU는 메모리 반환이 안 될까?

많은 개발자들이 del model을 호출해도 GPU 점유율이 떨어지지 않는 문제로 고통받습니다. 이는 파이썬의 참조 계수(Reference Counting) 때문입니다. 변수가 다른 곳(예: 리스트, 로그 객체 등)에 여전히 참조되고 있다면 GPU 메모리는 해제되지 않습니다. 따라서 로그를 기록할 때 반드시 loss.item()과 같이 스칼라 값으로 변환하여 텐서 전체가 메모리에 묶이지 않도록 해야 합니다.

---

5. 결론 및 요약

GPU 메모리 부족 문제는 단순한 하드웨어 한계가 아니라 소프트웨어적 최적화의 대상입니다. Gradient Accumulation으로 논리적 배치를 유지하고, Mixed Precision으로 데이터 크기를 최적화하며, Checkpointing으로 계산과 메모리를 트레이드오프하는 전략을 사용한다면 한정된 자원에서도 거대 모델을 성공적으로 학습시킬 수 있습니다.

 

참고 문헌 및 출처:

• PyTorch Documentation: Memory Management (pytorch.org)
• NVIDIA Apex: Mixed Precision Training Guide.
• Chen, T., et al. (2016). Training Deep Nets with Sublinear Memory Cost.
• TensorFlow Guide: GPU memory growth and allocation.