그라디언트 누적

그라디언트 누적 (Gradient accumulation)

큰 배치는 큰 활성화(activation)를 만들어내서 GPU 메모리를 고갈시켜요. 그라디언트 누적은 그라디언트 계산을 여러 mini-batch에 분산시켜서 더 큰 유효 배치 크기로 학습할 수 있게 해줘요.

출처: 문서

본문

옵티마이저가 가중치를 업데이트하기 전까지 그라디언트가 n개의 mini-batch에 걸쳐 누적돼요. 예를 들어 디바이스당 배치 크기가 8이고 누적 단계가 4라면 유효 배치 크기는 32가 돼요.

Step 1: mini-batch 1 → forward → backward → grads = G₁
Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂
Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃
Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄
        → optimizer.step()  ← same update as if batch_size × 4
        → zero_grad()

그라디언트 누적은 더 큰 배치가 메모리에 들어가지 않을 때만 사용해요. 진짜 큰 배치로 학습하는 것보다 처리량(throughput)을 높여주지는 않아요.

per_device_train_batch_size에 걸쳐 gradient_accumulation_steps 만큼 그라디언트를 누적해요.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
)

Loss 스케일링

커스텀 loss function을 사용한다면 num_items_in_batch를 포함해서 Trainer가 모든 mini-batch에 걸친 예측 대상 수로 loss를 나누게 해요. 이렇게 하면 gradient_accumulation_steps로 고정된 단계 수 대신 토큰 수 기준으로 정규화돼요. 그렇지 않으면 Trainer는 loss를 gradient_accumulation_steps로 나눠요.

import torch.nn.functional as F

def compute_loss(outputs, labels, num_items_in_batch=None):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

causal LM 모델에서는 num_items_in_batch가 shift된 라벨을 세요. loss가 라벨을 shift해서 위치 i의 예측이 위치 i + 1의 토큰을 타겟으로 하기 때문에, 모든 시퀀스의 위치 0에는 타겟이 없게 돼요. Trainer는 그 위치들을 제외하고 labels[..., 1:]에 대해 세므로, 분모가 loss가 실제로 사용하는 예측 대상 수와 일치하게 돼요. 데이터 콜레이터가 패딩 없는 콜레이터처럼 shift_labels를 직접 제공하면, Trainer는 그 대신 해당 텐서에 대해 세요. masked LM, classification 같은 다른 loss 유형은 전체 라벨 텐서를 세요.

더 알아보기 (Learn more)

학습 중 GPU 메모리를 무엇이 사용하고 있는지 이해하려면 GPU memory usage 문서를 읽어 보세요.
활성화를 캐싱하는 대신 다시 계산해서 활성화 메모리를 줄이는 방법은 Gradient checkpointing 가이드를 확인해 보세요.
더 낮은 정밀도 데이터 타입으로 메모리를 줄이고 학습을 빠르게 하는 방법은 Mixed precision training 가이드를 확인해 보세요.
그라디언트 누적이 어떻게 계산되는지 배우려면 Gradient Accumulation Fix 블로그 글을 읽어 보세요.