JAX: Difference between revisions
Jump to navigation
Jump to search
(Fix: remove --- horizontal lines (7 removed)) |
({{Status}}/{{TOC}} 템플릿 제거) |
||
| Line 1: | Line 1: | ||
== Overview == | == Overview == | ||
| Line 17: | Line 8: | ||
* 왜 필요한가? - HPC 및 서버 환경에서 필수 개념 | * 왜 필요한가? - HPC 및 서버 환경에서 필수 개념 | ||
* 언제 사용하는가? - 서버 구성, 성능 튜닝, 문제 해결 시 | * 언제 사용하는가? - 서버 구성, 성능 튜닝, 문제 해결 시 | ||
== Purpose == | == Purpose == | ||
| Line 26: | Line 16: | ||
* Scope: JAX의 개념, 사용법, 설정 | * Scope: JAX의 개념, 사용법, 설정 | ||
* Non-goals: 다른 주제로의 확장 | * Non-goals: 다른 주제로의 확장 | ||
== Key Concepts == | == Key Concepts == | ||
| Line 39: | Line 28: | ||
| [[Linux]], [[Server]] | | [[Linux]], [[Server]] | ||
|} | |} | ||
== Detailed Explanation == | == Detailed Explanation == | ||
| Line 46: | Line 34: | ||
[https://jax.readthedocs.io/en/latest/index.html JAX] stands for "Just Another XLA" where XLA stands for Accelerated Linear Algebra, which is a domain-specific compiler for linear algebra operations developed by Google. | [https://jax.readthedocs.io/en/latest/index.html JAX] stands for "Just Another XLA" where XLA stands for Accelerated Linear Algebra, which is a domain-specific compiler for linear algebra operations developed by Google. | ||
JAX provides a familiar NumPy-like API for array operations and supports automatic differentiation for arbitrary code, including control flow statements like loops and conditionals. This makes it easy to write and optimize complex numerical computations, including deep learning models, with high performance on CPU, GPU, and TPU hardware. | JAX provides a familiar NumPy-like API for array operations and supports automatic differentiation for arbitrary code, including control flow statements like loops and conditionals. This makes it easy to write and optimize complex numerical computations, including deep learning models, with high performance on CPU, GPU, and TPU hardware. | ||
JAX also provides a functional programming paradigm, enabling users to compose functions in a pure, side-effect-free way. This can make it easier to reason about the behavior of a program and to reuse code. so, org said '''JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research.''' | JAX also provides a functional programming paradigm, enabling users to compose functions in a pure, side-effect-free way. This can [[make]] it easier to reason about the behavior of a program and to reuse code. so, org said '''JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research.''' | ||
JAX looks not stable yet and the community seems not moving faster than other famous [[Deep Learning Frameworks|deep learning frameworks]] yet. | JAX looks not stable yet and the community seems not moving faster than other famous [[Deep Learning Frameworks|deep learning frameworks]] yet. | ||
* No [[kernel]] image is available for execution on the device in Google JAX | * No [[kernel]] image is available for execution on the device in Google JAX | ||
| Line 53: | Line 41: | ||
* Performance Issues - Another common problem with Google JAX is that of poor performance. This can manifest in a number of ways, such as slow execution times or poor GPU utilization | * Performance Issues - Another common problem with Google JAX is that of poor performance. This can manifest in a number of ways, such as slow execution times or poor GPU utilization | ||
* JAX 0.4.6 needs [[CUDA]] 11.8 is required for Hopper and Ada Lovelace series GPUs<ref>https://github.com/google/jax/pull/13614/commits/7bd6cc13e0641ce774a364c6dd4b298b0c3ea69a</ref> | * JAX 0.4.6 needs [[CUDA]] 11.8 is required for Hopper and Ada Lovelace series GPUs<ref>https://github.com/google/jax/pull/13614/commits/7bd6cc13e0641ce774a364c6dd4b298b0c3ea69a</ref> | ||
* as of March 2023, JAX has no dataset/dataloader API, nor standard datasets like MNIST. Thus you will have to use either TF or PyTorch for these tasks or implement everything by yourself. | * as of March 2023, JAX has no dataset/dataloader API, nor standard datasets like MNIST. Thus you will have to use either TF or [[PyTorch]] for these tasks or implement everything by yourself. | ||
* JAX is very TPU friendly and has built-in [[support]] for multiple devices | * JAX is very TPU friendly and has built-in [[support]] for multiple devices | ||
* It is still in the 0.x versions, the API might change. | * It is still in the 0.x versions, the API might change. | ||
| Line 59: | Line 47: | ||
* Apart from the TPU, there are few real advantages over PyTorch (or TF). | * Apart from the TPU, there are few real advantages over PyTorch (or TF). | ||
<references/> | <references/> | ||
== Best Practices == | == Best Practices == | ||
| Line 66: | Line 53: | ||
* 공식 문서 참고 | * 공식 문서 참고 | ||
* 테스트 환경에서 먼저 검증 | * 테스트 환경에서 먼저 검증 | ||
== References == | == References == | ||
* [https://wiki.hpcmate.com JAX] | * [https://wiki.hpcmate.com JAX] | ||
== Related Pages == | == Related Pages == | ||
| Line 79: | Line 64: | ||
* [[Hardware]] | * [[Hardware]] | ||
* [[Network]] | * [[Network]] | ||
[[Category:Server]] | [[Category:Server]] | ||
Latest revision as of 13:24, 2 October 2026
Overview
JAX에 대한 기술 문서입니다.
Summary
- 무엇인가? - JAX
- 왜 필요한가? - HPC 및 서버 환경에서 필수 개념
- 언제 사용하는가? - 서버 구성, 성능 튜닝, 문제 해결 시
Purpose
이 문서가 존재하는 이유
- Goal: JAX에 대한 기술 정보 제공
- Scope: JAX의 개념, 사용법, 설정
- Non-goals: 다른 주제로의 확장
Key Concepts
| Concept | Description | Related |
|---|---|---|
| JAX | HPC/서버 환경에서 중요한 기술 개념 | Linux, Server |
Detailed Explanation
섬네일|jaxlogo JAX stands for "Just Another XLA" where XLA stands for Accelerated Linear Algebra, which is a domain-specific compiler for linear algebra operations developed by Google. JAX provides a familiar NumPy-like API for array operations and supports automatic differentiation for arbitrary code, including control flow statements like loops and conditionals. This makes it easy to write and optimize complex numerical computations, including deep learning models, with high performance on CPU, GPU, and TPU hardware. JAX also provides a functional programming paradigm, enabling users to compose functions in a pure, side-effect-free way. This can make it easier to reason about the behavior of a program and to reuse code. so, org said JAX is NumPy on the CPU, GPU, and TPU, with great automatic differentiation for high-performance machine learning research. JAX looks not stable yet and the community seems not moving faster than other famous deep learning frameworks yet.
- No kernel image is available for execution on the device in Google JAX
- Compatibility Issues- Google JAX is a relatively new library and as such, it may not be compatible with all existing libraries. This can lead to compatibility issues when trying to use JAX alongside other libraries, such as NumPy.
- Problem: Memory Leak Issues - One of the most commonly reported problems with Google JAX is that of memory leaks. These occur when the program continues to hold on to memory resources even after they are no longer needed. This can lead to the program using up an excessive amount of memory and potentially crashing.
- Performance Issues - Another common problem with Google JAX is that of poor performance. This can manifest in a number of ways, such as slow execution times or poor GPU utilization
- JAX 0.4.6 needs CUDA 11.8 is required for Hopper and Ada Lovelace series GPUs[1]
- as of March 2023, JAX has no dataset/dataloader API, nor standard datasets like MNIST. Thus you will have to use either TF or PyTorch for these tasks or implement everything by yourself.
- JAX is very TPU friendly and has built-in support for multiple devices
- It is still in the 0.x versions, the API might change.
- Functional programming can be annoying for beginners.
- Apart from the TPU, there are few real advantages over PyTorch (or TF).
Best Practices
- 최신 버전 사용 권장
- 공식 문서 참고
- 테스트 환경에서 먼저 검증
References
Related Pages
Knowledge Graph
Related