beautifulsoups.dev
Published on

CUDA OOM? PyTorch 메모리 단편화 해결법

Authors

서버에서 학습을 돌리다 보면 CUDA out of memory가 가장 흔한 중단 원인입니다. 그런데 nvidia-smi를 보면 VRAM이 완전히 찬 것도 아닌데 OOM이 나기도 합니다. 이때는 “진짜로 메모리가 부족한 상황”이라기보다, PyTorch CUDA 캐시 할당자(caching allocator) 관점에서 연속된 큰 블록을 더 이상 만들 수 없는 상태, 즉 **메모리 단편화(fragmentation)**가 원인일 때가 많습니다.

이 글에서는 (1) 단편화가 왜 생기는지, (2) 어떻게 확인하는지, (3) 가장 효과적인 해결책(환경 변수, 코드 패턴, 운영 팁)을 우선순위대로 정리합니다.

1) OOM인데 VRAM이 남아 보이는 이유

PyTorch는 CUDA 메모리를 자주 cudaMalloc/cudaFree로 직접 요청/반납하지 않고, 성능을 위해 캐시 풀에 잡아두었다가 재사용합니다. 이 과정에서 다음이 겹치면 단편화가 커집니다.

  • 크기가 제각각인 텐서를 반복적으로 생성/파괴 (가변 길이 배치, 동적 shape, 조건 분기)
  • forward/backward마다 일시적인 버퍼가 생기고 사라짐
  • DDP/AMP/gradient checkpointing 등으로 메모리 패턴이 복잡해짐
  • 긴 실행 시간 동안 다양한 크기의 블록이 풀에 쌓임

결과적으로 “총 여유 메모리는 남아 있지만”, **연속된 큰 덩어리(large contiguous block)**가 부족해 특정 할당 요청이 실패하며 OOM이 발생합니다.

2) 단편화 의심 신호: 에러 메시지와 통계

2.1 에러 메시지에서 힌트 찾기

PyTorch OOM 메시지에는 종종 다음과 같은 문구가 포함됩니다.

  • reserved memory is >> allocated memory
  • If reserved memory is much larger than allocated memory try setting max_split_size_mb

여기서

  • allocated: 실제 텐서가 쓰는 메모리
  • reserved: 캐시 풀에 잡혀 있는 메모리(재사용을 위해 남겨둠)

reserved가 과도하게 크고, 동시에 큰 블록 할당이 실패하면 단편화를 강하게 의심할 수 있습니다.

2.2 torch.cuda.memory_summary()로 상태 확인

아래는 실행 중 단편화를 관찰할 때 가장 유용한 요약입니다.

import torch

print(torch.cuda.memory_summary(device=None, abbreviated=False))

여기서 특히 볼 포인트:

  • Allocated memory vs Reserved memory 차이
  • Active/Inactive split blocks
  • Largest free block이 작게 나오는지(버전에 따라 표시 다름)

2.3 실시간 메모리 지표 찍기

학습 루프에서 특정 스텝에만 OOM이 난다면, 스텝별로 피크를 기록해 원인 배치를 좁히는 게 좋습니다.

import torch

def log_cuda(step: int):
    torch.cuda.synchronize()
    alloc = torch.cuda.memory_allocated() / 1024**2
    reserv = torch.cuda.memory_reserved() / 1024**2
    peak = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[step={step}] alloc={alloc:.1f}MB reserved={reserv:.1f}MB peak={peak:.1f}MB")

# 루프 중간에
# log_cuda(step)

reserved가 계속 증가하고, alloc은 크게 변하지 않는데도 결국 OOM이 나면 단편화/캐시 누적 패턴을 의심합니다.

3) 가장 효과적인 해결책 1: PYTORCH_CUDA_ALLOC_CONF

단편화 대응의 핵심은 블록 분할/병합 정책을 조절하는 것입니다. PyTorch는 환경 변수로 CUDA 캐시 할당자 정책을 튜닝할 수 있습니다.

3.1 max_split_size_mb로 과도한 분할 방지

큰 블록을 잘게 쪼개 쓰다 보면, 나중에 큰 덩어리가 필요할 때 재조합이 어려워집니다. 이를 완화하는 대표 옵션이 max_split_size_mb입니다.

# 예: 128MB 이상 블록은 쪼개지 않도록
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
python train.py
  • 값이 너무 작으면 분할이 잦아 단편화가 늘 수 있음
  • 값이 너무 크면 작은 할당이 비효율적으로 큰 블록을 점유할 수 있음

실무에서는 64 / 128 / 256MB를 후보로 두고, 동일 workload에서 OOM 재현 여부와 throughput을 비교하는 방식이 안전합니다.

3.2 expandable_segments로 “세그먼트 확장” 활용(가능한 경우)

PyTorch 버전에 따라 expandable_segments:True가 도움이 되는 케이스가 있습니다. 이는 세그먼트 운용 방식으로 단편화를 줄이는 데 유리할 수 있습니다.

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

다만 모든 환경/버전에서 만능은 아니므로, OOM 재현 workload로 A/B 테스트를 권장합니다.

3.3 옵션은 조합 가능

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128,expandable_segments:True

4) 해결책 2: “가변 shape”를 줄여 메모리 패턴을 안정화

단편화는 결국 “할당 크기 분포가 넓고, 시간에 따라 달라지는 것”에서 시작합니다. 다음은 코드 레벨에서 체감 효과가 큰 방법들입니다.

4.1 배치의 token 길이/이미지 크기를 버킷팅(bucketing)

NLP에서 시퀀스 길이가 들쭉날쭉하면 매 step마다 activation 텐서 크기가 바뀌고, 캐시 풀에 다양한 크기의 블록이 쌓입니다.

  • 길이별로 버킷을 만들고
  • 비슷한 길이끼리 배치를 구성

이렇게 하면 할당 크기가 안정화되어 단편화가 크게 줄어듭니다.

4.2 torch.compile/그래프 고정화는 “형상 고정”과 궁합

가능한 경우 shape를 고정하거나 제한하면(예: padding을 일정 길이로) 메모리 할당 패턴이 반복되어 캐시 재사용 효율이 올라갑니다.

4.3 일시 텐서 생성 최소화

예를 들어 torch.cat/stack을 루프에서 반복하면 매번 새로운 큰 텐서가 생깁니다. 가능하면 미리 버퍼를 잡아두고 in-place로 채우는 방식이 더 안정적입니다.

5) 해결책 3: 학습 루프에서 “메모리 피크”를 낮추는 실전 팁

단편화만이 아니라, 피크 자체를 낮추면 OOM 확률이 내려갑니다.

5.1 AMP(autocast + GradScaler)

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

scaler = GradScaler()

for batch in loader:
    optimizer.zero_grad(set_to_none=True)
    with autocast():
        loss = model(**batch).loss
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

optimizer.zero_grad(set_to_none=True)는 불필요한 메모리 write를 줄이고, gradient 텐서의 라이프사이클을 단순화하는 데 도움이 됩니다.

5.2 Gradient accumulation으로 배치 크기 쪼개기

OOM이 특정 배치 크기에서만 난다면, 미니배치를 여러 번 누적해 “유효 배치”를 유지할 수 있습니다.

accum = 4
optimizer.zero_grad(set_to_none=True)

for i, batch in enumerate(loader):
    with torch.cuda.amp.autocast():
        loss = model(**batch).loss / accum

    scaler.scale(loss).backward()

    if (i + 1) % accum == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

5.3 Activation checkpointing

메모리를 줄이는 대신 연산을 늘리는 전략입니다. 단편화 이전에 “절대 메모리 부족”을 완화합니다.

6) torch.cuda.empty_cache()는 만능이 아니다

torch.cuda.empty_cache()는 PyTorch 캐시 풀에서 CUDA 드라이버로 메모리를 반환할 수 있지만,

  • 이미 활성(allocated/active)로 잡힌 텐서는 줄지 않음
  • 너무 자주 호출하면 성능이 크게 떨어질 수 있음

그럼에도 다음 상황에서는 유효합니다.

  • 학습/추론을 한 프로세스에서 번갈아 수행하며, 단계 사이에 큰 메모리 구성이 바뀔 때
  • 매우 큰 임시 작업(예: evaluation, big batch inference)을 “잠깐” 수행한 뒤 원래 루프로 돌아갈 때

권장 패턴은 “루프 내부 매 step 호출”이 아니라, 단계 전환 지점에서 제한적으로 호출하는 것입니다.

# 예: epoch 끝, 또는 evaluation 직후
import torch

torch.cuda.empty_cache()

7) DataLoader/입력 파이프라인이 만드는 숨은 메모리 변동

GPU OOM처럼 보이지만, 실제로는 입력 파이프라인이 GPU로 올리는 텐서 크기/타이밍을 흔들어 단편화를 악화시키는 경우가 있습니다.

  • pin_memory=True + 비동기 전송이 겹치며 순간 피크가 생김
  • 가변 크기 입력이 섞여 들어옴
  • prefetch_factor, num_workers에 따라 배치 준비 타이밍이 달라짐

해결 접근:

  • 가변 shape를 버킷팅
  • 전송 타이밍을 단순화(필요 시 non_blocking=False로 비교)
  • worker/prefetch를 줄여 피크를 낮추는 A/B 테스트

8) 컨테이너/EKS 환경에서의 추가 체크(멀티 테넌시)

K8s/EKS에서 여러 파드가 한 GPU를 공유(MIG 포함)하거나 노드에 여러 워크로드가 섞이면, “내 프로세스만의 문제”가 아닐 수 있습니다.

  • 다른 프로세스가 VRAM을 점유해 큰 연속 블록이 불리해짐
  • 파드 재시작/스케줄링으로 실행 조건이 자주 바뀜

이럴 때는 노드/파드 상태부터 안정화하는 것이 선행입니다. 예를 들어 파드가 Pending으로 오래 대기하거나 노드 리소스가 불안정하면 재현/튜닝이 어려워집니다. 관련해서는 EKS Pod Pending 0/XX nodes available 원인별 해결 같은 체크리스트가 도움이 됩니다.

또한 GPU 워크로드는 이미지 풀 실패로 재시작 루프에 빠지면(재기동마다 캐시 패턴도 달라짐) OOM 원인 분석이 더 어려워집니다. 운영 이슈가 겹친다면 EKS ImagePullBackOff 429 Too Many Requests 해결도 함께 점검해 두는 편이 좋습니다.

9) 단편화 해결을 위한 “우선순위” 체크리스트

9.1 가장 먼저 할 것(효과 대비 비용 최고)

  1. OOM 로그에서 reserved >> allocated 여부 확인
  2. torch.cuda.memory_summary()로 split/inactive 블록 상태 확인
  3. PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128부터 A/B 테스트

9.2 그다음(근본 개선)

  1. 가변 shape 버킷팅/패딩 정책 정리
  2. 배치 구성 안정화(동일 길이끼리 묶기)
  3. 일시 텐서 생성을 줄이는 코드 리팩터링

9.3 마지막(운영/성능 트레이드오프)

  1. 제한적 empty_cache()
  2. accumulation/checkpointing/AMP 조합 튜닝
  3. K8s 스케줄링/노드 혼잡도 개선

10) 재현 가능한 최소 예제로 단편화 감 잡기

아래 코드는 “크기가 다른 텐서를 반복 생성”해 단편화를 유도하는 장난감 예시입니다(환경에 따라 OOM이 나지 않을 수 있습니다). 핵심은 reserved가 커지고, 특정 시점에 큰 할당이 실패할 수 있다는 점을 관찰하는 것입니다.

import torch
import random

torch.cuda.init()

def alloc(sz_mb):
    # float16: 2 bytes
    numel = (sz_mb * 1024 * 1024) // 2
    return torch.empty(numel, device="cuda", dtype=torch.float16)

blocks = []
for step in range(1, 2000):
    # 4~256MB 사이를 랜덤하게 할당
    mb = random.choice([4, 8, 16, 32, 64, 96, 128, 192, 256])
    blocks.append(alloc(mb))

    # 랜덤하게 일부 해제(참조 제거)
    if len(blocks) > 50 and step % 3 == 0:
        for _ in range(10):
            blocks.pop(random.randrange(len(blocks)))

    if step % 50 == 0:
        torch.cuda.synchronize()
        a = torch.cuda.memory_allocated() / 1024**2
        r = torch.cuda.memory_reserved() / 1024**2
        print(f"step={step} allocated={a:.1f}MB reserved={r:.1f}MB")

# 마지막에 큰 블록을 요청해 실패하는지 확인
try:
    x = alloc(2048)
    print("allocated 2GB ok")
except RuntimeError as e:
    print("OOM:", e)

위 코드를 실행한 뒤, 다음과 같이 max_split_size_mb를 적용해 비교하면 차이를 체감할 수 있습니다.

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 python frag_demo.py

마무리

PyTorch의 CUDA OOM은 단순히 “VRAM이 부족해서”만 발생하지 않습니다. 특히 reserved가 큰데도 OOM이 난다면, 단편화로 인해 필요한 크기의 연속 블록을 확보하지 못하는 문제일 가능성이 큽니다.

가장 빠른 처방은 PYTORCH_CUDA_ALLOC_CONF(특히 max_split_size_mb)로 할당자 정책을 조절하는 것이고, 가장 확실한 근본 해결은 입력 shape와 배치 구성의 변동성을 줄여 메모리 패턴을 안정화하는 것입니다. 이를 통해 “가끔씩만 터지는” 악질 OOM을 재현 가능하게 만들고, 결국 제거할 수 있습니다.

Discuss on Twitter • View on GitHub