Ray on TPU 운영법: topology 한 줄이 LLM 서빙 실패를 가르는 이유

요약: Google Developers Blog가 2026년 7월 24일 공개한 Ray on TPU 2편은 “TPU를 잘 쓰는 법”보다 “TPU에서 LLM 서빙이 조용히 멈추는 이유”를 더 정확히 설명합니다. 핵심 키워드는 Ray on TPU, Ray Serve, TPU topology, JaxTrainer입니다. 실무 개발자 입장에서는 새 기능 소개보다 운영 사고를 줄이는 체크리스트로 읽는 편이 좋습니다.

왜 이 이슈가 중요한가

GPU 기반 LLM 서빙은 이제 많은 팀이 익숙합니다. 문제는 비용과 가용성입니다. 특정 리전에서 GPU가 부족하거나, 대형 모델 추론 비용이 계속 올라가면 TPU가 대안으로 올라옵니다. 다만 TPU는 GPU처럼 “칩 몇 개 잡으면 끝”이 아닙니다. TPU 칩은 slice라는 고정된 묶음으로 연결되고, 멀티 호스트 모델은 하나의 온전한 slice 안에 배치되어야 합니다.

이번 글에서 Google이 반복해서 강조한 지점도 여기입니다. 멀티 호스트 모델의 worker가 서로 다른 slice에 흩어지면 고속 인터커넥트인 ICI로 통신할 수 없습니다. 더 나쁜 점은 이 상황이 항상 명확한 에러로 터지지 않는다는 것입니다. 배포가 DEPLOYING 상태에 오래 머물고, 팀은 모델 코드나 vLLM 설정을 의심하면서 TPU 시간을 태웁니다.

즉 Ray on TPU 업데이트의 실무 가치는 “더 빠르다”가 아니라 “배치, 서빙, 학습에서 같은 placement 원칙을 라이브러리 레벨로 끌어올렸다”에 있습니다.

원인은 topology를 칩 수로 착각하는 데 있다

Ray Serve에서 TPU를 쓸 때 가장 중요한 필드는 accelerator_config.topology입니다. 예를 들어 4x4는 16칩이라는 숫자가 아니라 slice shape입니다. 이 값이 있어야 Ray Serve의 TPU backend가 replica 시작 시점에 slice placement group을 만들고, tensor-parallel worker를 같은 ICI mesh 안에 묶습니다.

운영 사고는 대체로 다음 흐름으로 생깁니다.

• LLM이 한 호스트에 들어가지 않아 tensor parallel을 켠다.
• TPU 리소스 요청은 통과한다.
• topology를 빠뜨리거나 칩 개수처럼 해석한다.
• worker가 서로 다른 slice에 배치된다.
• collective 통신이 시작되지 않아 배포가 멈춘다.
• 로그에는 모델 코드 문제처럼 보이는 흔적만 남는다.

이 패턴은 Kubernetes, Ray, vLLM, JAX가 함께 엮일 때 특히 찾기 어렵습니다. 한 레이어에서는 “리소스를 배정했다”고 보고하고, 다른 레이어에서는 “worker를 기다린다”고만 말하기 때문입니다. 그래서 이번 업데이트는 YAML 한 줄의 중요성을 크게 만든 뉴스입니다.

Ray Serve, Ray Data, Ray Train에서 바뀐 실무 포인트

Ray Serve는 vLLM 기반 LLM 서빙에서 topology 선언을 통해 멀티 호스트 모델을 한 slice에 묶습니다. 운영팀 입장에서는 autoscaling, load balancing, multi-model composition은 기존 Ray Serve 방식으로 가져가되 TPU 배치만 명시하면 됩니다.

Ray Data는 iter_jax_batches()가 핵심입니다. TPU가 빨라질수록 병목은 모델 연산보다 데이터 공급에서 먼저 터집니다. 기존 파이프라인이 NumPy 배열을 만들고, 이를 다시 JAX device array로 옮기는 구조라면 step마다 host-side copy가 쌓입니다. iter_jax_batches()는 JAX array와 device sharding을 전제로 batch를 넘기기 때문에 학습과 대규모 batch inference에서 손실이 줄어듭니다.

Ray Train 쪽에서는 JaxTrainer가 TPU slice shape를 ScalingConfig 안으로 가져옵니다. 여기서 실무자가 놓치기 쉬운 포인트는 import jax 위치입니다. Google 글은 worker 함수 내부에서 JAX를 import하라고 설명합니다. TPU context는 worker별로 초기화되기 때문에 모듈 최상단 import가 cryptic한 device-init 에러를 만들 수 있습니다.

바로 적용할 수 있는 배포 기준

TPU를 처음 도입하는 팀이라면 처음부터 대형 모델을 올리지 않는 게 좋습니다. 작은 slice에서 Ray task 하나를 실행하고, 그 다음 RayService로 간단한 vLLM endpoint를 띄운 뒤, 마지막에 학습 또는 batch inference로 확장하는 순서가 안전합니다.

운영 기준은 세 가지로 잡으면 됩니다.

첫째, topology를 코드 리뷰 필수 항목으로 둡니다. TPU 관련 PR에서 accelerator_type만 있고 topology가 없으면 반려하는 규칙을 둡니다. 둘째, 배포 상태가 DEPLOYING에 오래 머물 때 모델 로그만 보지 말고 placement와 slice label을 먼저 봅니다. 셋째, TPU 비용 알림을 배포 상태와 연결합니다. 조용히 멈춘 배포는 장애 알림보다 비용 알림으로 먼저 드러날 수 있습니다.

GPU 팀이 TPU로 넘어갈 때의 비교

GPU에서는 보통 num_gpus, node selector, MIG 여부, NCCL 설정을 봅니다. TPU에서는 topology와 slice가 그 자리를 차지합니다. GPU 사고가 “드라이버와 CUDA 버전”에서 자주 난다면, TPU 사고는 “worker가 같은 slice 안에 있는가”에서 납니다.

이 차이를 문서화하지 않으면 팀은 GPU식 사고방식으로 TPU를 다루게 됩니다. “16칩이면 아무 16개나 잡으면 되겠지”라는 가정이 실패의 시작입니다. TPU는 연결 구조가 리소스의 일부입니다. 칩 수와 네트워크 topology를 분리해서 보면 안 됩니다.

실행 체크리스트

• RayService YAML에 accelerator_config.kind: tpu와 topology가 함께 있는지 확인합니다.
• topology="4x4" 같은 값이 칩 수가 아니라 slice shape라는 설명을 운영 문서에 남깁니다.
• 멀티 호스트 LLM 서빙은 배포 후 첫 요청 전에 placement 상태를 확인합니다.
• DEPLOYING 장기 지속 알림을 만들고 TPU 사용량 비용 알림과 묶습니다.
• Ray Data를 쓰는 학습·batch inference는 iter_jax_batches()로 host copy 병목을 줄입니다.
• JAX import는 worker 함수 내부에서 수행하도록 샘플 코드를 고정합니다.
• 첫 도입은 작은 slice의 end-to-end 예제로 시작하고, 이후 모델 크기와 traffic을 올립니다.

출처: Google Developers Blog, “Run Ray on TPU, Part 2: Ray AI libraries”, 2026-07-24.