GPU 메모리가 8GB인 환경에서 batch size 256을 그대로 올리면 대부분 OOM(Out of Memory)이 발생합니다. VRAM 8GB로 batch size 256 효과 내기는 실제로는 물리적 배치를 작게 유지하면서 그래디언트를 여러 번 누적해 큰 배치와 동일한 업데이트를 만드는 방식으로 해결하는데, 이 글에서는 PyTorch 기준으로 정확히 동작하는 Gradient Accumulation 코드와 실무에서 자주 발생하는 3가지 실수를 구체적으로 다룹니다. 이론적 설명보다 실제로 돌려본 코드와 흔히 놓치는 함정 위주로 정리했습니다.
Gradient Accumulation의 원리와 VRAM 8GB의 한계
Gradient Accumulation은 한 번에 큰 배치를 forward/backward 하는 대신, 작은 micro-batch 여러 개를 순차적으로 backward만 수행해 그래디언트를 누적하고, 일정 횟수마다 한 번씩 optimizer.step()을 호출하는 방식입니다. 예를 들어 물리적 batch size 32로 8번 누적하면 옵티마이저 입장에서는 batch size 256과 수학적으로 동일한 그래디언트 합을 얻습니다.
VRAM 8GB에서 batch size 256이 안 되는 이유는 activation(중간 연산 결과)이 배치 크기에 비례해서 메모리를 차지하기 때문입니다. 모델 파라미터와 옵티마이저 상태(Adam의 경우 파라미터당 2배 추가 메모리)는 배치 크기와 무관하게 고정이지만, forward 시 저장해두는 activation은 배치가 커질수록 선형으로 늘어납니다. 따라서 batch size를 8분의 1로 줄이고 8번 누적하면 activation 메모리는 8분의 1 수준으로 유지하면서 그래디언트는 batch 256과 동일한 효과를 냅니다. 다만 연산 시간은 줄어들지 않고 오히려 forward/backward 호출 횟수가 늘어나는 만큼 늘어난다는 점은 감안해야 합니다.
PyTorch Gradient Accumulation 정확 구현 코드
아래는 batch size 32 micro-batch를 8번 누적해 effective batch size 256을 만드는 기본 코드입니다. 핵심은 loss를 accumulation_steps로 나눈 뒤 backward를 호출하고, optimizer.step()과 zero_grad()는 누적 주기가 끝났을 때만 실행한다는 점입니다.
accumulation_steps = 8 # micro-batch 32 x 8 = effective batch 256
model.train()
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / accumulation_steps # 누적 전 반드시 정규화
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
# 에폭 마지막에 accumulation_steps로 나눠떨어지지 않는 잔여 그래디언트 처리
if (i + 1) % accumulation_steps != 0:
optimizer.step()
optimizer.zero_grad()
Mixed Precision(AMP)을 함께 쓰는 경우 GradScaler와 조합해야 하며, scaler.step()과 scaler.update() 타이밍도 누적 주기에 맞춰야 합니다. PyTorch 공식 문서인 Automatic Mixed Precision 가이드에 예제가 정리되어 있으며, 이를 Gradient Accumulation과 결합하면 다음과 같습니다.
scaler = torch.cuda.amp.GradScaler()
optimizer.zero_grad()
for i, (inputs, labels) in enumerate(dataloader):
inputs, labels = inputs.to(device), labels.to(device)
with torch.autocast(device_type='cuda', dtype=torch.float16):
outputs = model(inputs)
loss = criterion(outputs, labels) / accumulation_steps
scaler.scale(loss).backward()
if (i + 1) % accumulation_steps == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
RTX 4060(8GB) 기준으로 이미지 분류 모델(ResNet-50급)을 학습할 때 micro-batch 32, autocast + GradScaler 조합이면 VRAM 사용량이 6~7GB 선에서 안정적으로 유지되는 경우가 많습니다. 다만 모델 크기와 입력 해상도에 따라 편차가 크므로, 실제 적용 시에는 micro-batch를 16부터 시작해 OOM이 안 나는 최대치를 찾은 뒤 accumulation_steps를 역산하는 방식을 권장합니다.
흔한 실수 3가지
첫째, loss 정규화 누락입니다. loss를 accumulation_steps로 나누지 않고 그대로 backward를 호출하면, 누적된 그래디언트 총합이 원래 의도한 batch 256의 평균 그래디언트보다 accumulation_steps배만큼 커집니다. 결과적으로 learning rate를 설정한 값보다 훨씬 크게 적용한 것과 같은 효과가 나며, 학습이 발산하거나 loss가 튀는 형태로 나타납니다. loss.backward() 직전에 반드시 loss = loss / accumulation_steps를 거쳐야 합니다.
둘째, optimizer.zero_grad() 호출 위치 오류입니다. for 루프 안에서 매 micro-batch마다 zero_grad()를 호출하면 그래디언트가 누적되지 않고 매번 초기화되어, 사실상 accumulation 없이 micro-batch 크기로만 학습하는 것과 동일해집니다. zero_grad()는 optimizer.step()을 호출한 직후, 즉 누적 주기가 끝난 시점에만 호출해야 합니다. 위 코드처럼 if (i + 1) % accumulation_steps == 0: 블록 안에 step()과 zero_grad()를 함께 둬야 합니다.
셋째, BatchNorm 레이어의 통계 왜곡입니다. Gradient Accumulation은 그래디언트 합산 관점에서는 큰 배치와 수학적으로 동일하지만, BatchNorm은 각 forward 시점의 micro-batch(예: 32개) 단위로 평균과 분산을 계산합니다. 즉 배치 256에서 계산했을 정규화 통계와는 다릅니다. 이 차이는 배치가 작을수록(예: 8 이하) 통계가 불안정해지면서 학습 결과에 영향을 줄 수 있습니다. micro-batch가 지나치게 작다면 BatchNorm 대신 GroupNorm이나 LayerNorm으로 교체하는 것이 대안이 될 수 있으며, 관련 배경 지식은 위키백과의 배치 정규화 문서에서 확인할 수 있습니다.
방식별 비교: micro-batch 크기 선택 가이드
동일한 effective batch size 256을 목표로 할 때 micro-batch 크기와 accumulation_steps 조합에 따라 VRAM 사용량과 속도, BatchNorm 안정성이 달라집니다.
| 구성 | micro-batch | accumulation_steps | VRAM 8GB 적합성 | BatchNorm 안정성 |
|---|---|---|---|---|
| A | 64 | 4 | 모델이 작을 때만 가능 | 양호 |
| B | 32 | 8 | 일반적인 CNN에 적합 | 양호 |
| C | 16 | 16 | 큰 모델/고해상도에 적합 | 다소 불안정 |
| D | 8 | 32 | 매우 큰 모델에서만 사용 | 불안정, GroupNorm 권장 |
일반적인 이미지 분류·객체 탐지 모델이라면 B 구성(micro-batch 32, accumulation 8)이 VRAM 여유와 BatchNorm 안정성의 균형이 가장 좋았습니다. Transformer 계열처럼 BatchNorm을 쓰지 않는 모델이라면 C, D 구성도 문제없이 사용할 수 있습니다.
정리 및 다음 단계
VRAM 8GB로 batch size 256 효과를 내려면 micro-batch를 줄이고 accumulation_steps만큼 그래디언트를 누적한 뒤 optimizer.step()을 호출하는 구조가 핵심입니다. loss 정규화, zero_grad() 위치, BatchNorm 통계 차이라는 3가지 실수만 피하면 대부분의 학습 파이프라인에 안전하게 적용할 수 있습니다. 지금 사용 중인 학습 스크립트가 있다면, loss.backward() 앞에 나누기 연산이 있는지, zero_grad()가 루프 안 매 스텝마다 호출되고 있지는 않은지부터 점검해 보는 것을 추천합니다. 이후에는 micro-batch 16, 32, 64를 바꿔가며 실제 VRAM 사용량을 nvidia-smi로 측정해 자신의 GPU에 맞는 최적 조합을 찾는 것이 다음 단계입니다.
자주 묻는 질문(FAQ)
Q1. Gradient Accumulation을 쓰면 학습 속도(시간)도 빨라지나요?
아닙니다. VRAM 사용량은 줄어들지만 forward/backward 연산 자체는 micro-batch 단위로 여러 번 수행되므로 총 연산량은 거의 동일하고, 오히려 Python 루프 오버헤드가 늘어나 전체 학습 시간은 비슷하거나 약간 늘어날 수 있습니다.
Q2. accumulation_steps를 8이 아니라 16, 32로 늘려도 되나요?
가능하지만 micro-batch가 지나치게 작아지면(8 이하) BatchNorm 통계가 불안정해질 수 있으므로, micro-batch는 최소 16 이상을 유지하는 것을 권장합니다.
Q3. learning rate도 batch size 256 기준으로 그대로 써도 되나요?
그래디언트 합산 관점에서는 batch 256과 동일하므로 원칙적으로는 batch 256에 맞춘 learning rate를 그대로 사용할 수 있습니다. 다만 BatchNorm 통계 차이로 인해 실제 수렴 양상이 완전히 동일하지는 않을 수 있으므로, 학습 초반 loss 추이를 보며 미세 조정하는 것이 안전합니다.