[JAX] 메모리 부족 문제 해결

JAX를 쓰다가 너무 많은 양의 데이터를 쓰다보니 메모리 부족(OOM: Out of memory) 현상을 겪었다.

 

근본적인 해결책은 달리 없다. 데이터가 너무 많아서 생기는 문제이니 데이터 양을 줄이던가 아니면 병렬 컴퓨팅을 하는 방법이 있다. 병렬 컴퓨팅을 간단한 코드에서는 실행해보았는데 큰 네트워크에서는 해본 적이 없다.

 

일단 임시방편으로는 다음과 같다.

 

메모리 부족이 쉽게 발생할 수 있는 이유는 JAX에서 처음에 import를 할 때 대부분의 메모리를 미리 할당해놓기 때문이다. 따라서 이 preallocation을 막거나 줄이면 도움이 된다.

 

 

1. Preallocation 중단

XLA_PYTHON_CLIENT_PREALLOCATE=false
# 구체적으로는 다음과 같이 구현한다.
import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

JAX를 import하기만 해도 

다음과 같이 메모리를 거의 대부분 잡아먹게 된다. 이는 JAX가 일부러 preallocation해놓은 것이다. 그래서 위의 명령을 통해 preallocation을 false로 해놓고 JAX를 import하면 이렇게 메모리를 잡아먹는 것을 피할 수 있다. (반드시 JAX import 전에 수행할 것!)

 

2. allocation fraction 제한

XLA_PYTHON_CLIENT_MEM_FRACTION=.XX

그러나 1번과 같이 preallocation을 false로 해놓으면 학습 속도가 느려지기 때문에 부분적으로 preallocation이 차지하는 비중을 위와 같이 줄일 수 있다. 위의 XX값이 원래는 90이라고 한다. 즉, 90%를 미리 할당해놓는다. 대신 60정도로 낮출 수 있을 것이다.

 

 

3. 필요없는 메모리는 deallocate

XLA_PYTHON_CLIENT_ALLOCATOR=platform

위와 같이 설정을 해놓으면 더이상 필요없는 메모를 deallocate할 수 있다. 그러나 이 경우에는 다시 쓰고 싶은 데이터를 deallocate했다가 다시 allocate하는 과정이 생길 것이다. 학습 속도가 느려지기 때문에 추천하지 않는다고 한다는데 그래도 최소한의 데이터를 사용하므로 메모리 여유가 생긴다.

 

 

그래서 나는 메모리 부족 문제가 생길 때 아래와 같이 설정해놓는다. 

 

import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"
# 아래 코드는 원하는 GPU 번호만 쓰도록 설정하는 코드
os.environ['CUDA_VISIBLE_DEVICES'] = "2"

 

 

 

---

https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

GPU memory allocation — JAX documentation