JAX를 실무 흐름으로 이해하기
JAX는 NumPy 스타일 API에 자동 미분, JIT 컴파일, 벡터화, 분산 실행을 결합한 고성능 연구/학습 프레임워크입니다. 이 가이드는 개념을 나열하기보다, 실제 프로젝트에서 판단해야 하는 순서대로 내용을 따라갈 수 있게 구성했습니다.
ML / AI 고성능 수치 연산 가이드
JAX 완전 가이드
JAX는 NumPy 스타일 API에 자동 미분, JIT 컴파일, 벡터화, 분산 실행을 결합한 고성능 연구/학습 프레임워크입니다.
가이드 사용법
AI / LLM 시스템
모델과 프롬프트만 보지 않고, 데이터 흐름, 평가, 배포 이후의 운영 지표까지 한 번에 연결해서 봅니다.
구조 다이어그램
글로 읽은 내용을 머릿속에 오래 남기려면 먼저 흐름을 그림으로 잡는 편이 좋습니다. 아래 두 그림은 JAX를 학습할 때 계속 되돌아볼 수 있는 기준 지도입니다.
학습 흐름
다이어그램 렌더링 중…
아키텍처 관점
다이어그램 렌더링 중…
JAX란?
JAX를 처음 펼칠 때는 세부 명령보다 큰 그림이 먼저입니다. 이 섹션에서는 앞으로 배울 개념들이 어떤 문제를 풀기 위해 등장했는지부터 잡아봅니다.
JAX는 NumPy와 비슷한 문법으로 GPU/TPU 가속, 자동 미분, JIT 컴파일을 제공하는 라이브러리입니다. 대규모 모델 연구와 고성능 수치 실험에 자주 사용됩니다.
| 기능 | 설명 |
|---|---|
| grad | 함수의 gradient를 자동 계산합니다. |
| jit | 함수를 XLA로 컴파일해 빠르게 실행합니다. |
| vmap | 배치 차원을 자동 벡터화합니다. |
| pmap | 여러 장치로 병렬 실행합니다. |
설치
여기서는 설치을 실제 코드와 함께 확인합니다. 예제를 그대로 따라 하기보다, 입력과 출력, 그리고 바뀌기 쉬운 부분이 어디인지 보면서 읽어보세요.
CPU 전용 환경에서는 jax[cpu] 패키지만으로 충분하며, GPU/TPU를 쓰려면 별도의 CUDA 빌드를 지정해 설치해야 합니다. jax.devices()로 실제 인식된 연산 장치를 확인하는 것이 첫 단계입니다.
BASH
uv venv
source .venv/bin/activate
uv pip install "jax[cpu]"
python -c "import jax; print(jax.devices())"자동 미분
여기서는 자동 미분을 실제 코드와 함께 확인합니다. 예제를 그대로 따라 하기보다, 입력과 출력, 그리고 바뀌기 쉬운 부분이 어디인지 보면서 읽어보세요.
jax.grad는 순수 함수를 입력받아 그 함수의 도함수를 계산하는 새로운 함수를 반환합니다. 원본 함수를 직접 수정할 필요 없이, 미분이 필요한 지점에서 grad()로 감싸기만 하면 됩니다.
PYTHON
import jax
import jax.numpy as jnp
def loss(w):
return jnp.sum((w - 3.0) ** 2)
grad_loss = jax.grad(loss)
print(grad_loss(jnp.array([1.0, 2.0, 4.0])))JIT 컴파일
여기서는 JIT 컴파일을 실제 코드와 함께 확인합니다. 예제를 그대로 따라 하기보다, 입력과 출력, 그리고 바뀌기 쉬운 부분이 어디인지 보면서 읽어보세요.
자주 호출되는 순수 함수는 jit로 컴파일해 실행 비용을 줄일 수 있습니다.
PYTHON
@jax.jit
def matmul(a, b):
return a @ b
x = jnp.ones((1024, 1024))
print(matmul(x, x).shape)벡터화
여기서는 벡터화을 실제 코드와 함께 확인합니다. 예제를 그대로 따라 하기보다, 입력과 출력, 그리고 바뀌기 쉬운 부분이 어디인지 보면서 읽어보세요.
vmap은 단일 샘플 함수를 배치 함수로 확장합니다.
PYTHON
def predict(w, x):
return jnp.dot(w, x)
batched_predict = jax.vmap(predict, in_axes=(None, 0))JAX 실무 설계
JAX 실무 설계은 선택지가 갈리는 지점입니다. 표를 기준으로 각 방법의 쓰임새와 운영상의 차이를 비교해두면 이후 판단이 훨씬 쉬워집니다.
JAX는 순수 함수와 불변 데이터를 전제로 설계해야 합니다. random key, model params, optimizer state를 명시적으로 전달해야 jit/vmap/pmap 확장이 쉬워집니다.
| 결정 지점 | 확인 질문 | 실무 기준 |
|---|---|---|
| 경계 | JAX 코드에서 바뀌기 쉬운 부분은 어디인가? | 입출력, 설정, 외부 연동, 핵심 규칙을 분리합니다. |
| 상태 | 상태가 어디서 생성되고 어디서 사라지는가? | 상태 소유자와 수명 주기를 코드로 드러냅니다. |
| 장애 | 실패했을 때 호출자는 무엇을 받는가? | timeout, fallback, error contract를 먼저 정합니다. |
JAX 운영 기준
이 섹션은 JAX 운영 기준을 실무 관점에서 정리합니다. 개념을 외우기보다, 어떤 상황에서 이 기준을 꺼내 쓸지에 초점을 맞춰보세요.
jit compile time과 실행 시간을 분리해서 측정해야 합니다. shape가 자주 바뀌면 재컴파일 비용이 커지므로 batch shape 안정화가 중요합니다.
JAX 검증 전략
JAX 검증 전략은 선택지가 갈리는 지점입니다. 표를 기준으로 각 방법의 쓰임새와 운영상의 차이를 비교해두면 이후 판단이 훨씬 쉬워집니다.
수치 안정성, dtype 차이, gradient check, deterministic PRNG key 사용을 검증해야 합니다.
| 품질 축 | 검증 방법 | 완료 기준 |
|---|---|---|
| 정확성 | 정상/실패 케이스를 자동화합니다. | 핵심 시나리오가 재현 가능하게 통과합니다. |
| 회귀 방지 | 버그 수정 시 동일 케이스를 테스트로 남깁니다. | 같은 장애가 다시 배포되지 않습니다. |
| 운영성 | 로그, 메트릭, 알림을 확인합니다. | 문제가 생겼을 때 원인 추적 경로가 있습니다. |