Ray on TPU 실전 가이드: KubeRay로 학습·평가 파이프라인을 옮길 때 보는 순서

Ray on TPU를 검토하는 팀이 처음부터 모델 서빙 전체를 옮기려고 하면 실패하기 쉽습니다. TPU는 강력하지만 운영 전제가 분명합니다. slice 단위 배치, 데이터 공급 속도, 체크포인트, 네트워크, JAX/XLA 친화성 같은 조건을 맞춰야 합니다. 그래서 실전에서는 온라인 트래픽보다 오프라인 학습·평가 파이프라인부터 옮기는 것이 안전합니다.

이 글은 Ray와 KubeRay를 이미 어느 정도 쓰고 있거나, GPU 기반 ML 파이프라인을 운영하는 팀이 TPU를 시험할 때 확인할 순서를 정리합니다. 핵심은 ‘TPU를 붙이면 빨라질까?’가 아니라 ‘어떤 병목이 TPU로 해결되고, 어떤 병목은 그대로 남는가?’를 측정하는 것입니다.

1단계: 옮길 작업을 좁게 고르기

처음부터 모든 학습 job을 옮기지 마세요. 가장 좋은 후보는 실패해도 사용자에게 직접 영향이 없고, 입력과 출력이 명확하며, 반복 실행 가능한 작업입니다.

예를 들면 다음과 같습니다.

• 야간 임베딩 생성
• 모델 후보별 오프라인 평가
• 대량 synthetic data 생성 후 필터링
• JAX 기반 파인튜닝 실험
• 고정된 데이터셋에 대한 batch inference

반대로 실시간 에이전트 요청, 도구 호출이 많은 복합 워크플로우, 입력 길이가 크게 흔들리는 서비스는 첫 후보로 좋지 않습니다. 이런 작업은 TPU 자체보다 세션 관리와 큐 설계가 더 큰 병목일 수 있습니다.

2단계: 현재 GPU 기준선을 남기기

TPU 테스트를 하기 전에 GPU에서 기준선을 남겨야 합니다. 기준선이 없으면 TPU가 빨라졌는지 비싸졌는지 판단할 수 없습니다.

최소한 다음 값을 기록합니다.

• 전체 실행 시간
• 처리한 샘플 수
• 시간당 처리량
• 총 클라우드 비용
• 실패율과 재시도 횟수
• 데이터 로딩 대기 시간
• 체크포인트 저장 시간
• 개발자 개입 횟수

여기서 중요한 지표는 처리량만이 아닙니다. TPU에서 실행 시간이 20% 줄었지만 실패 후 복구가 어려워 개발자 시간이 늘면 총비용은 오를 수 있습니다. 반대로 비용이 조금 비싸도 재현성과 처리량이 좋아지면 대규모 평가에서는 이득일 수 있습니다.

3단계: KubeRay와 TPU slice 배치 이해하기

Google 블로그에서 강조한 부분은 TPU slice를 완전한 단위로 예약하는 것입니다. TPU는 여러 칩과 호스트가 정해진 토폴로지로 묶여 동작합니다. 일부 리소스만 잡아서는 원하는 성능이 나오지 않습니다.

KubeRay on GKE는 TPU 하드웨어 레이아웃을 프로비저닝하고 라벨링합니다. Ray Core는 이 정보를 사용해 slice 단위 placement를 수행합니다. 개발자는 작업에서 필요한 토폴로지를 선언하고, Ray가 그에 맞는 리소스를 잡도록 구성합니다.

운영자가 봐야 할 것은 세 가지입니다. 첫째, 요청한 topology가 실제로 예약되는지. 둘째, job이 slice 일부만 잡고 대기 상태에 빠지지 않는지. 셋째, 여러 팀이 같은 클러스터를 쓸 때 리소스 경합이 어떻게 처리되는지입니다.

4단계: 데이터 공급 병목 제거하기

가속기를 바꿀 때 가장 자주 놓치는 부분이 데이터입니다. TPU가 빠르게 계산할 준비가 되어 있어도 입력 batch가 늦게 오면 의미가 없습니다.

먼저 데이터 위치를 확인합니다. Cloud Storage에서 읽는지, BigQuery에서 추출하는지, 로컬 디스크에 캐시하는지에 따라 병목이 달라집니다. 전처리 CPU가 느린 경우도 많습니다. 이때는 Ray Data를 써서 데이터 읽기와 전처리를 병렬화하고, accelerator에 맞는 batch 형태로 공급해야 합니다.

측정할 때는 전체 실행 시간만 보지 말고 stage별 시간을 나눠야 합니다. 데이터 읽기, 파싱, 토크나이징, batch 구성, accelerator transfer, 계산, 후처리, 저장을 분리합니다. TPU 도입 후 계산 시간이 줄었는데 토크나이징이 전체의 60%가 되면 다음 최적화 대상은 TPU가 아니라 CPU 전처리입니다.

5단계: 체크포인트와 재시작 설계하기

대규모 작업은 실패합니다. 네트워크가 끊기고, 노드가 재시작되고, job 설정이 잘못될 수 있습니다. TPU slice를 쓰는 큰 작업일수록 한 번 실패했을 때 손실이 큽니다.

그래서 체크포인트는 나중에 붙이는 기능이 아니라 초기 설계에 들어가야 합니다. 저장 주기, 저장 위치, 보존 개수, 재시작 시 읽을 checkpoint 선택 기준을 정합니다. Ray Train이나 JAX 학습 코드를 쓴다면 optimizer state와 random seed까지 복원되는지 확인해야 합니다.

오프라인 평가에서도 중간 결과 저장이 중요합니다. 100만 개 샘플 평가 중 80만 개에서 실패했는데 처음부터 다시 돌려야 한다면 운영 비용이 큽니다. 샘플 단위나 shard 단위로 완료 상태를 기록하는 방식이 필요합니다.

6단계: TPU와 GPU를 같은 표로 비교하기

최종 판단은 표로 해야 합니다. 감으로 ‘빠른 것 같다’고 말하면 다음 비용 회의에서 설득이 안 됩니다.

비교 표에는 다음 열을 넣습니다.

workload
hardware
topology
samples
runtime_minutes
throughput_per_hour
cloud_cost
failure_count
retry_cost
engineer_time_minutes
notes

여기서 engineer_time_minutes를 꼭 넣는 것이 좋습니다. 인프라 선택은 컴퓨팅 비용만이 아닙니다. 설정이 복잡해 운영자가 매번 개입해야 하면 실제 비용은 높습니다. 반대로 초기 설정은 어렵지만 반복 실행이 안정적이면 장기적으로 이득입니다.

운영 중 확인할 경고 신호

TPU job이 자주 대기 상태에 걸리면 slice 예약 정책을 봐야 합니다. 데이터 로딩 시간이 길면 Ray Data 파이프라인과 storage 위치를 봐야 합니다. 계산 utilization이 낮으면 batch size, sharding, model shape를 다시 봐야 합니다. 체크포인트 저장이 느리면 저장 주기와 checkpoint 크기를 조정해야 합니다.

또 하나의 경고 신호는 작은 실험이 과도하게 느려지는 것입니다. TPU는 큰 작업에서 효율이 나오는 경우가 많습니다. 빠른 실험 반복이 목적이라면 작은 GPU나 CPU path를 별도로 유지하는 것이 낫습니다. 모든 실험을 TPU로 보내면 개발 속도가 떨어질 수 있습니다.

실행 체크리스트

• 사용자 트래픽과 분리된 오프라인 workload를 첫 후보로 고른다.
• GPU 기준선: 실행 시간, 비용, 실패율, 개발자 개입 시간을 기록한다.
• KubeRay에서 TPU slice가 완전한 단위로 예약되는지 확인한다.
• 데이터 파이프라인을 stage별로 측정한다.
• 체크포인트와 shard 단위 재시작을 먼저 구현한다.
• TPU와 GPU 결과를 같은 표로 비교한다.
• 작은 실험용 경로와 대규모 실행 경로를 분리한다.

Ray on TPU의 핵심은 ‘TPU를 쓰자’가 아닙니다. 분산 실행과 하드웨어 토폴로지를 운영 가능한 수준으로 연결하는 것입니다. 오프라인 파이프라인부터 좁게 검증하면, TPU가 진짜 병목을 해결하는지 비용과 숫자로 판단할 수 있습니다.