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

Google Cloud와 Ray 팀이 Ray 2.55에서 TPU를 first-class accelerator로 다루는 흐름을 정리했다. 이미 Ray로 GPU 클러스터를 운영하는 팀이라면 익숙한 task, actor, Ray Serve, Ray Data, Ray Train 패턴을 TPU에서도 사용할 수 있다는 내용이다. 하지만 TPU에는 GPU와 다른 핵심 제약이 하나 있다. multi-host TPU 작업은 반드시 하나의 slice 안에 같이 배치되어야 한다.

이 제약을 모르면 LLM 서빙이나 학습 작업이 이상한 방식으로 멈춘다. 에러가 명확히 터지는 대신 deployment가 계속 DEPLOYING 상태에 머물거나, collective operation이 끝나지 않는다. 원인은 모델 코드가 아니라 placement일 수 있다. Ray on TPU에서 topology 설정이 중요한 이유가 여기에 있다.

이 글은 Ray on TPU 발표를 바탕으로, LLM 서빙·데이터 처리·JAX 학습을 TPU에서 운영할 때 개발자가 확인해야 할 포인트를 정리한다.

TPU slice를 이해하지 못하면 장애 원인을 잘못 본다

TPU는 여러 칩이 고속 인터커넥트로 묶인 고정 그룹, 즉 slice 단위로 동작한다. Google 문서에서는 이 연결을 ICI, Inter-Chip Interconnect라고 설명한다. multi-host 모델은 이 ICI가 연결된 하나의 slice 안에 배치되어야 한다. workers가 서로 다른 slice에 흩어지면 통신이 되지 않아 작업이 멈춘다.

GPU에 익숙한 개발자는 이를 “NVLink가 없는 서로 다른 박스에 tensor parallel worker를 나눠 놓은 상황”으로 이해하면 된다. 코드에는 문제가 없어도 all-reduce가 끝나지 않는다. TPU에서도 같은 일이 발생한다. 따라서 TPU 작업에서는 칩 개수만 요청하면 부족하다. slice의 모양, 즉 topology를 요청해야 한다.

예를 들어 4x4 topology는 16개 칩으로 구성된 slice를 의미한다. 중요한 것은 16이라는 숫자가 아니라 4x4라는 slice 형태다. Ray와 GKE는 이 topology 정보를 기반으로 필요한 worker를 같은 slice에 배치한다.

GKE와 Ray Core가 placement를 대신 처리한다

과거에는 TPU에서 distributed workload를 돌리려면 placement를 직접 많이 신경 써야 했다. Ray on TPU의 요지는 이 복잡도를 GKE와 Ray Core가 흡수한다는 점이다. GKE의 Ray Operator add-on은 TPU host에 slice 정보를 나타내는 label을 붙이고, Ray Core는 이 label을 읽어 하나의 slice를 atomically reserve한다.

Ray Core에는 slice_placement_group()이라는 primitive가 있다. 직접 커스텀 workload를 만들 때는 이 API로 “v6e 4x4 slice 하나를 통째로 예약”할 수 있다. 다만 일반적인 Ray Serve, Ray Data, Ray Train 사용자는 직접 호출할 일이 많지 않다. 상위 라이브러리가 topology를 받아 내부에서 placement를 처리한다.

운영자가 알아야 할 점은 하나다. TPU 작업에서 topology 설정이 빠지면 “칩은 잡혔는데 통신이 안 되는” 상태가 발생할 수 있다. 이런 장애는 로그만 보면 모델 초기화 문제, vLLM 문제, 네트워크 문제처럼 보인다. 실제로는 worker가 같은 slice에 있지 않은 placement 문제일 수 있다.

Ray Serve에서는 accelerator_config.topology를 확인한다

LLM 서빙에서는 Ray Serve를 먼저 쓰는 팀이 많다. Ray Serve는 autoscaling, load balancing, multi-model composition을 제공하고, TPU에서는 vLLM을 통해 LLM을 서빙할 수 있다. 한 host에 들어가는 작은 모델은 비교적 단순하다. 어려운 것은 16개 칩에 sharding되는 큰 모델이다.

Google 설명에 따르면 Ray Serve on TPU에서는 다음과 같은 설정이 핵심이다.

accelerator_type: TPU-V6E
accelerator_config:
  kind: tpu
  topology: "4x4"

이 topology 한 줄이 multi-host 모델의 worker를 하나의 slice에 묶는다. 빠지면 Serve가 per-chip bundle처럼 배치하려고 하면서 worker가 여러 slice에 흩어질 수 있다. 그러면 첫 collective에서 멈추고, 배포는 계속 진행 중처럼 보인다.

장애 대응 관점에서는 배포가 오래 DEPLOYING에 머물 때 다음을 먼저 확인해야 한다.

• topology가 모델 shard 수와 맞는가.
• TPU generation과 topology label이 GKE node pool과 일치하는가.
• Ray dashboard에서 TPU resource가 의도대로 잡혔는가.
• worker pod들이 같은 slice label을 공유하는가.
• vLLM 로그보다 placement 로그를 먼저 확인했는가.

이 순서만 지켜도 불필요한 모델 디버깅 시간을 줄일 수 있다.

Ray Data는 TPU를 굶기지 않기 위한 입력 파이프라인이다

TPU는 빠르다. 그래서 모델 계산보다 데이터 공급이 병목이 되기 쉽다. Ray Data의 iter_jax_batches()는 이 문제를 줄이기 위한 API다. batch를 JAX array 형태로, device-sharded 상태로 넘겨준다. 호스트에서 NumPy를 JAX로 바꾸는 복사 비용이나 sharding 처리를 줄일 수 있다.

큰 학습이나 배치 추론에서는 이 차이가 누적된다. accelerator가 비싼데 데이터 로더가 느려서 step 사이에 빈 시간이 생기면 비용이 바로 낭비된다. 특히 TPU slice를 예약해 놓고 입력 파이프라인이 느리면, 모델 최적화보다 데이터 공급 개선이 더 큰 효과를 낼 수 있다.

실무에서는 다음 지표를 봐야 한다.

• accelerator utilization.
• input pipeline 대기 시간.
• batch 생성 시간.
• host-to-device copy 시간.
• ragged final batch 처리 방식.
• checkpoint 전후 데이터 재개 속도.

Ray Data를 쓰는 이유는 코드가 예뻐서가 아니라, 비싼 accelerator를 놀리지 않기 위해서다.

JaxTrainer는 topology를 학습 루프 바깥으로 빼준다

분산 학습에서는 topology, checkpoint, fault tolerance가 같이 중요해진다. Ray Train의 JaxTrainer는 JAX 학습 함수를 받아 slice 형태에 맞게 worker를 띄우고, 체크포인트와 재시작을 관리한다. 설정 예시는 다음 구조를 갖는다.

trainer = JaxTrainer(
    train_loop_per_worker=train_loop_per_worker,
    scaling_config=ScalingConfig(
        use_tpu=True,
        topology="4x4",
        accelerator_type="TPU-V6E",
    ),
)
trainer.fit()

주의할 점도 있다. Google 예시에서는 jax import를 worker 함수 내부에서 하라고 설명한다. 각 worker가 자기 TPU context에서 JAX를 초기화해야 하기 때문이다. 모듈 최상단에서 import하면 초기화 오류를 만나기 쉽다.

또한 topology는 chip count가 아니라 slice shape다. “TPU 16개 주세요”가 아니라 “4x4 slice로 주세요”라고 선언해야 한다. 이 작은 표현 차이가 placement와 통신 구조를 결정한다.

운영 전 검증 체크리스트

Ray on TPU를 production에 넣기 전에는 작은 샘플부터 끝까지 돌려야 한다. 바로 대형 모델을 올리면 장애 원인 범위가 너무 넓어진다.

1. GKE 클러스터에 Ray Operator add-on이 활성화되어 있는지 확인한다.
2. TPU node pool의 accelerator generation과 topology label을 기록한다.
3. 작은 Ray task로 TPU resource scheduling이 되는지 먼저 본다.
4. multi-host workload는 topology를 명시하고 worker가 같은 slice에 배치되는지 확인한다.
5. Ray Serve 배포에서 accelerator_config.topology를 빠뜨리지 않는다.
6. 배포가 DEPLOYING에 멈추면 모델 로그보다 placement와 slice label을 먼저 본다.
7. Ray Data의 iter_jax_batches()로 입력 병목을 측정한다.
8. JaxTrainer에서는 jax import 위치와 checkpoint 재시작을 검증한다.
9. Ray Dashboard에서 TPU utilization과 memory를 모니터링한다.
10. 비용 테스트는 1시간 단위가 아니라 “성공한 job 1개당 비용”으로 계산한다.

Ray on TPU의 장점은 GPU용 Ray 개발 경험을 크게 버리지 않아도 된다는 점이다. 하지만 TPU slice와 topology를 모르면 작은 설정 하나가 긴 장애로 이어진다. LLM 서빙과 학습을 TPU로 옮기려는 팀이라면, 모델 코드보다 placement 모델을 먼저 이해하는 것이 빠른 길이다.