Gradient accumulation

큰 배치(batch)는 GPU 메모리를 소진하는 큰 활성화(activation)를 생성합니다. 그라디언트 축적(gradient accumulation)은 그라디언트 계산을 여러 미니 배치(mini-batch)에 걸쳐 분산시켜 더 큰 유효 배치(effective batch) 크기로 학습할 수 있게 합니다.

그라디언트는 옵티마이저가 가중치를 업데이트하기 전에 n개의 미니배치에 걸쳐 누적됩니다. 예를 들어, 장치별 배치 크기가 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()

메모리에 더 큰 배치(batch)가 들어가지 않을 때만 gradient accumulation을 사용하세요. 이는 실제 더 큰 배치로 학습하는 것보다 처리량(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,
)

손실 스케일링

사용자 정의 손실 함수의 경우, num_items_in_batch를 포함하여 Trainer가 모든 미니 배치에 걸친 예측 대상의 수로 손실을 나누도록 합니다. 이렇게 하면 gradient_accumulation_steps의 고정된 단계 수가 아닌 토큰 기준으로 정규화됩니다. 그렇지 않으면 Trainer는 손실을 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는 shifted된 라벨을 카운트합니다. 손실 함수는 라벨을 shift하여 i 위치의 예측이 i + 1 위치의 토큰을 대상으로 하도록 하며, 이로 인해 모든 시퀀스의 0번째 위치에는 대상이 남지 않습니다. Trainer는 해당 위치를 제외하고 labels[..., 1:]에 대해 카운트하므로, 분모는 손실 함수가 사용하는 예측 대상의 수와 일치합니다. padding-free collator와 같이 데이터 collator가 shift_labels를 직접 제공하는 경우, Trainer는 해당 텐서에 대해 카운트합니다. masked LM 및 분류와 같은 다른 손실 유형은 전체 라벨 텐서를 카운트합니다.

Trainer를 서브클래스로 상속하고 compute_loss()를 오버라이드하는 경우, 손실이 어떻게 정규화되는지 나타내기 위해 loss_is_scaled_for_ga를 설정합니다. 배치별 평균인 경우 False를 설정하면 Trainer가 gradient_accumulation_steps로 나누고, 이미 num_items_in_batch로 나눈 경우 True를 설정합니다.

class MyTrainer(Trainer):
    loss_is_scaled_for_ga = True

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        ...
        return loss_sum / num_items_in_batch

다음 단계

GPU 메모리 사용량이 학습 중 GPU에서 어떻게 발생하는지 이해하려면 GPU memory usage 문서를 읽어보세요.
활성화를 캐시하는 대신 다시 계산하여 활성화 메모리를 줄이는 방법을 배우려면 Gradient checkpointing 가이드를 참조하세요.
메모리를 줄이고 학습 속도를 높이기 위해 낮은 정밀도의 데이터 타입을 사용하는 방법을 배우려면 Mixed precision training 가이드를 참조하세요.
그라디언트 축적이 어떻게 계산되는지 알아보려면 Gradient Accumulation Fix 블로그 게시물을 읽어보세요.

Update on GitHub

Transformers

Gradient accumulation

손실 스케일링

다음 단계