Ray와 Google TPU의 만남, 왜 주목해야 할까?

딥러닝 모델의 크기가 점점 커지면서 단일 GPU로는 한계에 부딪히는 경우가 많아졌어요. 특히 최근 대규모 언어 모델(LLM)이나 멀티모달 모델을 직접 학습시키려는 시도가 늘면서, 분산 컴퓨팅의 중요성은 더욱 커지고 있죠.

그런데 분산 학습을 하려면 인프라 구성부터 시작해서 작업 분배, 장애 복구, 리소스 관리까지 신경 써야 할 부분이 한두 가지가 아닙니다. 이 과정에서 Ray는 복잡한 분산 시스템을 추상화해 주는 강력한 오픈소스 프레임워크로 자리 잡았어요.

여기에 **Google Cloud TPU(Tensor Processing Unit)**는 대규모 행렬 연산에 특화된 하드웨어로, 특히 Transformer 기반 모델 학습에서 뛰어난 성능을 보여줍니다. 이번에 Ray가 TPU를 완전 공식 지원하게 되면서, 이제 GKE 환경에서 Ray와 TPU를 함께 사용하는 것이 훨씬 매끄러워졌습니다.

이 글에서는 Ray와 TPU의 결합이 왜 강력한지 살펴보고, GKE에서 Ray 클러스터를 구성하고 분산 학습을 시작하는 방법을 단계별로 안내해 드릴게요.

Google Cloud GKE cluster with Ray nodes and TPU pods for distributed training Dev Environment Setup

본론 1: GKE에서 Ray + TPU 클러스터 구성하기

Ray의 TPU 공식 지원은 단순히 '돌아간다'는 수준을 넘어서, TPU Pod 슬라이싱, 멀티호스트 배포, 자동 복구 등을 Ray의 API로 자연스럽게 처리할 수 있게 되었다는 점이 핵심입니다. 이제 실제로 GKE에서 Ray 클러스터를 구성하는 방법을 살펴보겠습니다.

1. 사전 준비

Google Cloud 프로젝트가 준비되어 있고, gcloud CLI가 설치되어 있다고 가정할게요. 먼저 필요한 API를 활성화합니다.

# GKE, TPU, 그리고 Ray Operator에 필요한 API 활성화
gcloud services enable container.googleapis.com \
    tpu.googleapis.com \
    ray-operator.googleapis.com

2. GKE 클러스터 생성 (TPU 지원 노드 풀 포함)

TPU를 사용하려면 GKE 클러스터를 만들 때 노드 풀에 TPU 머신 타입을 지정해야 합니다. 여기서는 대표적인 TPU v5e-4 타입을 사용할게요.

# GKE 클러스터 생성 (버전 1.30 이상 권장)
gcloud container clusters create ray-tpu-cluster \
    --region=us-central1 \
    --cluster-version=1.31 \
    --release-channel=rapid

# TPU 노드 풀 추가 (v5e-4, 단일 호스트)
gcloud container node-pools create tpu-pool \
    --cluster=ray-tpu-cluster \
    --region=us-central1 \
    --machine-type=ct5lp-hightpu-4t \
    --num-nodes=1 \
    --tpu-topology=2x2

참고: ct5lp-hightpu-4t는 TPU v5e 한 개 칩(4코어)을 가진 머신 타입입니다. 더 큰 규모가 필요하면 ct5lp-hightpu-8t(8코어)나 멀티호스트 토폴로지를 고려할 수 있어요.

3. Ray Operator 설치 및 클러스터 배포

이제 GKE에 Ray Operator를 설치하고, TPU 노드 풀을 사용하는 RayCluster를 생성합니다.

# KubeRay Operator 설치
helm repo add kuberay https://ray-project.github.io/kuberay-helm/
helm install kuberay-operator kuberay/kuberay-operator \
    --version 1.2.0

그 다음, 아래와 같은 YAML 파일로 RayCluster를 정의합니다.

# ray-tpu-cluster.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-tpu-cluster
spec:
  headGroupSpec:
    serviceType: ClusterIP
    template:
      spec:
        containers:
          - name: ray-head
            image: rayproject/ray:2.40.0-py311
            ports:
              - containerPort: 6379
              - containerPort: 8265
            resources:
              requests:
                cpu: "2"
                memory: "4Gi"
              limits:
                cpu: "2"
                memory: "4Gi"
  workerGroupSpecs:
    - groupName: tpu-workers
      replicas: 2  # TPU 노드 수에 맞춰 조정
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: "tpu-v5-lite-podslice"
          tolerations:
            - key: "google.com/tpu"
              operator: "Exists"
          containers:
            - name: ray-worker
              image: rayproject/ray:2.40.0-py311
              resources:
                requests:
                  google.com/tpu: "4"  # v5e-4 기준
                limits:
                  google.com/tpu: "4"

# RayCluster 배포
kubectl apply -f ray-tpu-cluster.yaml

4. 분산 학습 스크립트 실행

이제 Ray Job을 제출하여 TPU에서 학습을 실행해 봅시다. Ray의 trainer API를 사용하면 TPU 설정을 자동으로 감지합니다.

# train.py
import ray
from ray import train
from ray.train.torch import TorchTrainer

def train_func(config):
    # TPU 초기화 (Ray가 자동으로 처리)
    import torch_xla
    import torch_xla.core.xla_model as xm
    
    device = xm.xla_device()
    
    # 간단한 모델 정의
    import torch
    import torch.nn as nn
    
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()
    
    # 더미 데이터
    data = torch.randn(100, 10).to(device)
    target = torch.randn(100, 1).to(device)
    
    # 학습 루프
    for epoch in range(10):
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        xm.optimizer_step(optimizer, barrier=True)
        
        if xm.is_master_ordinal():
            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# Ray Trainer 구성
ray.init()
trainer = TorchTrainer(
    train_func,
    scaling_config=train.ScalingConfig(
        num_workers=2,  # TPU 코어 수
        use_gpu=False,  # TPU 사용 시 False
        resources_per_worker={"google.com/tpu": 4}
    )
)

result = trainer.fit()
print("Training completed!", result)

# Ray Job 제출
ray job submit --address http://localhost:8265 -- python train.py

이렇게 하면 Ray가 자동으로 워커 간 통신을 구성하고, TPU 환경에 맞는 XLA 연산을 실행합니다.

Ray dashboard showing distributed training job metrics across multiple TPU nodes System Abstract Visual

본론 2: 운영 환경에서의 주의사항과 심화 팁

GKE에서 Ray와 TPU를 함께 운영할 때 몇 가지 중요한 포인트를 짚어볼게요.

TPU 특성 이해하기

TPU는 GPU와 달리 호스트 간 직접 통신이 제한적입니다. 특히 멀티호스트 토폴로지(예: v5e-8 이상)에서는 네트워크 토폴로지를 정확히 지정해야 해요. Ray는 TPU Pod 슬라이스를 자동으로 감지하지만, 커스텀 토폴로지가 필요하다면 placement group을 직접 정의해야 할 수 있습니다.

비용 관리 전략

TPU는 GPU보다 단위 성능당 비용이 낮을 수 있지만, **예약형(Reserved)**과 온디맨드(On-demand) 가격 차이가 큽니다. 장기 학습 작업이라면 예약형을, 테스트 단계라면 온디맨드를 사용하는 것이 비용 효율적이에요. GKE의 Node Auto Provisioning을 활용하면 워크로드에 맞춰 노드를 동적으로 조정할 수 있습니다.

장애 복구 및 체크포인팅

분산 학습에서 가장 중요한 것은 체크포인팅입니다. Ray의 train.Checkpoint 기능을 활용하면 주기적으로 모델 상태를 저장할 수 있어요. GCS(Google Cloud Storage)에 직접 체크포인트를 저장하도록 구성하면 TPU 노드가 재시작되어도 학습을 이어갈 수 있습니다.

# 체크포인트를 GCS에 저장하는 예시
from ray.train import Checkpoint

checkpoint = Checkpoint.from_dict({
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "epoch": epoch
})

# GCS 경로에 저장
train.report(metrics={"loss": loss.item()}, checkpoint=checkpoint)

모니터링과 알림

GKE 환경에서 TPU 사용량을 모니터링하려면 Cloud Monitoring과 Ray Dashboard를 함께 활용하는 것이 좋습니다. 특히 TPU의 온도, 전력 소모, utilization을 추적하면 병목 현상을 빠르게 발견할 수 있어요.

# Ray Dashboard 포트 포워딩
kubectl port-forward service/ray-tpu-cluster-head-svc 8265:8265

국내 클라우드 환경에서의 적용 맥락

국내에서는 아직 TPU를 사용하는 사례가 많지 않지만, GPU 비용 부담이 커지면서 TPU 도입을 검토하는 기업이 늘고 있어요. 특히 네이버클라우드, KT클라우드 등 국내 클라우드 서비스에서도 TPU 인스턴스를 제공하기 시작하면서 접근성이 높아지고 있습니다.

다만, 국내 환경에서는 아직 한국어 자료가 부족하고 커뮤니티가 활성화되지 않아서 초기 삽질이 불가피할 수 있어요. TPU 특유의 XLA 컴파일러 동작 방식에 익숙해지는 것이 첫 번째 관문이라고 볼 수 있습니다.

Architecture diagram of Ray on Google Kubernetes Engine with TPU accelerators Technical Structure Concept

결론: 이제 TPU 분산 학습의 문턱이 낮아졌다

Ray의 TPU 공식 지원은 분산 학습 생태계에서 중요한 이정표가 되었어요. 이전에는 TPU를 사용하기 위해 별도로 분산 학습 코드를 작성해야 했지만, 이제는 Ray의 추상화된 API를 통해 TPU와 GPU를 동일한 방식으로 다룰 수 있습니다.

실무 적용을 위한 제언

  1. GPU에서 개발, TPU에서 학습: 개발 사이클에서는 GPU 인스턴스를 사용해 코드를 빠르게 검증하고, 실제 학습은 TPU로 전환하는 전략이 비용 효율적입니다. Ray는 두 환경 모두 지원하므로 코드 수정이 거의 필요 없어요.

  2. 작은 규모부터 시작: 처음부터 대규모 클러스터를 구성하기보다는 단일 호스트 TPU(v5e-4)에서 파일럿 프로젝트를 진행해 보세요. TPU의 동작 방식과 디버깅 방법을 익힌 후 확장하는 것이 안전합니다.

  3. 커뮤니티와 자료 활용: Ray와 TPU 관련 공식 문서와 예제 코드가 계속 늘어나고 있습니다. 이 글에서 소개한 Ray와 GKE를 활용한 분산 학습 가이드를 기반으로 더 깊이 있는 학습을 이어가시면 좋겠어요.

이 기술의 한계 및 주의사항

  • TPU는 특정 모델 구조(특히 Transformer 계열)에 최적화되어 있어, CNN이나 비정형 데이터 처리에서는 GPU보다 성능이 떨어질 수 있습니다.
  • XLA 컴파일러가 지원하지 않는 일부 PyTorch 연산은 수정이 필요할 수 있습니다.
  • 멀티호스트 TPU 설정 시 네트워크 대역폭이 병목이 될 수 있어, 데이터 로딩 파이프라인 최적화가 필수적입니다.

함께 보면 좋은 글

다음 단계 학습 방향

  • Ray의 Trainer API와 Tune(하이퍼파라미터 튜닝)을 결합한 워크플로우 구성
  • GKE의 Autopilot 모드에서 Ray 클러스터 운영 방법 학습
  • TPU v6e(Trillium)의 성능 특성과 활용 사례 조사

이 글이 여러분의 분산 학습 여정에 실질적인 도움이 되었길 바랍니다. 궁금한 점이 있다면 댓글로 남겨주세요. 😊

본 콘텐츠는 신뢰할 수 있는 출처를 바탕으로 AI 도구를 활용하여 초안이 작성되었으며, 편집자의 검토를 거쳐 발행되었습니다. 전문가의 조언을 대체하지 않습니다.