なぜ Ray on TPU なのか

大規模な AI/ML ワークロードが増えるにつれ、GPU だけでなく Google の TPU（Tensor Processing Unit）を活用する事例が増えています。TPU は行列演算に特化したアクセラレータで、特に Transformer ベースのモデル学習で優れた性能を発揮します。

しかし TPU には スライス (Slice) という概念があります。複数の TPU チップが ICI（Inter-Chip Interconnect）という高速リンクで固定グループ化されており、ワーカープロセスは 同一スライス内 に配置されなければなりません。そうしないと all-reduce などの集合通信が完了せず、ジョブがハングします。

Ray は分散コンピューティングフレームワークで、複数マシン上で Python コードを task（ステートレス関数）と actor（ステートフルワーカー）として実行します。Ray 2.55 から TPU が ファーストクラスアクセラレータ として正式サポートされ、実験的なカスタムコンテナを用意する必要がなくなりました。

本記事では、GKE に Ray Operator を導入し、TPU スライスを取得して Ray ジョブを実行するまでの全手順を解説します。

このシリーズの Part 2 では、vLLM を使った LLM サービング、Ray Data によるデータパイプライン、JaxTrainer を用いた学習までを取り上げます。

Diagram of Ray architecture on Google Cloud TPU with GKE and slice placement group Technical Structure Concept

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

このコマンドで以下の 2 つがインストールされます：

  1. KubeRay Operator – RayCluster、RayService、RayJob の YAML を実際の Ray クラスタとしてデプロイする Kubernetes オペレーター
  2. 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  # スライス1つ = 1ワーカグループ
      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() 関数は スライス全体をアトミックに予約 します。つまり、すべてのホストが確保されるか、さもなくば失敗します。

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") とマークされており、将来変更される可能性がある点に注意してください。

Python code snippet showing ray.util.tpu.slice_placement_group usage for TPU training Algorithm Concept Visual

注意点と実践的なヒント

1. スライス境界を越えない

TPU スライスは ICI で接続された固定グループです。ワーカーを異なるスライスに分散させると all-reduce が永遠に待機し、ジョブが停止します。必ず slice_placement_group() または Ray AI ライブラリを通じて 1つのスライス に全ワーカーを配置してください。

2. トポロジーを正確に指定する

topology はスライスの形状（例：4x4、2x2x2）です。チップ数ではなく 形状 を要求する点を忘れないでください。誤ったトポロジーを指定すると GKE がスライスをプロビジョニングできません。

3. GKE Autopilot vs Standard

  • Autopilot: ノードプールを自動管理するため簡単ですが、カスタムノード設定に制限があります。
  • Standard: ノードプールを直接制御できるため、TPU スライス構成を細かく調整できます。

4. リソース制限の設定

google.com/tpu リソースは要求するチップ数と一致させる必要があります。4x4 スライスは 16 チップなので limits.google.com/tpu: 16 と設定します。

5. 日本での適用コンテキスト

日本では、NTT や Preferred Networks、東京大学などが TPU を活用した大規模言語モデルの研究を進めています。特に GKE と Ray の組み合わせは、Kubernetes ベースの MLOps パイプライン をすでに構築している組織にとって、インフラを変更せずに TPU を導入できる点がメリットです。ただし、TPU スライスはリージョンごとの利用可能容量に制限があるため、GCP コンソールで事前に TPU 割り当てを確認し、増量リクエストを出しておくことをお勧めします。

6. この技術の限界

  • TPU は GPU に比べて汎用性が低く、特定の演算（行列積）に最適化されています。
  • スライス単位でしか割り当てられないため、小規模な実験にはオーバースペックになる可能性があります。
  • slice_placement_group() API はまだ alpha 段階のため、プロダクション導入時は Ray のリリースノートを注視する必要があります。

Google Kubernetes Engine cluster with Ray Operator add-on managing TPU slice nodes Coding Session Visual

まとめ：TPU も Ray で簡単に使える時代に

本記事では、GKE に Ray Operator をインストールし、TPU スライスを要求する YAML を作成し、Ray Core の slice_placement_group() を使ってマルチホスト TPU で分散ジョブを実行する方法を解説しました。

要点まとめ：

  • TPU スライスは 1 つの単位 としてのみ動作し、ICI で接続されたホストを決して分離してはいけません。
  • GKE Ray Operator が TPU ホストに自動ラベルを付与し、スライス境界を識別します。
  • slice_placement_group() がスライス全体をアトミックに予約します。
  • 実際には Ray AI ライブラリ がこのプロセスをほとんど自動化するため、開発者はトポロジーを指定するだけで済みます。

次のステップ

  1. Part 2 で vLLM を使った LLM サービングと Ray Data パイプラインを TPU 上で実行する方法を学んでください。
  2. Ray Train の JaxTrainer を使って TPU 上で JAX ベースのモデル学習を実際に試してみてください。
  3. Ray Serve を TPU スライスにデプロイし、プロダクション推論パイプラインを構築してみてください。

合わせて読みたい

本コンテンツは、信頼性の高い情報源をもとにAIツールを活用して作成され、編集者によるレビューを経て公開されています。専門家によるアドバイスの代替となるものではありません。