왜 Ray + TPU인가?
GPU 기반으로 Ray를 이미 사용하고 있다면, 이제 Google TPU에서도 동일한 코드를 실행할 수 있습니다. 기존의 Ray 라이브러리(Data, Train, Serve)를 그대로 사용하면서, TPU 슬라이스를 마치 하나의 GPU 노드처럼 취급할 수 있게 되었어요.
Ray 2.55부터 TPU는 더 이상 실험 단계가 아닙니다. 공식 프리빌트 이미지와 릴리스 파이프라인에 포함되어, 커뮤니티 지원에 의존하지 않아도 됩니다. 이 글에서는 Ray on TPU의 기반이 되는 GKE 연동과 슬라이스 배치 개념을 중심으로 살펴봅니다.
TL;DR: TPU 슬라이스를 이해하면, Ray의 기존 개발 프로세스는 그대로 유지하면서 TPU 하드웨어를 활용할 수 있습니다.
TPU 슬라이스란?
TPU 칩들은 고정된 그룹으로 연결되어 있는데, 이 그룹을 **슬라이스(Slice)**라고 합니다. 여러 호스트 머신의 칩들이 ICI(Inter-Chip Interconnect)라는 고속 링크로 연결되어 있어요.
GPU 환경으로 비유하면, 하나의 슬라이스는 NVLink로 연결된 멀티 GPU 박스와 같습니다. 박스 외부로 나가면 고속 연결이 끊기기 때문에, 분산 학습 시 all-reduce 같은 집합 통신이 완료되지 않아 학습이 멈출 수 있습니다.
TPU에서 중요한 것은 **토폴로지(Topology)**입니다. 슬라이스의 형태를 나타내며, 예를 들어 16칩 슬라이스는 4x4로 표기합니다. Ray와 GKE가 이 토폴로지 정보를 사용하여 전체 슬라이스를 원자적으로 예약합니다.
GKE에서 Ray Operator 설정하기
Ray를 TPU에서 실행하려면 GKE 클러스터에 Ray Operator 애드온을 활성화해야 합니다. Autopilot 또는 Standard 모드에 따라 명령어가 조금 다릅니다.
Autopilot (완전 관리형 노드)
gcloud container clusters create-auto CLUSTER \
--enable-ray-operator --location=LOCATION
Standard (노드 풀 직접 관리)
gcloud container clusters create CLUSTER \
--addons=RayOperator --location=LOCATION &&
gcloud container node-pools create v6e-16-slice \
--cluster=CLUSTER \
--location=LOCATION \
--machine-type=ct6e-standard-4t \
--tpu-topology=4x4 \
--num-nodes=4
이 명령어는 KubeRay와 Ray TPU 웹훅을 설치합니다. 웹훅은 각 TPU 호스트에 ray.io/tpu-slice-name 같은 라벨을 붙여, Ray가 같은 슬라이스에 속한 머신을 식별할 수 있게 합니다.
RayCluster 매니페스트로 TPU 요청하기
RayCluster의 workerGroupSpec에서 nodeSelector를 사용하여 TPU 세대와 토폴로지를 지정하고, 리소스로 칩 수를 요청합니다.
# RayCluster workerGroupSpec 내부
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice # TPU 세대
cloud.google.com/gke-tpu-topology: "4x4" # 슬라이스 형태
# ... google.com/tpu 리소스 한도로 칩 수 요청
numOfHosts: 4 # 멀티 호스트: 슬라이스를 구성하는 호스트 VM 수
이렇게 하면 GKE가 슬라이스를 프로비저닝하고, 웹훅이 라벨을 붙이고, Ray가 라벨을 읽어 작업을 배치합니다. 코드는 그대로 Python으로 작성하면 됩니다.
Ray Core의 슬라이스 플레이스먼트 그룹
Ray Core는 분산 컴퓨팅의 기본 레이어로, 태스크와 액터를 스케줄링합니다. TPU 지원은 ray.util.tpu API에 포함되어 있으며, 핵심 함수는 slice_placement_group()입니다.
이 함수는 전체 슬라이스를 원자적으로 예약합니다. 즉, 모든 호스트가 성공하거나 아무것도 실행되지 않거나 둘 중 하나입니다.
from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
# v6e 4x4 슬라이스 하나 예약 (4개 호스트에 걸친 16칩)
spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.ready(), timeout=600)
@ray.remote(resources={"TPU": 4})
def worker(rank, world):
# 각 호스트에서 실행될 작업
pass
tasks = [
worker.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=spg.placement_group)
).remote(rank=i, world=spg.num_hosts)
for i in range(spg.num_hosts)
]
실무에서는 Ray AI 라이브러리(Data, Train, Serve)가 내부적으로 이 함수를 호출하므로, 직접 사용할 일은 많지 않습니다. 커스텀 분산 워크로드를 작성할 때만 직접 사용하면 됩니다.
⚠️ 이 API는 현재
alpha안정성 단계입니다. 공개 API지만 릴리스 간 변경될 수 있습니다.
주의사항 및 한계
- 슬라이스 단위 배치: TPU 슬라이스는 하나의 단위로 예약되므로, 슬라이스보다 작은 단위로는 실행할 수 없습니다. GPU처럼 유연한 스케일링이 어렵습니다.
- 알파 API:
slice_placement_group()은 아직 알파 단계이므로, 프로덕션 적용 시 API 변경 가능성을 염두에 두어야 합니다. - 비용: TPU는 GPU보다 시간당 비용이 낮을 수 있지만, 슬라이스 전체를 예약해야 하므로 유휴 시간이 발생할 수 있습니다.
한국 개발 생태계 적용 맥락
국내에서는 GPU 수급 문제로 인해 TPU 도입을 고려하는 기업이 늘고 있습니다. 특히 하이퍼스케일 AI 서비스를 운영하는 스타트업이나 대기업에서 Ray와 GKE를 함께 사용하는 사례가 증가할 것으로 예상됩니다. 다만, 국내 클라우드 환경에서 TPU 할당량을 확인하는 것이 우선입니다.
다음 단계 학습 방향
이 글에서는 기반을 다뤘습니다. 다음으로는 Ray Data, Ray Train, Ray Serve를 TPU에서 사용하는 방법을 학습해보세요. 특히 vLLM을 사용한 LLM 서빙과 JaxTrainer를 활용한 학습이 실무에서 유용합니다.
함께 보면 좋은 글
- Azure Databricks, 3년 ROI 331% 인증 받았다 — Forrester TEI 분석 요약
- JavaScript ShadowRealm API, 드디어 코드 격리의 시대가 온다
결론
Ray on TPU는 이제 실험 단계를 넘어 실무에서 활용할 수 있는 수준이 되었습니다. GKE가 프로비저닝과 라벨링을 담당하고, Ray Core가 슬라이스 단위로 예약하므로, 개발자는 Python 코드에 집중할 수 있습니다.
TPU 슬라이스의 개념과 제약을 이해하면, 기존 Ray 코드를 거의 수정 없이 TPU로 이전할 수 있습니다. GPU에서 이미 Ray를 사용하고 있다면, TPU로의 확장은 생각보다 간단합니다.
자세한 내용은 근거자료를 참고하세요. Happy building!