Ctrl+K

JAX 설치하기

JAX 설치하기#

JAX는 순수 파이썬으로 작성되어 있지만, XLA에 의존하며, 이는 jaxlib 패키지로 설치해야 합니다. pip 또는 conda를 사용해 이진 패키지를 설치, Docker 컨테이너를 사용, 또는 소스로부터 JAX 빌드를 위해서는 다음 지침을 확인하세요.

지원되는 플랫폼#

Linux x86_64

Linux aarch64

Mac x86_64

Mac ARM

Windows x86_64

Windows WSL2 x86_64

CPU

예

예

예

예

예

예

NVIDIA GPU

예

예

아니오

해당사항 없음

아니오

실험적

Google TPU

예

해당사항 없음

해당사항 없음

해당사항 없음

해당사항 없음

해당사항 없음

AMD GPU

실험적

아니오

아니오

해당사항 없음

아니오

아니오

Apple GPU

해당사항 없음

아니오

실험적

실험적

해당사항 없음

해당사항 없음

Linux(Ubuntu 20.04 이상) 및 macOS(10.12 이상) 플랫폼에서 jaxlib 설치 또는 빌드를 지원합니다. 또한 실험적으로 네이티브 Windows 지원도 있습니다.

Windows 사용자는 Windows Subsystem for Linux를 통해 CPU 및 GPU에서 JAX를 사용할 수 있으며, 대안으로 네이티브 Windows CPU만 지원을 사용할 수 있습니다.

CPU#

pip 설치: CPU#

현재 다음 운영 체제 및 아키텍처에 jaxlib 휠을 배포하고 있습니다:

  • Linux, x86-64

  • Mac, Intel

  • Mac, ARM

  • Windows, x86-64 (experimental)

랩탑에서 로컬 개발을 하는 데 유용할 수 있는 CPU 전용 버전의 JAX를 설치하려면, 다음을 실행할 수 있습니다

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

Windows에서는 이미 설치되어 있지 않다면 Microsoft Visual Studio 2019 Redistributable을 설치해야 할 수도 있습니다.

다른 운영 체제 및 아키텍처는 소스로부터 빌드해야 합니다. 다른 운영 체제 및 아키텍처에서 pip로 설치를 시도하면 jax는 성공적으로 설치될 수 있지만 jaxlib가 jax와 함께 설치되지 않을 수 있으며, 이 경우 실행 시 실패할 수 있습니다.

NVIDIA GPU#

JAX는 SM 버전 5.2(Maxwell) 이상을 가진 NVIDIA GPU를 지원합니다. Kepler 시리즈 GPU는 NVIDIA가 소프트웨어에서 Kepler GPU 지원을 중단함에 따라 더 이상 JAX에서 지원되지 않습니다.

먼저 NVIDIA 드라이버를 설치해야 합니다. NVIDIA에서 제공하는 가장 최신 드라이버를 설치하는 것을 권장하지만, CUDA 12의 경우 드라이버 버전이 >= 525.60.13, CUDA 11의 경우 Linux에서 >= 450.80.02 이어야 합니다. NVIDIA 드라이버를 쉽게 업데이트할 수 없는 클러스터 등에서 더 새로운 CUDA 툴킷을 이전 드라이버와 함께 사용해야 하는 경우, 이 목적으로 NVIDIA가 제공하는 CUDA forward compatibility packages를 사용할 수 있습니다.

pip 설치: GPU (CUDA, pip를 통한 설치, 더 쉬움)#

NVIDIA GPU 지원과 함께 JAX를 설치하는 두 가지 방법이 있습니다: pip 휠에서 설치한 CUDA와 CUDNN을 사용하는 방법과, 자체 설치한 CUDA/CUDNN을 사용하는 방법입니다. pip 휠을 사용하여 CUDA와 CUDNN을 설치하는 것이 훨씬 쉽기 때문에, 이 방법을 강력히 권장합니다! 이 방법은 NVIDIA가 aarch64 CUDA pip 패키지를 출시하지 않았기 때문에 x86_64에서만 지원됩니다.

pip install --upgrade pip

# CUDA 12 설치
# 주의: 휠은 리눅스에서만 가능.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11 설치
# 주의: 휠은 리눅스에서만 가능.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

JAX가 잘못된 버전의 CUDA 라이브러리를 감지할 경우, 확인해야 할 몇 가지 사항이 있습니다:

  • LD_LIBRARY_PATH가 설정되지 않았는지 확인하세요. LD_LIBRARY_PATH는 CUDA 라이브러리를 오버라이드할 수 있습니다.

  • JAX가 요청한 CUDA 라이브러리가 설치되어 있는지 확인하세요. 위의 설치 명령어를 다시 실행하면 작동해야 합니다.

pip 설치: GPU (CUDA, 로컬에 설치, 더 어려움)#

미리 설치된 CUDA를 사용하고 싶다면, 먼저 CUDA와 CuDNN을 설치해야 합니다.

JAX는 Linux x86_64만을 위한 미리 빌드된 CUDA 호환 휠을 제공합니다. 다른 운영 체제 및 아키텍처 조합도 가능하지만, 소스로부터 빌드가 필요합니다.

사용하는 NVIDIA 드라이버 버전은 CUDA toolkit’s corresponding driver version보다 적어도 같거나 새로워야 합니다. 예를 들어, NVIDIA 드라이버를 쉽게 업데이트할 수 없는 클러스터에서 더 새로운 CUDA 툴킷을 이전 드라이버와 함께 사용해야 한다면, 이 목적으로 NVIDIA가 제공하는 CUDA forward compatibility packages를 사용할 수 있습니다.

JAX는 현재 두 가지 CUDA 휠 변형을 제공합니다:

  • CUDA 12.3, cuDNN 8.9, NCCL 2.16

  • CUDA 11.8, cuDNN 8.6, NCCL 2.16

CUDA, cuDNN, 및 NCCL 설치의 주 버전이 일치하고, 부 버전이 동일하거나 더 새로운 경우 JAX 휠을 사용할 수 있습니다. JAX는 라이브러리 버전을 확인하며, 충분히 새롭지 않은 경우 오류를 보고합니다.

NCCL은 선택적 종속성이며, 여러 GPU를 사용한 계산을 수행하는 경우에만 필요합니다.

설치하려면, 아래의 내용을 실행하세요.

pip install --upgrade pip

# CUDA 12와 cuDNN 8.9 이상과 호환되는 휠을 설치.
# 주의: 휠은 linux에서만 가능.
pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11과 cuDNN 8.6 이상과 호환되는 휠을 설치.
# 주의: 휠은 linux에서만 가능.
pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

이 pip 설치는 Windows에서 작동하지 않고 안내 없이 실패할 것입니다; 위의 JAX 설치하기를 참고하세요.

아래 명령어를 통해 CUDA 버전을 확인할 수 있습니다:

nvcc --version

JAX는 CUDA 라이브러리를 찾기 위해 LD_LIBRARY_PATH를 사용하고 바이너리(ptxas, nvlink)를 찾기 위해 PATH를 사용합니다. 이 경로들이 올바른 CUDA 설치를 가리키는지 확인해 주세요.

미리 빌드된 휠과 관련하여 오류나 문제가 발생하는 경우, 이슈 트래커를 통해 알려주세요.

Docker 컨테이너: NVIDIA GPU#

NVIDIA는 JAX Toolbox 컨테이너를 제공하며, 이는 JAX의 나이틀리 릴리즈와 일부 모델/프레임워크를 포함한 최신 컨테이너입니다.

나이틀리 설치#

나이틀리 릴리스는 빌드된 시점에서 메인 저장소의 상태를 반영하며, 전체 테스트 스위트를 통과하지 않을 수 있습니다.

  • JAX:

pip install -U --pre jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html

  • Jaxlib CPU:

pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html

  • Jaxlib TPU:

pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
pip install -U --pre libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

  • Jaxlib GPU (Cuda 12):

pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda12_releases.html

  • Jaxlib GPU (Cuda 11):

pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_cuda_releases.html

Google TPU#

pip 설치: Google Cloud TPU#

JAX는 Google Cloud TPU를 위한 미리 빌드된 휠을 제공합니다. 적절한 버전의 jaxlib 및 libtpu와 함께 JAX를 설치하려면, 클라우드 TPU VM에서 다음을 실행할 수 있습니다:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

대화형 노트북 사용자를 위해: Colab TPU는 JAX 버전 0.4부터 JAX를 더 이상 지원하지 않습니다. 그러나 클라우드에서 대화형 TPU 노트북을 사용하고자 한다면, JAX를 완전히 지원하는 Kaggle TPU notebooks을 사용할 수 있습니다.

Apple GPU#

pip 설치: Apple GPUs#

Apple은 Apple GPU 하드웨어를 위한 실험적인 Metal 플러그인을 제공합니다. 자세한 내용은 Apple의 JAX on Metal 문서를 참고하세요.

Metal 플러그인과 관련된 몇 가지 주의사항이 있습니다:

  • Metal 플러그인은 새롭고 실험적이며 알려진 문제들이 여러 개 있습니다. 문제가 발생하면 JAX 이슈 트래커에 보고해 주세요.

  • Metal 플러그인은 현재 매우 구체적인 버전의 jax와 jaxlib을 요구합니다. 이 제한은 시간이 지나면서 플러그인 API가 성숙해짐에 따라 완화될 것입니다.

AMD GPU#

JAX has experimental ROCM support. There are two ways to install JAX:

Conda#

Conda 설치#

jax의 커뮤니티 지원 빌드가 있습니다. conda를 통해 설치하려면, 다음을 실행하세요.

conda install jax -c conda-forge

NVIDIA GPU가 있는 기계에 설치하려면, 다음을 실행하세요.

conda install jaxlib=*=*cuda* jax cuda-nvcc -c conda-forge -c nvidia

conda-forge에 의해 배포된 cudatoolkit은 JAX가 필요로 하는 ptxas가 누락되어 있습니다. 따라서 nvidia 채널에서 cuda-nvcc 패키지를 설치하거나, ptxas가 경로에 있도록 기계에 CUDA를 별도로 설치해야 합니다. 위의 채널 순서(conda-forge 이전에 nvidia)가 중요합니다.

JAX에 의해 사용되는 CUDA의 릴리스를 오버라이드하고 싶거나, GPU가 없는 기계에 CUDA 빌드를 설치하려면, conda-forge 웹사이트의 Tips & tricks 섹션의 지침을 따르세요.

자세한 내용은 conda-forge의 jaxlib 및 jax 저장소를 참고하세요.

소스로부터 JAX 빌드#

소스로부터 JAX 빌드를 참고하세요.

이전 jaxlib 휠 설치하기#

Python 패키지 인덱스의 저장 공간 제한으로 인해, 우리는 주기적으로 http://pypi.org/project/jax에서 이전 jaxlib 휠을 제거합니다. 이들은 여기 있는 URL을 통해 직접 설치할 수 있습니다; 예를 들어:

# 휠 아카이브를 통해 CPU용 jaxlib 설치
pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html

# jaxlib 0.3.25 CPU 휠을 직접 설치
pip install jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html

특정 이전 GPU 휠을 사용하려면 jax_cuda_releases.html URL을 사용하세요; 예를 들어

pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Contents