Colab에서 CUDA out of memory 뜰 때: 배치사이즈·gradient accumulation·AMP로 해결하기

by

Colab에서 CUDA out of memory 뜰 때: 배치사이즈·gradient accumulation·AMP로 해결하기

Colab에서 모델을 학습시키다가 RuntimeError: CUDA out of memory를 마주치면 진행 중이던 학습이 그대로 멈춰버립니다. 특히 무료 티어의 T4(16GB)나 Colab Pro의 A100 일부 사양에서도 배치 사이즈를 조금만 키우면 바로 터지는 경우가 많습니다. 이 글에서는 배치사이즈 조정, gradient accumulation, AMP(자동 혼합 정밀도) 세 가지를 실제 코드와 함께 조합하는 방법을 정리합니다. 단순히 “배치를 줄이세요”에서 끝나지 않고, 왜 그렇게 해야 하는지와 각 방법의 부작용까지 다룹니다.

왜 Colab GPU에서 유독 메모리가 부족한가

CUDA out of memory 오류는 GPU VRAM에 모델 가중치, 옵티마이저 상태, 그래디언트, 그리고 순전파 중 저장되는 activation이 모두 함께 올라가기 때문에 발생합니다. 모델 파라미터 자체는 몇 GB밖에 안 되더라도, Adam 옵티마이저는 파라미터마다 모멘텀과 분산 값을 추가로 저장하기 때문에 실제 필요 메모리는 파라미터 크기의 3~4배까지 뛸 수 있습니다.

Colab은 특히 세션마다 할당되는 GPU가 다르고(T4, L4, A100 등), 무료 티어에서는 세션이 다른 사용자와 공유되는 경우도 있어 동일한 코드라도 어제는 되던 배치 사이즈가 오늘은 안 되는 일이 흔합니다. Colab의 리소스 제한 정책은 Google Colab 공식 FAQ에 명시되어 있으며, 무료 티어는 GPU 가용성과 사용 한도가 유동적이라고 안내하고 있습니다. 이 때문에 “고정된 배치 사이즈”보다 “메모리에 맞춰 유동적으로 대응하는 코드”를 짜두는 것이 재발을 막는 핵심입니다.

1단계: 배치사이즈부터 절반씩 줄여보기

가장 먼저 시도할 조치는 배치 사이즈를 절반으로 줄이는 것입니다. 배치 사이즈가 32에서 16, 8로 줄어들면 activation 메모리도 거의 비례해서 줄어듭니다. 다만 무작정 줄이면 학습 안정성과 수렴 속도가 떨어질 수 있으므로, 아래처럼 OOM이 나면 자동으로 배치를 줄여 재시도하는 방식을 써두면 편합니다.

import torch

def train_step(model, batch, optimizer, criterion):
    optimizer.zero_grad()
    output = model(batch["input"])
    loss = criterion(output, batch["target"])
    loss.backward()
    optimizer.step()
    return loss.item()

batch_size = 32
while batch_size >= 1:
    try:
        batch = get_batch(batch_size)  # 사용자 정의 데이터 로더 함수
        loss = train_step(model, batch, optimizer, criterion)
        print(f"batch_size={batch_size} 성공, loss={loss:.4f}")
        break
    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        batch_size //= 2
        print(f"OOM 발생, batch_size={batch_size}로 재시도")

GPU 서버와 그래픽카드가 설치된 데이터센터 랙

여기서 torch.cuda.empty_cache()는 GPU 메모리를 즉시 회복시켜주진 않지만, PyTorch가 내부적으로 캐싱해둔 미사용 메모리 블록을 반환해 다음 할당 시 단편화 문제를 줄여줍니다. 배치 사이즈만으로 문제가 해결되지 않을 때는 다음 단계인 gradient accumulation으로 넘어갑니다.

2단계: gradient accumulation으로 실질 배치 유지하기

배치 사이즈를 줄이면 메모리는 아끼지만 학습 결과(특히 배치정규화, 학습률 스케줄링에 민감한 모델)가 달라질 수 있습니다. 이때 gradient accumulation을 쓰면 작은 배치로 여러 번 순전파·역전파를 수행하고, 옵티마이저 업데이트만 N번마다 한 번 실행해서 “실질적으로는 큰 배치로 학습한 것과 같은 효과”를 냅니다.

accumulation_steps = 4  # 배치 8 x 4 = 실질 배치 32
optimizer.zero_grad()

for step, batch in enumerate(dataloader):
    output = model(batch["input"])
    loss = criterion(output, batch["target"])
    loss = loss / accumulation_steps  # 누적 스케일 보정
    loss.backward()

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

핵심은 loss를 accumulation_steps로 나눈 뒤 backward()를 호출하는 부분입니다. 이렇게 하면 누적된 그래디언트의 평균이 원래 큰 배치로 한 번에 계산한 것과 동일한 스케일을 가지게 됩니다. 이 방식은 GPU 메모리 사용량은 작은 배치 수준(8)으로 유지하면서, 학습 안정성은 큰 배치(32) 수준을 유지할 수 있다는 점에서 배치 사이즈만 줄이는 것보다 우수한 선택입니다.

3단계: AMP(자동 혼합 정밀도)로 메모리와 속도 동시에 잡기

AMP는 FP32(32비트) 대신 FP16(16비트) 또는 BF16으로 연산 일부를 수행해서 activation과 그래디언트가 차지하는 메모리를 절반 가까이 줄여주는 기법입니다. PyTorch는 torch.cuda.amp 모듈로 이를 공식 지원하며, 정확한 사용법은 PyTorch 공식 AMP 문서에 정리되어 있습니다. FP16은 표현 범위가 좁아 그래디언트가 0으로 언더플로우될 위험이 있는데, 이를 막기 위해 GradScaler가 손실 값을 스케일업했다가 역전파 후 다시 스케일다운합니다.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
accumulation_steps = 4
optimizer.zero_grad()

for step, batch in enumerate(dataloader):
    with autocast():
        output = model(batch["input"])
        loss = criterion(output, batch["target"])
        loss = loss / accumulation_steps

    scaler.scale(loss).backward()

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

이 코드는 배치사이즈 축소, gradient accumulation, AMP 세 가지를 모두 결합한 형태입니다. autocast() 블록 안에서 순전파를 실행하면 PyTorch가 각 연산에 맞춰 FP16/FP32를 자동으로 선택하고, scaler가 그래디언트 스케일링을 관리합니다. Colab의 T4나 A100처럼 Tensor Core를 지원하는 GPU에서는 AMP를 켜면 메모리 절감뿐 아니라 연산 속도도 눈에 띄게 빨라집니다.

노트북으로 코드를 작성하며 데이터를 분석하는 개발자

세 가지 방법 비교와 상황별 선택 기준

세 방법은 서로 배타적이지 않고 조합해서 쓰는 것이 정석입니다. 아래 표로 각 방법의 특성을 비교합니다.

방법 메모리 절감 효과 학습 속도 영향 정확도 영향 적용 난이도
배치사이즈 축소 큼 (배치에 비례) 스텝당 빨라지나 전체 시간 증가 가능 배치정규화 등에서 변화 가능 매우 쉬움
Gradient Accumulation 없음(추가 절감 아님, 축소된 배치 유지 수단) 스텝 수 증가로 다소 느려짐 원래 배치와 거의 동일 쉬움
AMP (FP16/BF16) 큼 (activation 절반 수준) Tensor Core 활용 시 더 빨라짐 대부분 무시할 수준, 드물게 언더플로우 중간

실무에서는 먼저 AMP를 켜서 메모리 여유를 확보하고, 그래도 부족하면 배치를 줄이되 gradient accumulation으로 실질 배치를 보정하는 순서를 권장합니다. 세 가지를 모두 적용했는데도 OOM이 반복된다면 모델 구조 자체가 해당 GPU 메모리로는 불가능한 크기일 수 있으므로, activation checkpointing(torch.utils.checkpoint)이나 LoRA 같은 파라미터 효율적 미세조정 기법을 검토할 단계입니다.

자주 묻는 질문(FAQ)

Q1. torch.cuda.empty_cache()를 호출하면 OOM이 완전히 해결되나요?
아닙니다. 이 함수는 PyTorch가 미리 확보해둔 캐시 메모리를 OS에 반환할 뿐, 실제로 사용 중인 모델·옵티마이저·activation 메모리는 줄여주지 않습니다. 근본적인 해결은 배치사이즈 축소, gradient accumulation, AMP 적용입니다.

Q2. gradient accumulation을 쓰면 학습 시간이 그만큼 늘어나나요?
스텝 수가 늘어나는 만큼 순전파·역전파 횟수는 늘지만, 배치 사이즈가 작아진 만큼 GPU 활용률에 여유가 생겨 전체 wall-clock 시간 증가 폭은 생각보다 크지 않은 경우가 많습니다. 다만 옵티마이저 업데이트 빈도가 줄어들므로 학습률 스케줄러를 스텝 기준이 아닌 실제 업데이트 기준으로 맞춰야 합니다.

Q3. Colab 무료 티어에서 AMP를 써도 안전한가요?
T4 GPU는 Tensor Core를 지원하므로 AMP를 안전하게 사용할 수 있습니다. 다만 손실 값이 NaN으로 튀는 경우가 드물게 있으므로, GradScaler를 반드시 함께 사용하고 학습 로그에서 loss가 비정상적으로 커지거나 NaN이 되는지 주기적으로 확인하는 것이 좋습니다.

CUDA out of memory는 한 가지 원인이 아니라 배치 크기, 옵티마이저 메모리, activation 저장 방식이 겹쳐서 나는 문제입니다. 오늘 정리한 코드를 그대로 Colab 노트북에 붙여넣고, AMP → gradient accumulation → 배치 축소 순서로 하나씩 켜보면서 자신의 모델과 GPU 사양에 맞는 조합을 찾아보시기 바랍니다.

Leave a Comment