왜 Ray on TPU인가?
AI/ML 워크로드가 점점 거대해지면서 GPU뿐 아니라 Google의 TPU(Tensor Processing Unit)를 활용하는 사례가 늘고 있습니다. TPU는 대규모 행렬 연산에 특화된 가속기로, 특히 Transformer 기반 모델 학습에서 뛰어난 성능을 보여주죠.
하지만 TPU는 GPU와 달리 슬라이스(Slice) 라는 고정된 그룹으로 묶여 있어요. 여러 개의 TPU 칩이 ICI(Inter-Chip Interconnect)라는 고속 링크로 연결된 하나의 단위로만 동작합니다. 즉, 작업자(Worker) 프로세스가 반드시 같은 슬라이스 내에 위치해야 하며, 그렇지 않으면 all-reduce 같은 집합 통신이 영원히 완료되지 않고 멈춰버립니다.
Ray는 분산 컴퓨팅 프레임워크로, 여러 대의 머신에서 Python 코드를 task(stateless 함수)와 actor(stateful 워커)로 실행합니다. 이제 Ray 2.55부터 TPU가 일급(First-class) 가속기로 공식 지원되어, 더 이상 실험적인 커스텀 컨테이너를 만들 필요 없이 GKE와 함께 바로 사용할 수 있습니다.
이 글에서는 GKE에 Ray Operator를 설치하고, TPU 슬라이스를 할당받아 Ray 작업을 실행하는 전체 과정을 다룹니다. (참고: Ray 공식 문서)
이 시리즈의 Part 2에서는 vLLM을 이용한 LLM 서빙, Ray Data를 이용한 데이터 파이프라인, JaxTrainer를 이용한 학습까지 다룹니다.
GKE에 Ray Operator 설치하기
GKE 클러스터를 만들 때 --enable-ray-operator 플래그 하나면 됩니다. Autopilot(완전 관리형)과 Standard(수동 노드 풀) 모두 지원합니다.
# Autopilot 모드 (권장)
gcloud container clusters create-auto my-ray-cluster \
--enable-ray-operator \
--location=us-central1
# Standard 모드
gcloud container clusters create my-ray-cluster \
--addons=RayOperator \
--location=us-central1
이 명령어는 두 가지를 설치합니다:
- KubeRay Operator –
RayCluster,RayService,RayJobYAML을 실제 Ray 클러스터로 배포하는 쿠버네티스 오퍼레이터 - Ray TPU Webhook – 각 TPU 호스트에
ray.io/tpu-slice-name같은 레이블을 자동으로 붙여서, Ray가 같은 슬라이스에 속한 머신을 식별할 수 있게 해줌
TPU 슬라이스 요청 YAML
TPU를 사용하려면 nodeSelector로 TPU 세대와 토폴로지를 지정하고, google.com/tpu 리소스로 칩 개수를 요청합니다. 멀티호스트 슬라이스는 numOfHosts 필드를 추가합니다.
# ray-cluster-tpu.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
name: ray-tpu-cluster
spec:
rayVersion: '2.55.0'
headGroupSpec:
serviceType: ClusterIP
template:
spec:
containers:
- name: ray-head
image: rayproject/ray:2.55.0-py310
ports:
- containerPort: 6379
- containerPort: 8265
workerGroupSpecs:
- groupName: tpu-workers
replicas: 1 # 슬라이스 하나 = 하나의 worker 그룹
template:
spec:
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: "4x4"
containers:
- name: ray-worker
image: rayproject/ray:2.55.0-py310
resources:
limits:
google.com/tpu: 16 # 4x4 = 16칩
tolerations:
- key: "google.com/tpu"
operator: "Exists"
effect: "NoSchedule"
numOfHosts: 4 # 4대의 호스트 VM
kubectl apply -f ray-cluster-tpu.yaml로 배포하면, GKE가 TPU 슬라이스를 프로비저닝하고 Webhook이 레이블을 추가하며, Ray가 이를 읽어 워커들을 올바르게 배치합니다.
Ray Core: 슬라이스 배치 그룹 (Slice Placement Group)
Ray Core는 ray.util.tpu API를 통해 슬라이스 배치 그룹을 제공합니다. 이게 핵심입니다. slice_placement_group() 함수는 전체 슬라이스를 원자적(Atomic)으로 예약합니다. 즉, 모든 호스트가 확보되거나 아니면 아예 실패합니다.
from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
import ray
ray.init(address="auto")
# v6e 4x4 슬라이스(16칩, 4호스트)를 원자적으로 예약
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_size):
# 각 워커는 자신의 rank와 전체 world_size를 받음
print(f"Worker {rank+1}/{world_size} started on TPU slice")
return rank
# 4개의 워커를 슬라이스 배치 그룹에 스케줄링
tasks = [
worker.options(
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=spg.placement_group
)
).remote(rank=i, world_size=spg.num_hosts)
for i in range(spg.num_hosts)
]
results = ray.get(tasks)
print("All workers finished:", results)
중요: 실제로는 Ray AI 라이브러리(Ray Train, Ray Serve, Ray Data)가 내부적으로
slice_placement_group()을 호출합니다. 위 코드는 커스텀 분산 워크로드를 작성할 때만 필요합니다. 공식 문서에서는 이 API를@PublicAPI(stability="alpha")로 표시하고 있어, 향후 변경 가능성은 염두에 두어야 합니다.
주의사항 및 실무 팁
1. 슬라이스 경계를 절대 넘지 마세요
TPU 슬라이스는 ICI로 연결된 고정 그룹입니다. 워커를 서로 다른 슬라이스에 나누면 all-reduce가 영원히 대기하며 작업이 중단됩니다. 반드시 slice_placement_group() 또는 Ray AI 라이브러리를 통해 하나의 슬라이스에 모든 워커를 배치하세요.
2. 토폴로지(Topology)를 정확히 지정하세요
topology는 슬라이스의 형태(예: 4x4, 2x2x2)입니다. 칩 개수가 아니라 형태를 요청한다는 점을 기억하세요. 잘못된 토폴로지를 지정하면 GKE가 슬라이스를 프로비저닝하지 못합니다.
3. GKE Autopilot vs Standard
- Autopilot: 노드 풀을 자동 관리하므로 간편하지만, 커스텀 노드 설정이 제한됩니다.
- Standard: 노드 풀을 직접 제어할 수 있어 TPU 슬라이스 구성을 세밀하게 조정할 수 있습니다.
4. 리소스 한도 설정
google.com/tpu 리소스는 요청한 칩 수와 일치해야 합니다. 4x4 슬라이스는 16칩이므로 limits.google.com/tpu: 16으로 설정합니다.
5. 한국 개발 생태계에서의 적용 맥락
국내에서는 네이버, 카카오, LG AI연구소 등에서 TPU를 활용한 대규모 언어 모델 학습이 활발합니다. 특히 GKE와 Ray 조합은 쿠버네티스 기반 MLOps 파이프라인을 이미 구축한 조직에서 추가 인프라 변경 없이 TPU를 도입할 수 있다는 장점이 있습니다. 다만, TPU 슬라이스는 리전별 가용량이 제한적이므로, GCP 콘솔에서 미리 TPU 할당량을 확인하고 증량 요청을 해두는 것이 좋습니다.
6. 이 기술의 한계
- TPU는 GPU에 비해 범용성이 낮고, 특정 연산(행렬 곱)에 최적화되어 있습니다.
- 슬라이스 단위로만 할당되므로, 소규모 실험에는 오버스펙일 수 있습니다.
slice_placement_group()API는 아직 alpha 단계이므로, 프로덕션 도입 시 Ray 릴리스 노트를 주시해야 합니다.
결론: 이제 TPU도 Ray로 쉽게 쓴다
이 글에서는 GKE에 Ray Operator를 설치하고, TPU 슬라이스를 요청하는 YAML을 작성한 뒤, Ray Core의 slice_placement_group()을 이용해 멀티호스트 TPU에서 분산 작업을 실행하는 방법을 살펴봤습니다.
핵심 정리:
- TPU 슬라이스는 하나의 단위로만 동작하며, ICI로 연결된 호스트들을 절대 분리하면 안 됩니다.
- GKE Ray Operator가 TPU 호스트에 자동 레이블을 붙여 슬라이스 경계를 식별합니다.
slice_placement_group()이 전체 슬라이스를 원자적으로 예약해줍니다.- 실제로는 Ray AI 라이브러리가 이 과정을 대부분 자동화하므로, 개발자는 토폴로지만 지정하면 됩니다.
다음 단계 학습 방향
- Part 2에서 vLLM을 이용한 LLM 서빙과 Ray Data 파이프라인을 TPU에서 실행하는 방법을 학습하세요. (원문 링크 참조)
- Ray Train의
JaxTrainer를 사용해 TPU에서 JAX 기반 모델을 학습하는 예제를 직접 실행해보세요. - Ray Serve를 TPU 슬라이스에 배포하여 프로덕션 추론 파이프라인을 구성해보세요.