큰 배치(big batch)는 더 부드럽고 노이즈가 적은 그래디언트(gradient)를 제공하며, 레시피가 튜닝된 배치 크기로 학습할 수 있게 해줍니다. 하지만 순전파/역전파(forward/backward) 패스의 활성화 메모리(activation memory)는 그 안의 샘플 수에 비례하여 확장되므로, 크기가 B인 배치는 GPU에 단순히 들어가지 않을 수 있습니다. Gradient accumulation은 적은 메모리로 큰 배치를 흉내 냅니다. 즉, 유효 배치(effective batch)를 N개의 작은 마이크로 배치(micro-batches)로 나누고, 각 배치에 대해 순전파와 역전파를 실행하여 그래디언트가 쌓이게 하는 것입니다. 저는 이것이 수치적으로 동일함을 증명하는 데모를 만들었습니다. 그 아이디어는 다음과 같습니다.

.grad는 이미 누적됩니다

이 메커니즘은 비용이 들지 않습니다. PyTorch의 .backward()는 각 파라미터의 .grad에 값을 더합니다. 절대 덮어쓰지 않습니다. 보통은 매 스텝마다 zero_grad()를 호출하여 이를 숨깁니다. 누적(Accumulation)은 단지... 그렇게 하지 않을 뿐입니다. 여러 번의 역전파 패스가 동일한 버퍼에 쌓이도록 허용합니다.

loss_a.backward()   # p.grad = g_a
loss_b.backward()   # p.grad = g_a + g_b   <-- 교체되는 것이 아니라 더해짐
# zero_grad()를 건너뛰면 마이크로 배치 간에 그래디언트가 누적됩니다

N개의 마이크로 배치 동안 스텝을 지연시키기

크기가 m인 N개의 마이크로 배치를 공급하고, 각각에 대해 역전파를 수행하되, N번마다 한 번씩만 optimizer.step()을 호출합니다. 마이크로 배치 사이에서 버퍼는 커지지만, 단 하나의 마이크로 배치 활성화(m)만 살아있으므로, 피크 활성화 메모리(peak activation memory)는 N과 무관하게 일정하게 유지됩니다.

for i, (xb, yb) in enumerate(loader):   # 크기가 m인 마이크로 배치들
    criterion(model(xb), yb).backward() # .grad에 더함
    if (i + 1) % N == 0:                 # N개의 마이크로 배치마다...
...

1/N로 스케일링하기 — 모두가 잊어버리는 스텝

각 마이크로 배치 손실(loss)은 이미 m에 대한 *평균(mean)*입니다. 이것들을 N개 합치면 m·N에 대한 평균의 N× 배가 되어, N만큼 너무 커지게 됩니다. 각 손실을 N으로 나누면 누적된 버퍼는 유효 배치에 대한 진정한 평균이 됩니다. 동일한 크기 평균들의 평균은 전체 평균과 같기 때문에, 이는 B = m·N인 실제 큰 배치와 비트 단위로(bit-for-bit) 일치합니다. 이를 잊어버리면 여러분의 학습률(learning rate)을 조용히 N배만큼 곱하게 되는 셈입니다.

loss = criterion(model(xb), yb) / N    # <-- 합계(SUM)가 평균(MEAN)이 되도록 스케일링
# /N이 없는 경우:  buffer = Σ mean_j     = N × 실제 평균 (가장 흔한 누적 버그)
# /N이 있는 경우:  buffer = Σ mean_j / N = 큰 배치(big-batch)의 평균  ✓

비용과 얻는 것

전체 FLOPs(부동 소수점 연산량)는 변하지 않습니다. 어떤 방식이든 동일한 B개의 샘플을 처리하기 때문입니다. 따라서 여러분이 지불하는 것은 _시간_입니다. 업데이트당 N번의 순차적인 작은 패스(pass)를 수행해야 하며, 업데이트 빈도는 N배만큼 줄어듭니다. 이는 재계산 패스(recompute pass)를 추가하는 48일 차의 그래디언트 체크포인팅 (gradient checkpointing)과는 극명한 대조를 이룹니다. 체크포인팅은 고정된 배치 크기에서 연산량을 희생하여 활성화 메모리 (activation memory)를 확보하는 방식인 반면, 누적 (accumulation)은 고정된 메모리에서 배치 크기를 위해 단계(steps)/시간을 희생하는 방식입니다. 이 두 기술은 깔끔하게 결합될 수 있습니다.

실제로 작성하게 될 루프

이를 종합하면 일반적인 루프 주변에 다섯 줄 정도만 추가하면 됩니다. 메모리는 m만큼 사용하면서, m·N 크기의 큰 배치와 동일한 가중치(weights)를 가질 수 있습니다. 이는 혼합 정밀도 (mixed precision, autocast + GradScaler) 및 48일 차의 체크포인팅과도 잘 어우러집니다.

opt.zero_grad()
for i, (xb, yb) in enumerate(loader):        # 마이크로 배치 (micro-batches), 크기 m
    loss = criterion(model(xb), yb) / N       # 합계가 평균이 되도록 스케일링
...

단 하나의 진짜 주의사항: BatchNorm

이 등가성은 샘플별 손실 (per-sample losses)에 대해서만 정확합니다. BatchNorm (BN)은 유효한 m·N이 아니라 현재 마이크로 배치인 m에 대해 정규화(normalize)를 수행하므로, BN이 포함되어 있다면 누적은 실제 큰 배치와 동일하지 않습니다. GroupNorm, LayerNorm, 또는 SyncBN을 사용하세요. 또한 마지막 부분(tail)도 주의해야 합니다. 만약 로더(loader)의 크기가 N의 배수가 아니라면, 루프가 끝난 후 남은 그룹을 플러시(flush)해야 합니다. 그렇지 않으면 해당 샘플들을 조용히 누락하게 됩니다.

제가 얻은 교훈은 이것입니다: 큰 배치는 수학의 문제가 아니라 메모리의 문제입니다. .grad는 이미 합산되므로, N개의 마이크로 배치를 전달하고, 1/N을 기억하며, 한 번만 스텝(step)을 수행하세요. 그러면 작은 배치 하나만큼의 비용으로 큰 배치를 얻을 수 있습니다.

m과 N을 설정하고, 누적된 그래디언트 (gradient)가 큰 배치 (batch)와 기계 정밀도 0(machine zero) 수준까지 일치하는 동안 메모리 바 (memory bar)가 평평하게 유지되는 것을 확인하세요:
https://dev48v.infy.uk/dl/day49-gradient-accumulation.html