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

2023. 4. 26. 11:22·연구 Research/인공지능 Artificial Intelligent

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

GPU memory allocation JAX will preallocate 90% of the total GPU memory when the first JAX operation is run. Preallocating minimizes allocation overhead and memory fragmentation, but can sometimes cause out-of-memory (OOM) errors. If your JAX process fails

jax.readthedocs.io

 

 

저작자표시 비영리 변경금지 (새창열림)

'연구 Research > 인공지능 Artificial Intelligent' 카테고리의 다른 글

[JAX] 병렬컴퓨팅 예제 - jax.pmap으로 신경망 학습 예제  (0) 2023.06.13
[인공지능] 딥러닝, 머신러닝에서 uncertainty/error 개념  (0) 2023.05.16
[인공지능] Ubuntu 18.04에서 CUDA, CuDNN 설치  (1) 2023.04.16
[JAX] device 확인, default device 설정  (0) 2023.04.13
[JAX] JAX와 Torch, CUDA, cudnn 버전 맞추기  (0) 2023.04.12
'연구 Research/인공지능 Artificial Intelligent' 카테고리의 다른 글
  • [JAX] 병렬컴퓨팅 예제 - jax.pmap으로 신경망 학습 예제
  • [인공지능] 딥러닝, 머신러닝에서 uncertainty/error 개념
  • [인공지능] Ubuntu 18.04에서 CUDA, CuDNN 설치
  • [JAX] device 확인, default device 설정
보통의공대생
보통의공대생
수학,프로그래밍,기계항공우주 등 공부하는 기록들을 남깁니다.
  • 보통의공대생
    뛰는 놈 위에 나는 공대생
    보통의공대생
  • 전체
    오늘
    어제
    • 분류 전체보기 (468)
      • 공지 (1)
      • 영어 공부 English Study (40)
        • 텝스 TEPS (7)
        • 글 Article (21)
        • 영상 Video (10)
      • 연구 Research (99)
        • 최적화 Optimization (3)
        • 데이터과학 Data Science (7)
        • 인공지능 Artificial Intelligent (40)
        • 제어 Control (45)
      • 프로그래밍 Programming (103)
        • 매트랩 MATLAB (25)
        • 파이썬 Python (33)
        • 줄리아 Julia (2)
        • C++ (3)
        • 리눅스 우분투 Ubuntu (6)
      • 항공우주 Aeronautical engineeri.. (21)
        • 항법 Navigation (0)
        • 유도 Guidance (0)
      • 기계공학 Mechanical engineering (13)
        • 열역학 Thermodynamics (0)
        • 고체역학 Statics & Solid mechan.. (10)
        • 동역학 Dynamics (1)
        • 유체역학 Fluid Dynamics (0)
      • 수학 Mathematics (34)
        • 선형대수학 Linear Algebra (18)
        • 미분방정식 Differential Equation (3)
        • 확률및통계 Probability & Sta.. (2)
        • 미적분학 Calculus (1)
        • 복소해석학 Complex Analysis (5)
        • 실해석학 Real Analysis (0)
      • 수치해석 Numerical Analysis (27)
      • 확률 및 랜덤프로세스 Random process (2)
      • 추론 & 추정 이론 Estimation (3)
      • 기타 (26)
        • 설계 프로젝트 System Design (8)
        • 논문작성 Writing (55)
        • 세미나 Seminar (2)
        • 생산성 Productivity (3)
      • 실험 Experiment (1)
      • 유학 생활 Daily (8)
  • 블로그 메뉴

    • 홈
    • 태그
    • 방명록
  • 링크

  • 공지사항

  • 인기 글

  • 태그

    수치해석
    Linear algebra
    JAX
    IEEE
    생산성
    obsidian
    텝스
    teps
    Numerical Analysis
    pytorch
    Dear abby
    matplotlib
    텝스공부
    에러기록
    ChatGPT
    LaTeX
    논문작성
    MATLAB
    Statics
    옵시디언
    Julia
    고체역학
    우분투
    WOX
    서버
    인공지능
    논문작성법
    Zotero
    딥러닝
    Python
  • 최근 댓글

  • 최근 글

  • hELLO· Designed By정상우.v4.10.3
보통의공대생
[JAX] 메모리 부족 문제 해결
상단으로

티스토리툴바