Ray on TPU 운영법: topology 한 줄이 LLM 서빙 장애를 막는 이유

요약: Google Developers Blog의 Ray on TPU Part 2는 TPU를 쓰는 LLM 서빙·데이터·학습 워크로드에서 가장 자주 터지는 문제를 짚습니다. 핵심은 TPU 칩이 slice라는 고정된 묶음으로 연결되어 있고, 멀티호스트 모델은 하나의 온전한 slice 위에 올라가야 한다는 점입니다. Ray Serve에서는 accelerator_config.topology 한 줄이 장애와 정상 배포를 가르는 경우가 있습니다.

TPU는 GPU 여러 장과 다르게 배치된다

GPU 클러스터를 쓰던 팀이 TPU로 넘어갈 때 가장 먼저 헷갈리는 부분은 배치 방식입니다. GPU에서는 필요한 카드 수를 맞추는 감각이 강합니다. 하지만 TPU는 칩들이 slice라는 고정된 그룹으로 묶이고, 같은 slice 안의 호스트들이 ICI라는 고속 링크를 공유합니다.

멀티호스트 모델은 이 링크가 중요합니다. 텐서 병렬이나 샤딩된 모델의 worker들이 서로 통신해야 하는데, 서로 다른 slice에 흩어지면 첫 collective부터 끝나지 않을 수 있습니다. 문제는 이 상황이 명확한 크래시로 보이지 않을 때가 많다는 점입니다. 배포가 DEPLOYING에 머물고 TPU-hours만 태우는 식으로 나타납니다.

Google 글은 이 문제를 “topology field makes the difference”라고 설명합니다. Ray Core에는 slice_placement_group() 같은 primitive가 있고, GKE Ray Operator add-on은 slice host를 프로비저닝하고 라벨링합니다. 상위 라이브러리인 Ray Serve, Ray Data, Ray Train은 결국 같은 패턴을 따릅니다. topology를 선언하고 Core가 slice를 예약하게 하는 것입니다.

Ray Serve에서 topology를 빼먹으면 생기는 일

LLM 서빙은 많은 팀이 TPU를 처음 쓰는 구간입니다. 큰 모델을 올리려면 여러 칩을 묶어야 하고, vLLM 같은 서빙 엔진을 통해 throughput을 뽑아야 합니다. Ray Serve는 autoscaling, load balancing, multi-model composition을 제공합니다. TPU에서는 vLLM을 통해 LLM 서빙을 지원합니다.

문제는 모델이 한 호스트에 안 들어가는 경우입니다. 예를 들어 16칩에 걸친 tensor parallel 모델이라면 worker들이 같은 slice에 있어야 합니다. 이때 설정에 accelerator_config와 topology: "4x4" 같은 선언이 필요합니다. 여기서 4x4는 칩 개수 숫자가 아니라 slice shape입니다.

이 설정이 있으면 Serve의 TPU backend가 replica 시작 시 slice placement group을 만들고, worker를 하나의 ICI mesh 위에 둡니다. 설정이 없으면 per-chip bundle처럼 흩어질 수 있고, worker들이 서로 닿지 않는 slice에 배치될 수 있습니다. 결과는 배포가 끝나지 않는 장애입니다. 로그만 보면 모델 로딩 문제인지 네트워크 문제인지 헷갈리기 쉽습니다.

Ray Data의 iter_jax_batches()가 중요한 이유

빠른 accelerator는 데이터가 느리면 아무 소용이 없습니다. TPU 학습이나 대량 batch inference에서 흔한 병목은 모델이 아니라 입력 파이프라인입니다. 호스트 메모리에서 NumPy 배열을 만들고, 다시 JAX 배열로 바꾸고, device shard를 맞추는 과정이 step time을 잡아먹습니다.

Ray Data의 iter_jax_batches()는 이 구간을 줄입니다. Parquet 같은 데이터셋을 읽어 JAX array 형태, device-sharded batch로 넘겨줍니다. 학습 step이나 batch inference job은 바로 이 batch를 받아 처리할 수 있습니다. 마지막 batch가 batch size로 딱 나누어떨어지지 않을 때도 drop, pad, raise 같은 정책을 명시할 수 있습니다.

운영 관점에서 이 기능은 “3시간 돌다가 shape error로 죽는” 유형의 장애를 줄이는 데 의미가 큽니다. 대규모 학습과 후처리는 비용이 큰 작업입니다. 실패가 늦게 날수록 비용도 커집니다. ragged final batch 처리를 명시하고, device sharding을 프레임워크에 맡기는 것은 단순 편의가 아니라 비용 관리입니다.

JaxTrainer에서 import 위치까지 신경 써야 하는 이유

Ray Train의 JaxTrainer는 JAX 기반 TPU 학습 루프를 Ray Train 방식으로 실행하게 합니다. checkpointing, fault tolerance, multi-slice scale-out 같은 기능을 가져갈 수 있고, ScalingConfig에서 use_tpu=True, topology="4x4", accelerator_type="TPU-V6E"처럼 배치를 선언합니다.

Google 글에서 실무적으로 중요한 팁은 import jax 위치입니다. JAX import를 모듈 상단에서 하지 말고 worker 함수 내부에서 하라고 설명합니다. 각 worker가 자기 TPU context 안에서 JAX를 초기화해야 하기 때문입니다. 모듈 스코프에서 import하면 첫 step 전에 device-init 오류와 싸울 수 있습니다.

이런 디테일은 문서 한 줄처럼 보이지만 운영에서는 반나절을 날립니다. 특히 TPU는 익숙하지 않은 팀이 많아 에러 메시지를 보고 바로 원인을 찾기 어렵습니다. 따라서 TPU 학습 코드는 “일반 Python 스타일”보다 runtime 초기화 순서를 우선해야 합니다.

모니터링은 TPU utilization과 memory를 같이 봐야 한다

Ray는 TPU 지원의 일부로 rayproject/ray:*-tpu 이미지를 제공합니다. JAX/TPU stack, flax, optax, orbax-checkpoint, profiling tooling이 들어 있어 환경 조립 비용을 줄입니다. 또한 Ray Dashboard는 Cluster tab에서 CPU, GPU뿐 아니라 TPU utilization과 memory를 보여줍니다.

서빙 장애를 줄이려면 대시보드 지표를 배포 체크리스트에 넣어야 합니다. deployment 상태, replica 수, TPU utilization, memory, request latency, queue length를 같이 봐야 합니다. utilization이 낮은데 latency가 높으면 데이터 입력이나 request queue 문제일 수 있습니다. memory가 한계에 가까우면 batch size, tensor parallel 설정, 모델 shard 전략을 봐야 합니다.

프로파일러도 중요합니다. Ray가 제공하는 JAX profiler 연결을 통해 per-worker 병목을 확인할 수 있습니다. TPU를 쓰는 이유는 비용 대비 처리량인데, utilization을 보지 않으면 “비싼 accelerator를 놀리는” 상황이 됩니다.

TPU 도입 전에 계산해야 할 비용 구조

TPU가 GPU보다 항상 싼 것은 아닙니다. 더 정확한 표현은 “특정 워크로드에서 더 구하기 쉽고 비용 효율적일 수 있다”입니다. 계산 기준은 모델 크기, batch size, 지연 시간 요구사항, accelerator availability, 운영팀 숙련도입니다.

LLM inference라면 p95 latency와 tokens/sec를 함께 봐야 합니다. 학습이라면 step time, checkpoint overhead, 실패 재시작 비용을 봐야 합니다. batch inference라면 데이터 로딩 throughput과 accelerator idle time이 중요합니다. TPU slice를 예약했는데 데이터 파이프라인이 못 따라오면 실제 비용은 높아집니다.

따라서 첫 도입은 작은 end-to-end 샘플로 시작해야 합니다. Google이 제공하는 GKE Ray TPU get-started 예제처럼 serve, data, train 중 하나를 그대로 실행하고, 팀이 배포·모니터링·삭제까지 해보는 것이 좋습니다. 바로 프로덕션 모델을 옮기면 배치와 네트워크, 이미지, 권한, 비용을 동시에 디버깅하게 됩니다.

실행 체크리스트

• TPU를 “칩 개수”가 아니라 slice topology 기준으로 이해한다.
• 멀티호스트 LLM 서빙에서는 Ray Serve 설정에 accelerator_config.topology를 반드시 명시한다.
• topology="4x4" 같은 값은 칩 수가 아니라 slice shape임을 팀 문서에 적는다.
• 배포가 DEPLOYING에 오래 머물면 모델 코드보다 slice placement부터 확인한다.
• Ray Data의 iter_jax_batches()로 JAX-native, device-sharded batch를 공급한다.
• ragged final batch 정책을 drop, pad, raise 중 하나로 명시한다.
• JaxTrainer에서는 import jax를 worker 함수 내부로 옮긴다.
• rayproject/ray:*-tpu 이미지를 기본 후보로 써 환경 조립 리스크를 줄인다.
• Ray Dashboard에서 TPU utilization, memory, replica 상태, latency를 같이 본다.
• 프로덕션 이전에 get-started 예제로 serve, data, train 중 하나를 끝까지 실행해 삭제까지 검증한다.