성능 프로파일링과 디버깅

학습이 느릴 때는 병목이 어디인지 먼저 찾아야 해. JAX는 이를 위한 강력한 프로파일링 도구를 제공해.

jax.profiler, TensorBoard 추적

import jax
import jax.numpy as jnp

# trace 시작
jax.profiler.start_trace("/tmp/profile-data")

# 측정할 코드
for _ in range(100):
    state, loss = train_step(state, batch)
    state.params["w"].block_until_ready()

# trace 종료
jax.profiler.stop_trace()

그 후, TensorBoard로 열기:

pip install tensorboard tensorboard-plugin-profile
tensorboard --logdir /tmp/profile-data
# Profile 탭 → trace_viewer 또는 op_profile

각 연산의 시간, 장치 활용도, 호스트와 장치의 통신 등을 시각화해. 학습 스텝의 병목을 찾는 데 결정적인 자료야.

TraceAnnotation으로 코드 구간에 이름 붙이기

@jax.profiler.annotate_function
def model_forward(params, x):
    return ...

# 또는 context manager
with jax.profiler.TraceAnnotation("data_loading"):
    batch = next(loader)

with jax.profiler.TraceAnnotation("train_step"):
    state = train_step(state, batch)

프로파일 결과에서 이름 붙인 구간이 따로 표시되므로 어디가 느린지 시각적으로 확인할 수 있어.

perfetto / Chrome 추적

TensorBoard가 무거우면 더 가벼운 perfetto:

jax.profiler.start_trace("/tmp/perfetto-trace")
# code
jax.profiler.stop_trace()

https://ui.perfetto.dev/에서 직접 추적 파일 열 수 있어.

컴파일 시간 분석

import time

@jax.jit
def step(x):
    # 복잡한 계산
    return ...

# 첫 호출 — compile
t = time.time()
y = step(x).block_until_ready()
print(f"compile + first run: {time.time()-t:.2f}s")

# 두 번째 — pure run
t = time.time()
y = step(x).block_until_ready()
print(f"warm run: {time.time()-t:.4f}s")

컴파일이 너무 길면 함수를 더 작게 쪼개거나 Python 루프 펼치기를 scan으로 대체해.

메모리 사용량

# peak memory 측정
peak = jax.live_arrays()
total_bytes = sum(arr.nbytes for arr in peak)
print(f"live array memory: {total_bytes / 1e9:.2f} GB")

OOM 발생 시, 어떤 배열이 살아있는지 확인해. 보통, 체크포인트 안 한 활성값, 큰 데이터셋이 장치에 머물러 있는 경우야.

백엔드 정보

print(jax.devices())
print(jax.default_backend())
print(jax.local_device_count())

# XLA HLO 보기 — compile 결과의 IR
print(jax.xla_computation(step)(x).as_hlo_text())

HLO IR을 보면 어떤 연산으로 컴파일됐는지 정확히 알 수 있어. 연산 융합이 의도대로 되었는지 확인할 수 있어.

📊 프로파일링 우선순위

(1) 학습 스텝 시간 측정 (block_until_ready), 1차 신호. (2) tensorboard 추적, 호스트 vs 장치 시간. (3) 만약 호스트가 병목, 데이터 로더, 호스트 ↔ 장치 전송. (4) 장치가 병목, op 프로파일, 연산 융합 검증해. (5) 컴파일이 길면 함수 분할, scan 사용해. 근거 없는 미세 최적화보다 프로파일 결과에 따라 움직여.

경험상 학습 코드를 처음 작성한 뒤 프로파일링하면 80% 정도의 경우 1.5~3배의 속도 향상을 얻을 수 있어. 그 이후에는 모델 구조나 하드웨어 한계를 살펴봐야 해.

Code

import jax
import jax.numpy as jnp

# Check what's being compiled
@jax.jit
def my_function(x):
    return jnp.sin(x) + jnp.cos(x)

# Inspect the jaxpr (JAX's intermediate representation)
jaxpr = jax.make_jaxpr(my_function)(jnp.ones(3))
print(jaxpr)
# Shows the operations JAX will compile — useful for checking
# that your function isn't doing unexpected work

# Inspect compiled HLO
compiled = jax.jit(my_function).lower(jnp.ones(3)).compile()
print(compiled.cost_analysis())
# Shows estimated FLOPs, memory usage, etc.

import jax

# Profile with TensorBoard integration
jax.profiler.start_trace('/tmp/jax_profile')

# Run your training code
for step in range(100):
    params, loss = train_step(params, batch)

jax.profiler.stop_trace()
# Then: tensorboard --logdir=/tmp/jax_profile

# Or use the context manager
with jax.profiler.trace('/tmp/jax_profile'):
    for step in range(100):
        params, loss = train_step(params, batch)

# Set this to log when JIT recompiles
jax.config.update("jax_log_compiles", True)

# Now you'll see messages like:
# "Compiling my_function (..." every time JIT compiles a new variant

# Check XLA compilation logs for detailed info
# JAX_LOG_COMPILES=1 python my_script.py

# Timing best practices
import time

@jax.jit
def fast_fn(x):
    return jnp.linalg.svd(x, full_matrices=False)

x = jax.random.normal(jax.random.key(0), (1000, 500))

# First call includes compilation time
start = time.time()
result = fast_fn(x)
jax.block_until_ready(result)
print(f"First call (includes compile): {time.time() - start:.4f}s")

# Second call is the real runtime
start = time.time()
result = fast_fn(x)
jax.block_until_ready(result)
print(f"Second call (actual runtime): {time.time() - start:.4f}s")

Exercise

50스텝 학습 실행을 jax.profiler.start_trace로 프로파일링해. trace를 TensorBoard나 Perfetto에서 열고 가장 느린 연산 하나를 찾아. 행렬 곱셈 재배치처럼 작은 변경이라도 하나만 적용해 다시 측정해. 프로파일링이 먼저고 최적화는 그다음이야.

📊 프로파일링 우선순위

Code

External links

Exercise

Progress

댓글 0