Pytorch CUDA Out of Memory 해결하기

이번 공모전에 출품할 작품에는 요약, 키워드 추출 등 인공지능에 대한 기능이 많아 API 호출이 자주 발생합니다. 이를 위해 OpenAI 사에서 제공하는 ChatGPT API를 사용하게 되면 하루에 10$ 이상 크레딧을 사용하게 되었기에 LLM 서버를 직접 구축하고자 했습니다. 이 과정에서 LLaMA보다 OrionStarAI/Orion-14B-Chat 모델의 한국어 처리 성능이 뛰어나다는 조언을 받고 해당 모델을 채택하게 되었습니다.

 

그러나, 실행 과정에서 아래와 같은 오류를 마주하게 되었습니다.

 

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 300.00 MiB. GPU 0 has a total capacity of 22.09 GiB of which 147.44 MiB is free. Including non-PyTorch memory, this process has 21.94 GiB memory in use. Of the allocated memory 21.70 GiB is allocated by PyTorch, and 18.32 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.
(.venv) first@server:~/10521/subvencion/ai $

 

GPU mem이 24GiB였음에도 불구하고, 모델이 너무 무거운 탓인지 계속해서 OutOfMemory 에러가 발생했습니다.

검색해보니, 환경 변수 설정, 모델 크기 조정 등의 방법이 있었지만, 필자는 Pytorch 내의 옵션을 사용하는 방법을 선택했습니다.

 

1.  torch_dtype=torch.float16 옵션과 low_cpu_mem_usage=True 옵션을 사용하여 모델 로드 시 메모리 사용량을 줄였습니다.

2. gc.collect()와 torch.cuda.empty_cache()를 삽입하여 캐시를 자주 정리했습니다.

3. PYTORCH_CUDA_ALLOC_CONF 환경 변수의 값을 max_split_size_mb:128로 설정하여 한 번에 사용할 수 있는 메모리 크기를 설정했습니다.

 

해당 옵션을 설정하고 python3 main.py를 실행하니 아래와 같은 결과를 얻을 수 있었습니다.

 

아래 링크에서 최종 코드를 확인해보실 수 있습니다.

https://github.com/ninejuan/pytorch-orion

GitHub - ninejuan/pytorch-orion: Pytorch를 이용한 Orion Model 백엔드 코드입니다.