PyTorch에서 그래디언트 누적을 구현하는 방법
이 글에서는 짧은 튜토리얼과 코드, 인터랙티브 시각화를 통해 PyTorch에서 gradient accumulation을 구현하는 방법을 배웁니다. 직접 따라 해 보세요. 이 글은 AI 번역본입니다. 오역이 의심되면 댓글로 알려 주세요.
Created on September 15|Last edited on September 15
Comment
I conducted some experiments, so now it is possible to win NLP competitions on old single GPUs such as Tesla P100. Training `deberta-v3-large` with 512 input sequences length and batch size of 4.
— Vadim (@vad13irt) June 19, 2022
TensorFlow와 달리 PyTorch는 계산 효율적인 학습 루프를 손쉽게 작성할 수 있는 방법을 제공하며, 여기에 몇 줄의 코드만으로 gradient accumulation을 간단히 추가할 수 있습니다.
💡
목차
코드를 보여주세요
대부분의 PyTorch 학습 루프는 다음과 같은 형태입니다:
optimizer = ...for epoch in range(...):for i, sample in enumerate(dataloader):inputs, labels = sample# Forward Passoutputs = model(inputs)# Compute Loss and Perform Back-propagationloss = loss_fn(outputs, labels)loss.backward()# Update Optimizeroptimizer.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 Passoutputs = model(inputs)# Compute Loss and Perform Back-propagationloss = loss_fn(outputs, labels)# Normalize the Gradientsloss = loss / NUM_ACCUMULATION_STEPSloss.backward()if ((idx + 1) % NUM_ACCUMULATION_STEPS == 0) or (idx + 1 == len(dataloader)):# Update Optimizeroptimizer.step()optimizer.zero_grad()
이것으로 충분합니다!
- 기울기 누적 단계 수에 맞춰 손실을 정규화합니다
- 옵티마이저는 각 청크마다 한 번만 업데이트합니다. 청크의 개수는 스텝 수를 누적 단계 수로 나눈 값입니다. 또는 데이터 로더의 끝에서 한 번 업데이트합니다.
요약
이 글에서는 계산 효율적인 학습 루프를 작성하기 위해 PyTorch에서 gradient accumulation을 구현하는 방법을 살펴보았습니다.
W&B의 모든 기능을 확인하려면 이 짧은 내용을 살펴보세요 5분 가이드수학적 내용과 처음부터 구현한 코드까지 다루는 더 많은 리포트를 원하신다면, 아래 댓글이나 저희의 채널에서 알려주세요 포럼 ✨!
추천 읽을거리
Preventing The CUDA Out Of Memory Error In PyTorch
A short tutorial on how you can avoid the "RuntimeError: CUDA out of memory" error while using the PyTorch framework.
How To Use GradScaler in PyTorch
In this article, we explore how to implement automatic gradient scaling (GradScaler) in a short tutorial complete with code and interactive visualizations.
How To Use Autocast in PyTorch
In this article, we learn how to implement Tensor Autocasting in a short tutorial, complete with code and interactive visualizations, so you can try it yourself.
How to Set Random Seeds in PyTorch and Tensorflow
Learn how to set the random seed for everything in PyTorch and Tensorflow in this short tutorial, which comes complete with code and interactive visualizations.
Setting Up TensorFlow And PyTorch Using GPU On Docker
A short tutorial on setting up TensorFlow and PyTorch deep learning models on GPUs using Docker.
How To Calculate Number of Model Parameters for PyTorch and TensorFlow Models
This article provides a short tutorial on calculating the number of parameters for TensorFlow and PyTorch deep learning models, with examples for you to follow.
Add a comment
There is a minor error here, you should zero out your gradients after the optimizer.step().
or else, the weights will never update
Reply
optimizer.zero_grad()
Don't perform zero_grad() on each step, only on steps where you call optimizer.step().
4 replies