PyTorch CUDA OutOfMemoryError 해결

HL·2023년 8월 23일

PyTorch cuda

Transformer 모델 학습 과정에서 CUDA out of memory 에러가 발생했다.

로컬 컴퓨터에서든 코랩에서든 OOM(Out Of Memory error)가 계속 발생해서, 여러 시도도 해보고 해결책도 찾아봤다.

나는 현재 하나의 GPU만을 사용하는 상황인지라, 여러개의 GPU를 사용하는 경우는 잘 고려하지 않았다.

CUDA OOM이 발생한다면

CUDA OOM이 뜨는 경우 두가지 가능성이 있다.

실제로 GPU 메모리가 부족한 상황
실제로는 GPU 메모리가 부족하지 않은데 부족하다고 에러가 뜨는 상황

내가 받은 에러 메시지를 보면,

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 58.00 MiB (GPU 0; 5.80 GiB total capacity; 4.22 GiB already allocated; 75.00 MiB free; 4.35 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

이걸로 보면 나는 아마 두번째 경우인 듯하다.
그래도 로컬 컴퓨터 GPU 메모리가 워낙 작아서(5기가) 첫번째 경우도 많이 맞이할 것이기에... 미래를 위해 그냥 다 알아봤다.

첫번째 경우라면 말그대로 메모리가 부족한 것이기에 아래 해결책 파트에서 1, 2, 3번을 주로 봐야 할 것 같고
두번째 경우라면 주로 4번을 봐야할 것 같다.

해결책

실제 시도해보거나, 찾아본 CUDA OOM의 해결책은 크게 다음과 같다.

학습 데이터 크기 조절
모델 크기 조절
그래디언트 데이터 크기 조절
GPU 메모리를 직접 관리
GPU말고 CPU에서 돌리기

1. 학습 데이터 크기 조절

데이터 배치 사이즈 줄이기

배치 사이즈 128일 때 OOM 뜬다면, 64로 해보는 등
그래도 메모리 부족하면 batch_size=32, ...

데이터를 나눠서 다운받기

큰 데이터가 한번에 GPU 메모리에 올라가지 않도록
감당할 수 있을만큼만 나눠서 메모리에 올리기
divide-and-conquer 느낌으로

Hugging Face의 IterableDataset, PyTorch의 DataLoader 활용할 수 있음

from datasets import load_dataset
from torch.utils.data import DataLoader

batch_size = 1000

de_en_dataset = load_dataset("wmt14", "de-en", streaming=True, split="train")
de_en_dataloader = DataLoader(de_en_dataset, batch_size=batch_size)

데이터 나눠서 학습시키기

위의 데이터 나눠서 다운받기 와도 연결

모델을 백업해놓고 이어서 학습시키든
예기치 못하게 에러로 학습 중단되는 문제에도 대비할 수 있음 (중간중간 백업본 저장. "checkpoint")

데이터 사이즈 줄이기

예시를 들자면

이미지의 경우, 픽셀 수 줄이기
텍스트의 경우, 텍스트 임베딩 사이즈 줄이기

이와 같이 중요성 떨어지는 데이터/피쳐부터 포기함으로써 데이터 사이즈 줄이기

데이터 변형이기에 모델 성능에 영향 미칠 수 있음을 주의해야 한다.

데이터 없애기

필요 없어진 변수/텐서는 제거하거나, gpu 연산 필요없다면 애초에 cpu에 올리거나

# `x` usage end
del x

2. 모델 크기 조절

파라미터 줄이기

레이어 수 줄인다든가

파라미터 줄이는 건 아래의 그래디언트 데이터 크기 조절와도 이어짐

Mixed Precision Training

자료형이 차지하는 데이터 아껴서 (대신 숫자 정확도 좀 포기) 메모리 사용량 아낌
16-bit float 쓴다든가

torch.cuda.amp

model = model.half()
input = input.half()

제대로 알아보지 않아 잘 모르겠지만, 숫자 정확도를 좀 포기하게 되니 모델 성능에도 영향이 없진 않을 듯하다.

3. 그래디언트 데이터 크기 조절

그래디언트 자체도 데이터이므로 메모리 공간을 차지한다.

쓸데없는 그래디언트 저장하지 않기

그래디언트 구할 필요 없을 때는 안구하도록

require_grad가 True인 변수에 대해서는, 어떤 연산을 하든 그 변수에 대해 gradient를 계산한다. 그래디언트 저장할 필요 없는 변수 혹은 연산이라면 이런 경우 require_grad가 False인지 확인해야 연산도 의도대로 되고 메모리도 아낄 수 있을 것이다.

# after `optimizer.step()`
# 그래디언트가 누적되지 않도록 학습 스텝마다 초기화
optimizer.zero_grad()

# disables gradient computation
with torch.no_grad():
    ...

알고리즘

당연히 모델에 영향 미친다는 점 주의해야 함

optimizer의 경우, Adam 알고리즘은 그래디언트의 1st moment와 2nd moment 둘다 필요로 한다. 모멘트를 사용하지 않는 알고리즘에 비하면 연산량과 용량을 그만큼 더 필요로 할 것
model의 경우, RNN에 너무 긴 sequence 넣는다거나, linear layer 너무 큰 것을 사용한다면 메모리 용량 많이 먹을 것

Gradient Checkpointing

"Don't need to store it, can compute it again"

x1 = f1(x0)
x2 = torch.utils.checkpoint(f2, x1) # f2(x1)
x3 = f3(x2)

체크포인트 사용하는 경우 x2.grad_fn에 실제 연산 결과를 저장하지 않고 이후 필요할 때 계산하도록 한다. 그래서 이후 f3(x2) 연산을 위한 메모리 공간을 그만큼 더 확보할 수 있다.

4. GPU 메모리를 직접 관리

PyTorch는 caching memory allocator를 사용한다.
이것때문에 실제로는 사용되지 않는 메모리가 nvidia-smi에 의해서는 샤옹되고 있는 것처럼 보여질 수 있다.

실제로 텐서가 차지하고 있는 메모리 공간 확인하려면 memory_allocated(), max_memory_allocated() 사용하고
caching allocator가 관리하는 메모리 공간 총 얼마인지 보고 싶으면 memory_reserved(), max_memory_reserved() 사용하면 된다.

empty_cache() 호출하면 파이토치의 cached memory지만 사용되지 않는 부분을 release하게 된다.
그러나 실제로 텐서가 이미 차지하고 있는 메모리는 비워지지 않는다.

# nvidia
# this is terminal command
nvidia-smi

# pytorch cuda
memory_allocated()
max_memory_allocated()
memory_reserved()
max_memory_reserved()

empty_cache()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if device=="cuda": torch.cuda.empty_cache()

max_split_size_mb 설정

메모리 조각화 방지
torch.cuda.memory_allocated()로 현재 할당된 메모리 확인하고, 현재 할당된 메모리보다 크게 설정하면 됨

5. GPU 말고 CPU에서 돌리기

이것도 하나의 방법임, 그러나 당연히 모델 학습이라면 많은 경우 시간 더 걸릴 것

GPU 굳이 필요 없는 연산은 CPU에서 돌리는 것도 포함
코드 전체를 CPU에서 돌릴 수도 있고

더 알아볼 것

그 외에도 조사하면서 나온 해결책이지만 잘 모르겠어서 분류 안한 것,

gc.collect()

잘 모르겠는 파이토치 특성

PyTorch's caching memory allocator

출처

다음 포스트