Ray 2.55, TPU를 공식 가속기로 지원…GKE서 학습·서빙 통합

Google Developers Blog 공식 발표에 따르면 분산 컴퓨팅 프레임워크 Ray가 2.55 버전부터 Google Cloud TPU를 공식 가속기로 지원한다. TPU 지원이 Ray의 릴리스 파이프라인과 사전 빌드 이미지, 핵심 라이브러리 지원 범위에 들어가면서 기존의 실험적 경로에서 정식 통합 단계로 옮겨갔다.

Ray는 파이썬 작업을 여러 노드에 분산하는 오픈소스 프레임워크다. 함수 형태의 작업과 상태를 유지하는 액터를 클러스터에 배치하고, 학습과 데이터 처리, 모델 서빙을 각각 Ray Train·Data·Serve로 다룬다. 이번 변화는 GPU 중심으로 Ray를 사용해 온 개발자가 같은 작업·액터 모델을 TPU 환경에도 적용할 수 있도록 한 데 의미가 있다.

TPU에서는 여러 호스트의 칩이 고속 인터칩 연결로 묶인 슬라이스를 하나의 단위로 유지해야 한다. 분산 학습 작업이 서로 다른 슬라이스에 나뉘면 집단 통신이 멈출 수 있기 때문이다. GKE의 Ray Operator는 KubeRay와 TPU 전용 웹훅을 설치해 각 호스트에 슬라이스 이름, 작업자 번호, 토폴로지 등의 라벨을 붙이고, Ray가 같은 물리 슬라이스 안에서 작업자를 예약하도록 한다.

Ray에는 이를 추상화하는 ray.util.tpu 패키지가 추가됐다. slice_placement_group 함수는 요청한 TPU 세대와 토폴로지에 맞는 슬라이스를 찾아 전체 호스트를 한꺼번에 예약한다. Ray 공식 문서는 이 방식이 다중 호스트와 다중 슬라이스 구성의 배치 복잡성을 줄인다고 설명한다. 다만 해당 API는 공개돼 있지만 안정성 표시는 alpha여서 후속 릴리스에서 인터페이스가 바뀔 수 있다.

일반적인 작업에서는 개발자가 슬라이스 배치 함수를 직접 호출할 필요가 크지 않다. Google은 Ray Data·Train·Serve가 필요한 예약 절차를 내부에서 처리하도록 구성했다고 밝혔다. 학습은 JaxTrainer, 데이터 처리는 Ray Data, 모델 서빙은 Ray Serve와 vLLM을 조합하는 방식이다.

Google Cloud가 공개한 실행 예제는 GKE에 Ray Operator를 설치하고 TPU v6e 2x4 슬라이스 8개 칩을 구성한다. Qwen3-4B-Instruct-2507 모델을 대상으로 데이터 준비, LoRA 미세조정, 일괄 추론, vLLM 서빙까지 한 환경에서 진행하도록 설계됐다. 클러스터는 Terraform으로 만들 수 있으며 모니터링 구성도 포함된다.

이번 통합은 AI 인프라 선택지를 넓히지만, GPU 코드를 TPU로 옮기기만 하면 모든 워크로드가 그대로 동작한다는 뜻은 아니다. TPU 토폴로지와 프레임워크 호환성, GKE 구성, 리전별 자원 제공 범위와 할당량을 함께 검토해야 한다. Google Cloud의 책임 분담 문서도 Ray Operator 자체의 관리와 별개로 이용자가 워커 이미지와 애플리케이션 구성을 유지해야 한다고 안내한다.

출처: Google Developers Blog https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/
Google Cloud 공식 GKE Ray on TPU 예제 https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started
Google Cloud 공식 Ray on GKE 문서 https://docs.cloud.google.com/kubernetes-engine/docs/add-on/ray-on-gke/concepts/overview
Ray 공식 KubeRay TPU 문서 https://docs.ray.io/en/latest/cluster/kubernetes/user-guides/tpu.html