jit의 일: 추적, 컴파일, 캐싱

jax.jit이 실제로 무슨 일을 하는지 한 번 살펴보면 모든 게 명확해져. 과정은 세 단계야.

1. 추적, 첫 호출 때 함수를 추상적인 Tracer 객체로 한 번 실행해. 어떤 연산이 어떤 순서로 일어나는지 IR (XLA HLO)로 기록해.

2. 컴파일, IR을 XLA 컴파일러로 넘김. 가속기별 기계어 생성해.

3. 캐시, 컴파일 결과를 (입력 shape, dtypes, 정적 args) 키로 캐시에 저장해. 다음 호출 때 같은 키면 재사용해.

import jax
import jax.numpy as jnp
import time

@jax.jit
def f(x):
    return jnp.sum(x ** 2 + jnp.sin(x))

x = jnp.arange(1_000_000.0)

# 첫 호출 — trace + compile + run
t = time.time()
y = f(x).block_until_ready()
print(f"첫 호출: {time.time()-t:.3f}s")  # ~ 0.3s

# 두 번째 — cache hit, 빠름
t = time.time()
y = f(x).block_until_ready()
print(f"두 번째: {time.time()-t:.3f}s")  # ~ 0.001s

# 다른 shape — 새 trace + compile
y = f(jnp.arange(2_000_000.0)).block_until_ready()

캐시 키의 핵심:

Shape: (1000,)와 (2000,)는 다른 키 → 다시 컴파일해.
Dtype: float32와 float64는 다른 키.
정적 args: static_argnames로 표시한 인자의 값 변화도 다시 컴파일해.
장치: 같은 함수도 CPU와 GPU는 별도 캐시에 저장해.

실험으로 보면:

def trace_count_demo():
    n_traces = 0
    @jax.jit
    def f(x):
        nonlocal n_traces
        n_traces += 1   # trace 시점에만 +1
        return x * 2

    f(jnp.arange(10.))   # trace + compile + run
    f(jnp.arange(10.))   # cache hit
    f(jnp.arange(20.))   # 다른 shape → recompile
    f(jnp.arange(20., dtype=jnp.float64))  # 다른 dtype → recompile
    print(f"trace 횟수: {n_traces}")

⚡ "컴파일 한 번, 호출 여러 번" 모델

JAX의 성능 모델 핵심, 함수를 처음 한 번 컴파일하고 그 후 수천 번 빠르게 부른다. 학습 루프라면 train_step 함수가 첫 step에서 컴파일되고, 나머지 1만 스텝은 캐시된 코드를 호출해. 컴파일 비용 (한 번 1초)은 학습 시간 (1시간)에 비해 무시할 만해. 그래서 jit은 거의 항상 이득이야.

중요한 함정, shape이 호출마다 바뀌면 매번 다시 컴파일해. 학습 마지막 배치가 미세하게 작은 거 (drop_last=False), 그것 때문에 매 에포크마다 컴파일 한 번 더 할 수도 있어. 해결: 패딩해서 shape을 동일하게 유지해.

캐시를 비우려면 jax.clear_caches()를 호출할 수 있지만 보통은 필요 없어. 메모리가 부족할 때만 사용해.

Code

import jax
import jax.numpy as jnp

def slow_fn(x):
    """Each operation launches a separate kernel."""
    y = jnp.sin(x)
    z = jnp.cos(x)
    return jnp.sum(y * z + y ** 2)

# Wrap with jit
fast_fn = jax.jit(slow_fn)

# Or use as a decorator
@jax.jit
def fast_fn_v2(x):
    y = jnp.sin(x)
    z = jnp.cos(x)
    return jnp.sum(y * z + y ** 2)

x = jnp.ones(10000)
result = fast_fn(x)  # First call: trace + compile + execute
result = fast_fn(x)  # Second call: execute cached compiled code (much faster)

import jax
import jax.numpy as jnp

@jax.jit
def add_one(x):
    print("TRACING!")  # This print helps us see when tracing happens
    return x + 1

# Call 1: traces + compiles + executes
add_one(jnp.array([1.0, 2.0]))  # Prints "TRACING!"

# Call 2: same shape → uses cache, no retracing
add_one(jnp.array([3.0, 4.0]))  # No print — cached!

# Call 3: different shape → must retrace
add_one(jnp.array([1.0, 2.0, 3.0]))  # Prints "TRACING!" — new shape

Exercise

같은 행렬 곱셈을 세 방식으로 측정해. jit을 쓰지 않은 jnp 실행, jit의 첫 실행, 캐시를 재사용하는 이후 실행을 행렬 크기 64, 256, 1024, 2048에서 각각 비교해. 결과를 그래프나 표로 정리하고, 컴파일 비용을 감수할 만해지는 지점을 설명해.

⚡ "컴파일 한 번, 호출 여러 번" 모델

Code

External links

Exercise

Progress

댓글 0