Skip to main content

PyTorch에서 그래디언트 누적을 구현하는 방법

이 글에서는 짧은 튜토리얼과 코드, 인터랙티브 시각화를 통해 PyTorch에서 gradient accumulation을 구현하는 방법을 배웁니다. 직접 따라 해 보세요. 이 글은 AI 번역본입니다. 오역이 의심되면 댓글로 알려 주세요.
Created on September 15|Last edited on September 15

이 글에서는 gradient accumulation을 구현하는 방법을 살펴봅니다 PyTorch 계산 효율적인 학습 루프를 작성하기 위해서.
TensorFlow와 달리 PyTorch는 계산 효율적인 학습 루프를 손쉽게 작성할 수 있는 방법을 제공하며, 여기에 몇 줄의 코드만으로 gradient accumulation을 간단히 추가할 수 있습니다.
PyTorch에서 사용할 수 있는 다양한 혼합 정밀도 방법을 더 자세히 살펴보려면 다음을 참조하세요 공식 문서를 참조하세요.
💡

목차





코드를 보여주세요

대부분의 PyTorch 학습 루프는 다음과 같은 형태입니다:
optimizer = ...

for epoch in range(...):
for i, sample in enumerate(dataloader):
inputs, labels = sample

# Forward Pass
outputs = model(inputs)
# Compute Loss and Perform Back-propagation
loss = loss_fn(outputs, labels)
loss.backward()
# Update Optimizer
optimizer.step()

optimizer.zero_grad()
또한 PyTorch로 딥러닝 모델, 특히 대규모 언어 모델(LLMs)을 일정 기간 학습해 보았다면, 아마 다음과 같은 CUDA 오류를 겪어 보셨을 것입니다:
RuntimeError: CUDA out of memory. Tried to allocate .. MiB (.. GiB total capacity; ... GiB already allocated;
... MiB free; ... GiB reserved in total by PyTorch)
CUDA OOM 오류를 방지하는 방법에 대한 훨씬 더 자세한 개요는 다음 보고서를 참조하세요:



대규모 언어 모델을 학습할 때는 배치 크기를 크게 할수록 수렴이 더 잘되는 경향이 있음이 알려져 있습니다. 하지만 대부분의 경우 더 큰 배치 크기를 GPU 메모리에 넣을 수 없습니다. 그럼 해결책은 무엇일까요? 바로 “Gradient Accumulation”을 사용하는 것입니다. 위의 코드 스니펫에서처럼 일반적으로는 데이터 배치를 받아서 forward pass를 수행한 뒤 곧바로 backpropagation을 실행합니다. 그런데 매 배치마다 바로 업데이트하는 대신, 여러 번의 forward pass에서 발생한 기울기를 저장해 두었다가 backpropagation을 수행하면 어떨까요? 이런 방식으로 몇 단계 동안 기울기를 “누적”한 후에 backpropagation을 실행하는 것이 Gradient Accumulation입니다.
자, 이제 Gradient Accumulation이 무엇인지 알았으니, 이를 PyTorch 학습 루프에서 어떻게 적용하는지 살펴보겠습니다:
optimizer = ...
NUM_ACCUMULATION_STEPS = ...

for epoch in range(...):
for idx, sample in enumerate(dataloader):
inputs, labels = sample

# Forward Pass
outputs = model(inputs)
# Compute Loss and Perform Back-propagation
loss = loss_fn(outputs, labels)

# Normalize the Gradients
loss = loss / NUM_ACCUMULATION_STEPS
loss.backward()

if ((idx + 1) % NUM_ACCUMULATION_STEPS == 0) or (idx + 1 == len(dataloader)):
# Update Optimizer
optimizer.step()

optimizer.zero_grad()
이것으로 충분합니다!
  1. 기울기 누적 단계 수에 맞춰 손실을 정규화합니다
  2. 옵티마이저는 각 청크마다 한 번만 업데이트합니다. 청크의 개수는 스텝 수를 누적 단계 수로 나눈 값입니다. 또는 데이터 로더의 끝에서 한 번 업데이트합니다.

요약

이 글에서는 계산 효율적인 학습 루프를 작성하기 위해 PyTorch에서 gradient accumulation을 구현하는 방법을 살펴보았습니다.
W&B의 모든 기능을 확인하려면 이 짧은 내용을 살펴보세요 5분 가이드수학적 내용과 처음부터 구현한 코드까지 다루는 더 많은 리포트를 원하신다면, 아래 댓글이나 저희의 채널에서 알려주세요 포럼 ✨!
다음 다른 리포트도 확인해 보세요 완전 연결 GPU 활용도와 모델 저장과 같은 다른 기본 개발 주제도 다룹니다.

추천 읽을거리



이 글은 AI로 번역되었습니다. 오역이 의심되면 댓글로 알려주세요. 원문 리포트 링크는 다음과 같습니다: 원문 리포트 보기
aurko
aurko •  
There is a minor error here, you should zero out your gradients after the optimizer.step(). or else, the weights will never update
Reply
Alexey Zaytsev
Alexey Zaytsev •  
optimizer.zero_grad() Don't perform zero_grad() on each step, only on steps where you call optimizer.step().
4 replies