[PyTorch 에러 해결] CUDA out of memory (feat. 주피터 노트북)

실행 환경

• OS: Ubuntu 18.04.5
• PyTorch 버전: 2.0.1
• CUDA 버전: 11.0

 

문제 상황

모델 크기가 23GB인 LLM을 GPU 메모리에 로드하던 중, CUDA out of memory 에러가 발생했다. GPU 메모리 크기가 32GB인데다, 주피터 노트북 파일에서 테스트 했을 때 문제 없이 실행되었기 때문에 코드 구현 상에 메모리 누수가 있는지 확인했으나, 별다른 문제를 찾지 못했다.

문제 해결

nvidia-smi 명령어로 GPU의 상태를 확인해 본 결과, 파이썬 스크립트를 동작시키지 않은 상황에서 GPU 메모리가 23GB가 할당되어 있는 것을 확인할 수 있었다. 알고보니 주피터 노트북 파일을 실행시켰을 때 GPU 메모리에 올라갔던 모델을 해제하지 않아서 메모리가 반환되지 않고 있었다.

 

따라서 주피터 커널을 재시작하여 메모리를 반환하여 문제를 해결했다.