TPU와 GPU 위의 JAX

JAX는 같은 코드를 CPU, GPU, TPU 어디서나 실행할 수 있다고 약속해. 실전에서는 하드웨어마다 사용 방식에 약간의 차이가 있어.

장치 확인

import jax
print(jax.devices())
# [CpuDevice(id=0)]                          ← CPU only
# [GpuDevice(id=0, process_index=0), ...]    ← CUDA GPU
# [TpuDevice(id=0, ...), ...]                 ← TPU

print(jax.default_backend())   # 'cpu' / 'gpu' / 'tpu'
print(jax.device_count())      # 8 (8 TPU cores)

설치 차이

# CPU
pip install -U jax

# CUDA 12 (NVIDIA)
pip install -U "jax[cuda12]"
# 또는 specific
pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# TPU (Google Cloud)
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# Apple Silicon (실험적)
pip install -U jax-metal

Colab에서

Colab, JAX TPU가 무료야. 실행 환경에서 런타임 유형 변경을 열고 TPU를 선택해. 그러면:

import jax
print(jax.devices())   # 8 개 TPU core
print(jax.device_count())   # 8

메모리 관리

JAX는 기본값으로, 프로세스를 시작할 때 GPU 메모리의 90%를 미리 할당 (메모리 fragmentation 방지). PyTorch와 같이 쓰면 충돌할 수 있어:

export XLA_PYTHON_CLIENT_PREALLOCATE=false
# 또는 specific fraction
export XLA_PYTHON_CLIENT_MEM_FRACTION=0.5

Python 안에서:

import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
import jax

다중 GPU

# 자동 — 모든 GPU 가 visible
import jax
print(jax.device_count())   # 4 (4 GPU)

# 특정 device 만 사용
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"   # GPU 0, 1 만

TPU의 특성

TPU pod는 여러 칩으로 구성되고 각 칩에는 여러 코어가 있어.
TPU의 systolic 배열 덕분에 행렬 곱셈이 매우 빨라.
매우 큰 배치 크기에 유리해.
bfloat16이 기본 데이터 타입이야.
동적 shape에 약해서 shape이 바뀌면 다시 컴파일해.

GPU의 특성

NVIDIA CUDA, 가장 많이 쓰는 백엔드.
유연성이 높고 커널을 직접 작성할 수 있음 (Pallas, Triton).
멀티프로세스와 멀티호스트 학습에서는 NCCL을 사용해.
보통은 TPU보다 메모리가 적어.

multi-호스트 학습

여러 머신이 각각 여러 GPU를 사용하는 구성이야. jax.distributed.initialize()로 클러스터를 설정해:

import jax

jax.distributed.initialize(
    coordinator_address="10.0.0.1:1234",
    num_processes=4,
    process_id=process_idx,   # 각 host 에서 다른 값
)

# 이제 모든 host 의 device 가 jax.devices() 에 보임
# pmap / sharding 이 자동으로 cross-host 통신 처리

SLURM이나 Kubernetes에서 실행할 때는 각 프레임워크가 권장하는 시작 방식을 따라.

🔋 하드웨어별 sweet spot

(1) 학습과 연구에서는 TPU가 단일 작업당 비용 효율적인 경우가 많아 (Google Cloud). (2) 추론과 배포에서는 GPU가 더 친숙한 경우가 많아 (NVIDIA Triton 등 범용 인프라). (3) 온디바이스 / 모바일, Apple Silicon의 jax-metal 또는 jax2tf → TF Lite. (4) CPU만 있어도 작은 모델을 학습하고 추론할 수 있으니 하드웨어가 부족할 때 시작점으로 삼아.

같은 코드가 하드웨어만 바꾸면 그대로 도는 게 JAX의 약속. 완전히 자동인 것은 아니므로 비동기 디스패치와 메모리 관리처럼 하드웨어별로 약간의 튜닝이 필요해.

Code

import jax

# Check available devices
print(jax.devices())
# [CudaDevice(id=0), CudaDevice(id=1), ...]

# JAX automatically uses GPU if available
x = jax.numpy.ones((1000, 1000))
# x is already on GPU — no .to('cuda') needed!

# Multi-GPU: use sharding (see Track 12)
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

devices = jax.devices()  # all GPUs
mesh = Mesh(devices, ('data',))

# Data parallelism: shard batch across GPUs
data_sharding = NamedSharding(mesh, P('data'))
batch = jax.device_put(x, data_sharding)

# As of JAX 0.6.0: requires CUDA 12.8+
# pip install jax[cuda12]

# On Google Cloud TPU VMs:
# pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# TPU pods: automatic multi-host setup
print(jax.device_count())  # e.g., 8 for a v4-8, 32 for a v4-32

# TPU-specific tips:
# 1. Use bfloat16 — TPUs have native bfloat16 support
x_bf16 = x.astype(jax.numpy.bfloat16)

# 2. Pad batch sizes to multiples of 128 (TPU-friendly)
# 3. Avoid scalar operations — TPUs are designed for large tensor ops
# 4. Use jax.profiler for TPU-specific profiling

# Common performance pitfalls:

# 1. UNNECESSARY RECOMPILATION
# BAD: different shapes cause recompilation
for batch in variable_size_batches:
    result = jax.jit(fn)(batch)  # recompiles every new shape!

# GOOD: pad to fixed size
max_batch_size = 256
for batch in batches:
    padded = pad_to_size(batch, max_batch_size)
    result = jax.jit(fn)(padded)  # compiled once, reused

# 2. HOST-DEVICE TRANSFER
# BAD: pulling values back to CPU in a loop
for step in range(1000):
    loss = train_step(params, batch)
    print(float(loss))  # blocks! transfers to CPU every step

# GOOD: only transfer periodically
for step in range(1000):
    loss = train_step(params, batch)
    if step % 100 == 0:
        print(float(loss))  # transfer only every 100 steps

# 3. USE jax.block_until_ready() for timing
import time
x = jax.numpy.ones((1000, 1000))
start = time.time()
y = x @ x
y.block_until_ready()  # wait for computation to finish
print(f"Time: {time.time() - start:.4f}s")

Exercise

GPU나 Colab TPU를 사용할 수 있다면 jax.devices()로 장치를 확인해. 이전 스크립트 하나를 코드 변경 없이 옮겨 같은 행렬 곱셈을 CPU, GPU, TPU에서 측정해. 성능 차이를 과장하지 말고 그대로 해석한 뒤 JAX의 이식성이 주는 가치를 적어.

🔋 하드웨어별 sweet spot

Code

External links

Exercise

Progress

댓글 0