2024년 Google은 JAX 기반 머신러닝 도구를 묶은 JAX AI Stack을 발표했어. Flax NNX (모델), Optax (옵티마이저), Orbax (체크포인트), Grain (데이터 로더), 기타. 한 번에 설치:
pip install jax-ai-stack
새 프로젝트를 시작할 때 이 한 줄로 호환성이 검증된 라이브러리 묶음을 설치해 버전 충돌을 줄일 수 있어.
Pallas, JAX의 사용자 정의 커널 언어
Pallas는 표준 JAX 연산으로 표현하기 어려운 하드웨어 친화적 커널을 작성하는 도구로, PyTorch 생태계의 Triton과 비슷해.
from jax.experimental import pallas as pl
@pl.pallas_call(out_shape=jax.ShapeDtypeStruct((1024,), jnp.float32))
def add_kernel(x_ref, y_ref, out_ref):
'''단순 vector add — Pallas 식으로'''
x = x_ref[...]
y = y_ref[...]
out_ref[...] = x + y
# 사용 — 표준 JAX 처럼
x = jnp.ones(1024)
y = jnp.ones(1024)
result = add_kernel(x, y)
위 예제는 단일 프로그램 블록에서 실행되는 벡터 덧셈이야. 실전에서는 더 복잡한 융합 커널, 특히 attention 변형에 사용해:
# Flash Attention — Pallas 로 구현 (의사 코드)
@pl.pallas_call(...)
def flash_attention_kernel(q_ref, k_ref, v_ref, out_ref):
# block-wise loading
# softmax 의 online algorithm
# weighted sum 를 in-place 누적
...
왜 Pallas가 필요한가:
- 메모리 계층 활용: HBM과 SRAM 사이의 이동을 명시적으로 제어해. attention의 입출력 비용은 메모리 계층에 크게 좌우돼.
- 연산 융합의 한계 돌파: XLA의 자동 연산 융합이 놓치는 패턴을 직접 융합해.
- 새 하드웨어 활용: TPU의 sparse core와 GPU의 tensor core 같은 기능을 직접 활용해.
실전, Pallas가 빛나는 곳
| 분야 | Pallas의 역할 |
|---|---|
| Flash Attention | 긴 sequence의 attention 메모리 감소 |
| Sparse layers | 희소 행렬 곱셈을 위한 효율적인 커널 |
| Quantization | int8 / fp8 행렬 연산 |
| MoE (Mixture of Experts) | scatter/gather의 효율 패턴 |
| LSH / approximate ops | 해시 기반 샘플링, near-neighbor |
대부분의 사용자에게는 Pallas가 필요하지 않아. 표준 JAX op으로 충분. Pallas가 결정적인 경우는 매우 큰 모델, 특수 하드웨어 활용, 새 알고리즘 연구야.
JAX 진영의 다른 보조 라이브러리
- chex, 단위 test, 디버깅, dataclass
- clu (Common 루프 Utils), metric 누적, summary writer
- einshape / einops, 텐서 shape 조작
- xpilot, distributed coordinator (Google 내부 → 외부 release)
- distrax, probability distribution (DeepMind)
🛠 프로덕션 stack 추천
(1) 프레임워크: Flax NNX 또는 Equinox. (2) 옵티마이저: Optax. (3) 체크포인트: Orbax. (4) 데이터: Grain (또는 PyTorch DataLoader). (5) 로깅: wandb. (6) test: chex. (7) 기타: jax-ai-stack의 의존성. Pallas는 표준 op이 막힐 때만. 이 stack으로, Llama-클래스 모델 학습할 수 있어.
Pallas를 비유하자면 JAX의 인라인 어셈블리야. 99%의 코드는 안 쓰지만, 나머지 1%에서는 결정적일 수 있어. 존재와 역할을 알아 두면 충분해.